[QUESTION] Calling tape.backward multiple time leads to incorrect gradient

[tom1484]

I'm currently debugging the gradient flow of FeatherstoneIntegrator so I wrote the following code:

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

    vars = []

    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):
            integrator = wsim.FeatherstoneIntegrator(model)
            state = model.state()

            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)
            vars.append(state.joint_tau)
    
    return tape, states, vars


tape, states, vars = rollout(model)

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

    tape.zero()

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

    tape.zero()

The result of the gradient calculation should be [0. 0. 0.2937096] for every var since I reset integrator and state before every step. But weird thing happens, the gradients I got was

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

Only the first result is correct while others keep growing every a new loss variable is involved.

How did this happen? I did call tape.zero() every time after backward...