Warp pathwise gradients are different in reproduction of code

[DINHQuangDung1999]

Hello,

I am working on an algorithm, and I am trying it out with the Bounce task. I notice that every time I run the optimization code, I receive the gradient values that is slightly different, leading to a potentially completely different loss trajectory and convergence. I wonder if this is a sort of numerical error due to using float32, or is there anything more structured and related to the gradient computing mechanism of Warp. I attach my code below. To run the code, move inside the extracted folder and run

python grad_bounce_optimize.py --config-name conf_multicontact

You can switch between the values for grad_type in the conf_multicontact.yaml config file.

Update 1: As soon as i comment out the collision handling wp.collide inside the simulation loop, the gradient is consistent.

    def compute_loss(self):
        # run control loop
        for i in range(self.sim_steps):
            self.states[i].clear_forces()
            # wp.sim.collide(self.model, self.states[i], requires_grad = True)
            self.integrator.simulate(self.model, self.states[i], self.states[i + 1], self.sim_dt)
        
        # compute loss on final state
        wp.launch(loss_kernel, dim=self.num_envs, inputs=[self.states[-1].particle_q, self.target, self.loss])
        return self.loss

Can you help me explain this behavior?

Update 2: In the previous experiment, I intentionally put the wall above the particle at a low altitude and give the particle a large enough upward initial velocity so that it makes a number of contacts before reaching the target. However, when I return to the original bounce experiment (which have only 1-2 contacts), the pathwise gradient is stable. Does the number of contact have a relationship to the non-deterministic property of the contact handling function?
Update 3: Additional experiments were carried out. Notably, reducing the task horizon make the gradient stable.
bounce_exp.zip

[shi-eric]

@eric-heiden @mmacklin can probably provide additional information, but contact handling in Warp is not deterministic unless you take additional care in ensuring that the contacts are stored and processed in a stable order that's not changing from run to run due to how threads on a GPU complete at different speeds.

Also, whenever you see an atomic reduction involving floating point numbers, this will also generally be non-deterministic.

[DINHQuangDung1999]

Hi @shi-eric,
Thank you for replying. This raise a question for me of whether I need to use wp.sim.collide in my simulation loop or not if I want my code to be 1. reproducible and 2. correct ? I notice that if I do not, then at least, the particle collision will not be handled properly. For example, the particle might experience contact even before it reaches any where near the object.