[QUESTION] (Bug?) Leaking gradients through cloned states

Hi, I am trying to intentionally break the computational graph in a warp simulation. I want to simulate n frames, but I only want the backpropagation to consider a horizon of h states at a time (basically a simplified version of SHAC/AHAC)

My idea is simple: During computation, I just "clone" the state where I want to interrupt the computational graph. However, unexpectedly, this does not entirely break the gradient backpropagation. For some reason, I also need to re-allocate auxiliary model variables (I am using the Featherstone integrator).

I wondered why this is necessary and/or if this is a bug. As far as I can see, the method integrator.allocate_model_aux_vars(model) does not make use of any state variable, neither does the model or the integrator depend on the state, so I can't see how it affects the computational graph here. Lastly, when I perform a "normal" simulation (not "detaching" anything), a call to this method doesn't change my gradients at all.

Here's a minimal example to reproduce the issue:

import numpy as np
import trimesh
import warp as wp
import warp.sim

@wp.kernel
def loss_fun(joint_q: wp.array(dtype=float), loss: wp.array(dtype=float)):
    wp.atomic_add(loss, 0, joint_q[1] ** 2.)

def clone_state(state: wp.sim.State) -> wp.sim.State:
    """Clone a state, breaking the computational graph."""
    new_state = wp.sim.State()
    for attribute in dir(state):
        if attribute.startswith('__') or attribute == "requires_grad":
            continue
        try:
            setattr(new_state, attribute, getattr(state, attribute))
        except AttributeError as exe:
            if 'object has no setter' in str(exe):
                continue  # This is a derived property
        if isinstance(getattr(new_state, attribute), wp.array):
            requires_grad = getattr(new_state, attribute).requires_grad
            if requires_grad:
                setattr(new_state, attribute, wp.clone(getattr(new_state, attribute), requires_grad=False))
                getattr(new_state, attribute).requires_grad = True
    if getattr(new_state, '_featherstone_augmented', False):
        new_state._featherstone_augmented = False  # Avoid leaking gradients in auxiliary featherstone variables
    return new_state

tm = trimesh.creation.uv_sphere(.5)
builder = wp.sim.ModelBuilder()
b = builder.add_body()
builder.add_shape_mesh(b, pos=wp.vec3(0., 0., 0.), rot=wp.quat_from_axis_angle(wp.vec3(0., 0., 1.), wp.pi),
                          mesh=wp.sim.Mesh(vertices=(tm.vertices + np.array([1., 0., 0.])).tolist(),
                                           indices=tm.faces.tolist()))
builder.add_joint_free(b)
builder.joint_q[1] = 5.
model = builder.finalize(requires_grad=True)
integrator = wp.sim.FeatherstoneIntegrator(model=model)
states = list()

steps = 3
dt = 0.01
for _ in range(steps):
    states.append(model.state(requires_grad=True))

loss_val = wp.array([0.], dtype=float, requires_grad=True)
tape = wp.Tape()
with tape:
    for step in range(steps - 1):
        if step == steps // 2:
            s0 = clone_state(states[step])
            integrator.allocate_model_aux_vars(model)  # TODO: This should not be necessary - why is it?
        else:
            s0 = states[step]
        s0.clear_forces()
        s1 = states[step + 1]
        integrator.simulate(model, s0, s1, dt)

    wp.launch(loss_fun, dim=1, inputs=[states[-1].joint_q], outputs=[loss_val])

tape.backward(loss_val)
print('Final position:', s1.joint_q.numpy()[:3])
print('Position gradient (last state):', s1.joint_q.grad.numpy()[:3])
print(f'Position gradient (cloned state, nr {steps // 2}):', states[steps // 2].joint_q.grad.numpy()[:3])
print('Position gradient (first state):', states[0].joint_q.grad.numpy()[:3])

This prints (as it should be)

Final position: [6.4252388e-11 4.9970579e+00 2.3770302e-17]
Position gradient (last state): [0.       9.994116 0.      ]
Position gradient (cloned state, nr 1): [0. 0. 0.]
Position gradient (first state): [0. 0. 0.]

However, removing the line integrator.allocate_model_aux_vars(model) , it prints

Final position: [ 2.0269795e-11  4.9970579e+00 -1.4608122e-17]
Position gradient (last state): [0.       9.994116 0.      ]
Position gradient (cloned state, nr 1): [0. 0. 0.]
Position gradient (first state): [-2.9500473e-01 -1.2510548e-09 -2.3670459e-09]

Note that the gradient for the first state is neither zero nor the one it would be if the states were not "detached" at all.

Start with the minimal Python reproducer, especially clone_state(), FeatherstoneIntegrator.simulate(), and allocate_model_aux_vars(model); compare the tape's gradients with and without the allocation call. Done means the source of the gradient difference and the expected graph-detachment behavior are established, with the observed result verified against the example.