[BUG] FeatherstoneIntegrator causes wrong gradient accumulation

[tom1484]

Bug Description

I'm recently trying to obtain the gradient of model state w.r.t. model parameter, for example, the gradient of joint_q w.r.t. body_mass. The derivation of gradient works, but I found that the gradient improperly explodes.

I think their might be some incorrect accumulation during the steps so I wrote the code below to ensure that the derived gradients are the same as long as the initial states don't vary.

But it shows that the gradients do not remain the same, and I think it has some thing to do with the FeatherstoneIntegrator because if I re-initialize the integrator every time I reset the state instead of using the same integrator every time, the gradients change (though they are still not all the same, but this might be caused by #249 ).

import math
import os

from typing import Optional

import numpy as np

import warp as wp
import warp.sim as wsim
import warp.sim.render as wrender

import torch as th
import gymnasium as gym


wp.init()
wp.set_device("cuda")


articulation_builder = wsim.ModelBuilder(up_vector=wp.vec3(0.0, 0.0, 1.0))
urdf_path = "inverted_pendulum.xml"
wsim.parse_urdf(
    urdf_path,
    articulation_builder,
    xform=wp.transform_identity(),
    density=100.0,
    armature=0.0,
    stiffness=0.0,
    damping=0.0,
    limit_ke=1.0e4,
    limit_kd=1.0e1,
    enable_self_collisions=False,
)

builder = wsim.ModelBuilder()
builder.add_builder(
    articulation_builder,
    xform=wp.transform_identity(),
)
builder.joint_q[-1] = np.pi - 0.1

# finalize model
model = builder.finalize(requires_grad=True)
model.ground = False


def rollout(model):
    fps = 60
    substeps = 10
    total_steps = 5

    tape = wp.Tape()
    with tape:
        states = []
        
        # NOTE: Reusing the integrator causes incorrect gradients
        integrator = wsim.FeatherstoneIntegrator(model)
        sim_dt = 1 / fps / substeps
        for _ in range(total_steps):
            state = model.state()
            # NOTE: Different result if the integrator gets initialized every time
            # integrator = wsim.FeatherstoneIntegrator(model)

            for _ in range(substeps):
                state.clear_forces()
                state_next = model.state()

                integrator.simulate(
                    model, state, state_next, sim_dt
                )
                state = state_next

            states.append(state)
    
    return tape, states, vars


tape, states, vars = rollout(model)

for state in states:
    joint_q = state.joint_q
    selection = wp.ones(joint_q.shape)
    tape.backward(grads={joint_q: selection})
    print(model.body_mass.grad)

    tape.zero()

Output

[0.        0.        0.2937096]
[0.        0.        0.5874192]
[0.        0.        0.8811288]
[0.        0.        1.1748384]
[0.        0.        1.468548]

System Information

• CPU: Intel(R) Xeon(R) W-2255
• Ubuntu 20.04
• Python 3.10.14
• WARP 1.2.1