Consufing Gradients on a Simple Scene

[Wilbur-Django]

Confusing Gradient Output

I got confusing output gradients from Nimble on a simple scene. The scene is about two balls with the same mass making fully elastic collision. In this scene, Nimble gives inconsistent gradients with the analytical gradients.

Scene description

Two balls are allowed to move horizontally. There is no friction or gravity. The two balls are of the same mass 1kg and have the same radius r = 0.1m. In the beginning, the left ball at x1 = 0 (shown in blue) moves at v0 = 1m/s to the right, while the right ball at x2 = 0.52m (shown in green) has velocity u0 = 0. Since there is no friction, the blue ball would make the uniform motion. At t = 0.5s, the two balls would collide. Then the two balls would exchange their speeds since they are of the same mass and the collision is fully elastic. The blue ball would then stay still while the green ball moves at 1m/s to the right. At t = T = 1s, the green ball would appear at xf = 1.2m.

Gradients computation

It is easy to show that the analytical form of xf w.r.t. (x1, x2, v0, u0) is:
xf = v0 T + x1 + 2r

So the analytical gradient of xf w.r.t. (x1, x2, v0, u0) is (1, 0, 1, 0).
However, the output gradients from Nimble is (0.75, 0.25, 0.91, 0.08), which is obviously inconsistent with the analytical gradients.

Reproduce

System configuration:

• OS: Ubuntu 20.04 LTS
• CPU: AMD Ryzen Threadripper 3970X 32-Core Processor
• GPU: NVIDIA GeForce RTX 3090
• Nimblephysics version: 0.8.38
• Pytorch version: 1.13.0
• Python: 3.9.13

Source code:

import nimblephysics as nimble
import torch


def create_ball(radius, color):
    ball = nimble.dynamics.Skeleton()
    sphereJoint, sphereBody = ball.createTranslationalJoint2DAndBodyNodePair() 

    sphereShape = sphereBody.createShapeNode(nimble.dynamics.SphereShape(radius))
    sphereVisual = sphereShape.createVisualAspect()
    sphereVisual.setColor([i / 255.0 for i in color])
    sphereShape.createCollisionAspect()
    sphereBody.setFrictionCoeff(0.0)
    sphereBody.setRestitutionCoeff(1.0)
    sphereBody.setMass(1)

    return ball


def create_world():
    world = nimble.simulation.World()
    world.setGravity([0, 0, 0]) # No gravity

    radius = 0.1
    world.addSkeleton(create_ball(radius, [68, 114, 196]))
    world.addSkeleton(create_ball(radius, [112, 173, 71]))

    return world


def simulate_and_backward(world, x1, x2, v0, u0):
    # Ball 1 is initialized to be at x1 on the x-axis, with velocity v0.
    # Ball 2 is initialized to be at x2 on the x-axis, with velocity u0.
    # The zeros below mean that the vertical positions and velocities are all zero.
    # So the balls woul only move in the horizontal direction.
    init_state = torch.tensor([x1, 0, x2, 0, v0, 0, u0, 0], requires_grad=True)
    control_forces = torch.zeros(4) # No external forces
    total_simulation_time = 1.0 # simulate for 1 second
    num_time_steps = 1000       # split into 1000 discrete small time steps
    # Each time step has length 0.001 seconds
    world.setTimeStep(total_simulation_time / num_time_steps)
    state = init_state
    states = [state]
    for i in range(num_time_steps):
        state = nimble.timestep(world, state, control_forces)
        states.append(state)

    # xf is the final x-coordinate of ball 2
    xf = state[2]
    xf.backward()

    # The gradients on the y-axis are irrelevant, so we exclude them.
    grad = (init_state.grad)[0:8:2] 
    print(f"xf = {xf.detach().item()}")
    print(f"gradients of xf = {grad}")

    return states


if __name__ == "__main__":
    world = create_world()
    gui = nimble.NimbleGUI(world)
    gui.serve(8080)
    states = simulate_and_backward(world, x1=0, x2=0.52, v0=1, u0=0)
    gui.loopStates(states)
    input()
    gui.stopServing()

Execution results:

xf = 1.1989999809264922
gradients of xf = tensor([0.7500, 0.2500, 0.9197, 0.0803])

[gonultasbu]

modifying the input parameters to function as (increasing the u0 to 0.2 from 0.0)
states = simulate_and_backward(world, x1=0.0, x2=0.52, v0=1.0, u0=0.2)

changes the gradients as follows:
gradients of xf = tensor([0.7500, 0.2500, 0.8997, 0.1002])

the derived analytical equation does not imply such behavior, maybe there is something I missed?

[Wilbur-Django]

modifying the input parameters to function as (increasing the u0 to 0.2 from 0.0) states = simulate_and_backward(world, x1=0.0, x2=0.52, v0=1.0, u0=0.2)

changes the gradients as follows: gradients of xf = tensor([0.7500, 0.2500, 0.8997, 0.1002])

the derived analytical equation does not imply such behavior, maybe there is something I missed?

Yes, I agree with you. The analytical gradients of xf should still be [1, 0, 1, 0] in the case of u0 = 0.2.

I'm not so familiar with the gradient computation inside Nimble. Could someone explain whether it is the expected behavior or not? I really appreciate any help you can provide!

[gonultasbu]

I did not rigorously work through the analytical equation, but my immediate intuition is that it is true only locally. If we set the input parameters s.t. the collision would never occur e.g. u0 > v0, wouldn't that break the analytical formula? The collision case does not necessarily show a smooth and differentiable behavior w.r.t input parameters, therefore I cannot conclusively say anything about the gradient computation. However, I would expect to have a gradient w.r.t the u0.