warp
warp copied to clipboard
Published
2 months ago
•
NVIDIA
NVIDIA
[BUG] Gradients for Inertial Properties in FeatherstoneIntegrator
Bug Description
EDIT: This issue does not only appear for long time horizons. Upon closer inspection, even for short horizons, gradients for inertial properties are wrong using the FeatherstoneIntegrator. The example posted in the first reply to this post reproduces these issues "better" even for short time horizons.
Issue: When running ~~longer~ simulations in warp, I have repeatedly encountered missing and/or what I believe to be wrong gradients for inertial properties.
Minimal Example:
To analyze this, I have set up a simple simulation: A Pendulum attached to a revolute joint is actuated by a constant momentum around the joint axis. The "goal" is to rotate the pendulum as far as possible in limited simulation time by changing its inertial properties. Intuitively, the lower the I_{xx} attribute of the pendulum, the further it will swing. The gradient of the loss should inform me how to change the inertial properties to achieve this.
However, if the simulation is sufficiently long, the gradient switches sign (edit: and even for short time horizons, the gradients seem to be wrong). Note that the experiment is designed s.t. the Pendulum does not perform a full swing (not even a half swing for the longest horizon). Also, the objective is measured as the absolute joint angle, i.e., a full rotation would still result in a smaller loss.
Here's a plot of the gradient vs. simulation time behavior (evaluated running 8 independent simulations):
Ablations: The following ablations/changes to the simulation still resulted in the same behavior for the gradient:
- Changing the number of substeps per frame between a minimum of 15 and a maximum of 2000 (at a certain number of substeps, the gradient just exploded though)
- Starting from different initial joint angles
- Note that warp by default does not consider the symmetry of the inertia matrix when computing gradients (#585 ) : Manually fixing this as described in the issue does not affect this experiment significantly
Fixes: After a while, I was able to get a similar example running using the SemiImplicitIntegrator.
Here are source code and a Video of the simulation:
from pathlib import Path
from typing import Dict, List, Optional
from matplotlib import pyplot as plt
import torch
import warp as wp
import warp.sim
import warp.sim.render
@wp.kernel
def pendulum_loss_fn(
joint_q: wp.array(dtype=float),
loss: wp.array(dtype=float),
):
wp.atomic_add(loss, 0, -(joint_q[0]) ** 2.)
@wp.kernel
def set_momentum(m: float, idx: int, act: wp.array(dtype=float)):
"""Set the momentum."""
act[idx] = m
class PendulumSim:
momentum: float = 65
def __init__(self, num_frames, stage_path: Optional[str] = None) -> None:
builder = wp.sim.ModelBuilder(up_vector=(0., 0., 1.))
body = builder.add_body(name='Pendulum')
builder.add_shape_box(body, hx=.1, hy=.2, hz=.4)
builder.add_joint_revolute(-1, body,
parent_xform=wp.transformf(wp.vec3(0., 0., 1.), wp.quat_identity()),
child_xform=wp.transformf(wp.vec3(0., 0., .3), wp.quat_identity()),
armature=10)
self._store_intermediate_states: bool = True
self.name: str = "pendulum"
self.num_frames: int = num_frames
self.substeps: int = 15
self.current_frame: int = 0
self.current_substep: int = 0
self.fps: float = 60
self._requires_grad: bool = True
self.latest_gradients: Dict[wp.array, torch.Tensor] = {}
self.builder: wp.sim.ModelBuilder = builder
self.model: wp.sim.Model = builder.finalize(requires_grad=True)
self.control: wp.sim.Control = self.model.control()
self.states: List[List[wp.sim.State]] = [[]]
self._reset_states()
self.integrator: wp.sim.FeatherstoneIntegrator = wp.sim.FeatherstoneIntegrator(self.model)
if self.model.joint_count > 0:
wp.sim.eval_fk(self.model, self.model.joint_q, self.model.joint_qd, None, self.current_state)
self._do_render: bool = stage_path is not None
self.stage_path: Optional[str] = stage_path
if self._do_render:
Path(stage_path).unlink(
missing_ok=True
) # It seems like usda sometimes just appends instead of overwriting.
self.renderer = wp.sim.render.SimRenderer(self.model, stage_path, fps=self.fps)
self.loss = wp.array((0.,), dtype=float, requires_grad=True)
def launch(self):
"""Top level entry point to launch the simulation."""
if self._do_render: # I like to render the initial state as well.
self.render()
with wp.ScopedTimer('Simulation'):
tape = wp.Tape()
with tape:
self.integrator.allocate_model_aux_vars(
self.model
) # Otherwise, there is no gradient for model inertial properties.
f = 0
while not self._break_condition():
self.current_frame = f
self.step_frame()
f += 1
self.loss_fn()
with wp.ScopedTimer(name='backward'):
tape.backward(self.loss)
if self._do_render:
self.renderer.save()
def step_frame(self):
"""Perform a single simulation step (substep) and render the frame."""
for ss in range(self.substeps):
self.current_substep = ss
self.simulation_step()
if self._do_render:
self.render()
def simulation_step(self):
m = 50
wp.launch(set_momentum, 1, inputs=[m, 0], outputs=[self.control.joint_act])
s0 = self.current_state
s0.clear_forces()
self.integrator.simulate(self.model, s0, self.next_state, self.sim_dt, control=self.control)
def loss_fn(self):
"""Computes the distance travelled"""
wp.launch(
pendulum_loss_fn, 1, inputs=[self.current_state.joint_q], outputs=[self.loss])
def render(self):
"""Render the current frame."""
self.renderer.begin_frame(self.sim_time)
self.renderer.render(self.current_state)
self.renderer.end_frame()
@property
def current_state(self) -> wp.sim.State:
"""Returns the currently active state."""
if not self._store_intermediate_states:
return self.states[0][0]
return self.states[self.current_frame][self.current_substep]
@property
def next_state(self) -> wp.sim.State:
"""Returns the state that the integrator will be writing to next."""
if not self._store_intermediate_states:
return self.states[0][1]
if self.current_substep < self.substeps - 1 or self.current_frame == self.num_frames - 1:
return self.states[self.current_frame][self.current_substep + 1]
else:
return self.states[self.current_frame + 1][0]
@property
def frame_dt(self) -> float:
"""The time that passes between two frames."""
return 1 / self.fps
@property
def sim_dt(self) -> float:
"""The time that passes within one integratino step."""
return self.frame_dt / self.substeps
@property
def sim_time(self) -> float:
"""The current simulation time."""
return self.current_frame * self.frame_dt + self.current_substep * self.sim_dt
def _break_condition(self):
"""Called after every frame to see if the simulation should stop."""
return self.current_frame >= self.num_frames - 1
def _reset_states(self):
"""Sets all the states of the simulation to the initial state."""
if not self._store_intermediate_states:
self.states: List[List[wp.sim.State]] = [[self.model.state() for _ in range(2)]]
else:
self.states: List[List[wp.sim.State]] = [[self.model.state() for i in range(self.substeps)]
for _ in range(self.num_frames)]
self.states[self.num_frames - 1].append(self.model.state()) # A final state to integrate into at the end.
for fr in self.states:
for s in fr:
s.clear_forces()
def main():
mass_gradients = list()
inertia_gradients = list()
losses = list()
num_frames = range(20, 60, 5)
for frames in num_frames:
sim = PendulumSim(frames)
sim.launch()
mass_gradients.append(sim.model.body_mass.grad.numpy()[0])
inertia_gradients.append(sim.model.body_inertia.grad.numpy()[0, 0, 0])
losses.append(sim.loss.numpy()[0])
fig, ax = plt.subplots(layout='constrained')
# plt.plot(num_frames, mass_gradients, label='Mass')
ax.plot(num_frames, inertia_gradients, label='Inertia')
ax.set_xlabel('Frames')
ax.set_ylabel('Gradient')
ax2 = ax.twinx()
ax2.plot(num_frames, losses, label='Loss', color='orange')
ax2.set_xlabel('angle [rad]')
plt.title('Pendulum Gradient')
plt.legend()
plt.tight_layout()
plt.show()
if __name__ == '__main__':
main()
https://github.com/user-attachments/assets/90b06a48-d690-4456-9d78-f8e5d6332e11
Is there anything I can do to figure out where the gradient computation fails? The options
wp.config.verbose = True
wp.config.verbose_warnings = True
do not return any warnings. The tape looks correct. Still, there seems to be a threshold simulation time that I don't know in advance and that seems to affect the gradient.
System Information
- warp 1.7.0
- Python 3.12.2
- Driver Version: 570.124.04
- CUDA Version: 12.8
I wrote up an even simpler use case (no gravity, armature, ...) and compared the warp gradients with analytical gradients that are easy to obtain for this simplified model.
Based on this even smaller experiment and anecdotal evidence, the main issue is gradients w.r.t. inertial properties. Gradients for control inputs, for instance, are way more accurate, even over long time horizons. (I believe this is unrelated to #585 , see original post)
Running a minimal pendulum-like test (actually, it is just a sphere rotating around its com...):
import warp as wp
import warp.sim
@wp.kernel
def loss_fun(
joint_q: wp.array(dtype=float),
loss: wp.array(dtype=float),
):
"""Loss function for the pendulum."""
wp.atomic_add(loss, 0, - joint_q[0] ** 2.)
def test(steps: int, silent: bool = False):
builder = wp.sim.ModelBuilder()
b = builder.add_body()
builder.add_shape_sphere(b, radius=1.0)
builder.add_joint_revolute(-1, b,
axis=wp.vec3(1., 0., 0.),
parent_xform=wp.transform(wp.vec3(0., 2., 0.), wp.quat_identity()),
child_xform=wp.transform(wp.vec3(0., 0., 0.), wp.quat_identity()),
limit_lower=-float('inf'),
limit_upper=float('inf'),
armature=0.
)
model = builder.finalize(requires_grad=True)
model.gravity = wp.vec3(0., 0., 0.)
states = list()
dt = 0.005
for _ in range(steps):
states.append(model.state(requires_grad=True))
wp.sim.eval_fk(model, states[0].joint_q, states[0].joint_qd, None, states[0])
actuation = wp.array([30.], dtype=float, requires_grad=True)
control = model.control()
control.joint_act = actuation
loss_val = wp.array([0.], dtype=float, requires_grad=True)
tape = wp.Tape()
with tape:
integrator = wp.sim.FeatherstoneIntegrator(model=model) # Inside tape to capture auxiliary variable computation
for step in range(steps - 1):
s0 = states[step]
s1 = states[step + 1]
s0.clear_forces()
integrator.simulate(model, s0, s1, dt, control=control)
wp.launch(loss_fun, dim=1, inputs=[states[-1].joint_q], outputs=[loss_val])
tape.backward(loss_val)
"""
L = - q^2 = - (q0 + t * qd0 + (1 / 2) alpha * t^2)^2 = - [(1 / 2) (tau / Ixx) t^2]^2
dL/dIxx = (1 / 2) tau^2 t^4 / (Ixx^3)
"""
def q(tau, Ixx, t):
return (1 / 2) * (tau / Ixx) * t ** 2
def grad(tau, Ixx, t):
return (1 / 2) * tau ** 2 * t ** 4 / (Ixx ** 3)
def grad_tau(tau, Ixx, t):
return - (1 / 2) * tau * t ** 4 / (Ixx ** 2)
if not silent:
print(f"{steps} steps")
print(f"Joint angle (simulated):", states[-1].joint_q.numpy()[0])
print(f"Joint angle (theory):", q(actuation.numpy()[0], model.body_inertia.numpy()[0, 0, 0], dt * steps))
g_wp = model.body_inertia.grad.numpy()[0, 0, 0]
g_a = grad(actuation.numpy()[0], model.body_inertia.numpy()[0, 0, 0], dt * steps)
print("Warp gradient I_xx:", g_wp)
print("Analytical gradient:", g_a)
print("Ratio:", g_wp / g_a)
gt_wp = actuation.grad.numpy()[0]
gt_a = grad_tau(actuation.numpy()[0], model.body_inertia.numpy()[0, 0, 0], dt * steps)
print("Warp gradient tau:", gt_wp)
print("Analytical gradient:", gt_a)
print("Ratio:", gt_wp / gt_a)
for i in (10, 50, 100, 200):
test(i)
print('\n' * 2)
Prints the following:
10 steps
Joint angle (simulated): 2.0143052e-05
Joint angle (theory): 2.2381164e-05
Warp gradient I_xx: 1.3851539e-11
Analytical gradient: 5.97925e-13
Ratio: 23.166016
Warp gradient tau: -2.7049502e-11
Analytical gradient: -3.3394433e-11
Ratio: 0.8100003
50 steps
Joint angle (simulated): 0.00054833863
Joint angle (theory): 0.0005595291
Warp gradient I_xx: 8.407696e-07
Analytical gradient: 3.7370312e-10
Ratio: 2249.833
Warp gradient tau: -2.0045011e-08
Analytical gradient: -2.087152e-08
Ratio: 0.96040016
100 steps
Joint angle (simulated): 0.0022157363
Joint angle (theory): 0.0022381165
Warp gradient I_xx: 0.00010363942
Analytical gradient: 5.97925e-09
Ratio: 17333.18
Warp gradient tau: -3.2729895e-07
Analytical gradient: -3.339443e-07
Ratio: 0.9801004
200 steps
Joint angle (simulated): 0.008907697
Joint angle (theory): 0.008952466
Warp gradient I_xx: 0.013010821
Analytical gradient: 9.5668e-08
Ratio: 135999.72
Warp gradient tau: -5.2898117e-06
Analytical gradient: -5.343109e-06
Ratio: 0.99002504
As can be seen, the joint angle after time t matches the analytical joint angle pretty accurately. The gradients w.r.t. applied torque are pretty close too (ignoring what I believe to be just numerical inaccuracies). However, the gradients for inertial properties (here: Ixx) do not match the analytical ones at all.
I'm happy to run more tests if someone has an idea what might cause this...
For the sake of completeness, here's a version of the code where analytical, autograd, and numeric gradients match, using the SemiImplicitIntegrator
from typing import Optional
import warp as wp
import warp.sim
from physics_kernels import recompute_inverse_inertial_properties # For details, see issue #585
@wp.kernel
def loss_fun(
body_q: wp.array(dtype=wp.transformf),
loss: wp.array(dtype=float),
):
"""Loss function for the pendulum."""
pos = wp.transform_get_translation(body_q[0])
rot = wp.transform_get_rotation(body_q[0])
angle = 2. * wp.atan2(wp.sqrt(rot[0] ** 2. + rot[1] ** 2. + rot[2] ** 2.), rot[3])
wp.atomic_add(loss, 0, angle)
def closure(steps: int, eps: Optional[float] = None, integrator: str = 'SemiImplicit'):
silent = eps is not None
requires_grad = eps is None
builder = wp.sim.ModelBuilder()
b = builder.add_body()
builder.add_shape_sphere(b, radius=1.0)
builder.add_joint_revolute(-1, b,
axis=wp.vec3(1., 0., 0.),
parent_xform=wp.transform(wp.vec3(0., 2., 0.), wp.quat_identity()),
child_xform=wp.transform(wp.vec3(0., 0., 0.), wp.quat_identity()),
limit_lower=-float('inf'),
limit_upper=float('inf'),
armature=0.
)
model = builder.finalize(requires_grad=requires_grad)
model.gravity = wp.vec3(0., 0., 0.)
dt = 0.005
if eps is not None:
I = model.body_inertia.numpy()
I[0, 0, 0] += eps
model.body_inertia = wp.array(I, dtype=wp.mat33)
states = list()
for _ in range(steps):
states.append(model.state(requires_grad=True))
wp.sim.eval_fk(model, states[0].joint_q, states[0].joint_qd, None, states[0])
actuation = wp.array([30.], dtype=float, requires_grad=True)
control = model.control()
control.joint_act = actuation
loss_val = wp.array([0.], dtype=float, requires_grad=True)
tape = wp.Tape()
with tape:
if integrator == 'SemiImplicit':
integrator = wp.sim.SemiImplicitIntegrator()
wp.launch(recompute_inverse_inertial_properties, dim=model.body_count,
inputs=[model.body_mass, model.body_inertia],
outputs=[model.body_inv_mass, model.body_inv_inertia]
)
else:
integrator = wp.sim.FeatherstoneIntegrator(model=model) # Inside tape to capture auxiliary variable computation
for step in range(steps - 1):
s0 = states[step]
s1 = states[step + 1]
s0.clear_forces()
integrator.simulate(model, s0, s1, dt, control=control)
wp.launch(loss_fun, dim=1, inputs=[states[-1].body_q], outputs=[loss_val])
if requires_grad:
tape.backward(loss_val)
"""
L = q = q0 + t * qd0 + (1 / 2) alpha * t^2 = (1 / 2) (tau / I) t^2
dL/dI = - (1 / 2) (tau / I^2) t^2
"""
def q(tau, Ixx, t):
return (1 / 2) * (tau / Ixx) * t ** 2
def grad(tau, Ixx, t):
return - (1 / 2) * tau * t ** 2 / (Ixx ** 2)
def grad_tau(tau, Ixx, t):
return (1 / 2) * t ** 2 / Ixx
if not silent:
print(f"{steps} steps")
print(f"Joint angle (simulated):", loss_val.numpy()[0])
print(f"Joint angle (theory):", q(actuation.numpy()[0], model.body_inertia.numpy()[0, 0, 0], dt * steps))
gt_wp = actuation.grad.numpy()[0]
gt_a = grad_tau(actuation.numpy()[0], model.body_inertia.numpy()[0, 0, 0], dt * steps)
print("Warp gradient tau:", gt_wp)
print("Analytical gradient tau:", gt_a)
print("Ratio:", gt_wp / gt_a)
g_wp = model.body_inertia.grad.numpy()[0, 0, 0]
g_a = grad(actuation.numpy()[0], model.body_inertia.numpy()[0, 0, 0], dt * steps)
print("Warp gradient I_xx:", g_wp)
print("Analytical gradient I_xx:", g_a)
print("Ratio:", g_wp / g_a)
return loss_val.numpy()[0]
epsilon = 1e-1
for i in (10, 50, 100, 200):
closure(i)
left = closure(i, -epsilon / 2)
right = closure(i, epsilon / 2)
print("Finite difference gradient m:", (right - left) / epsilon)
print('\n' * 2)