Hi,

thanks a lot for releasing a 3rd party implementation of the paper. Neverthless, I'm afraid there is some problem with your code or at least the hyperparameters are not chosen well. This can be shown by looking at the validation error as follows:

n_epoch = 100
batches_per_epoch = 100

data_test = gen(n_objects, True)

losses = []
losses_test = []
for epoch in range(n_epoch):
    for _ in range(batches_per_epoch):
        objects, sender_relations, receiver_relations, relation_info, target = get_batch(data, 30)
        predicted = interaction_network(objects, sender_relations, receiver_relations, relation_info)
        loss = criterion(predicted, target)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        losses.append(np.sqrt(loss.data[0]))

    objects, sender_relations, receiver_relations, relation_info, target = get_batch(data_test, 30)
    predicted = interaction_network(objects, sender_relations, receiver_relations, relation_info)
    losses_test.append(np.sqrt(criterion(predicted, target).data[0]))

    clear_output(True)
    plt.figure(figsize=(20,5))
    plt.subplot(131)
    plt.title('Epoch %s RMS Train Error %s' % (epoch, np.mean(losses[-100:])))
    plt.plot(losses)
    plt.subplot(132)
    plt.title('Epoch %s RMS Test Error %s' % (epoch, np.mean(losses_test[-100:])))
    plt.plot(losses_test)        
    plt.show()

I got train RMS 2.3 but validation RMS 209.6. Update: This is mostly because you train on just a single "scene", so that the network actually never sees any other masses and thus cannot generalise to those.