NVIDIA
[BUG] Featherstone with Free 6DOF joint results in incorrect simulation and gradients
Hello,
My Setup
I'm trying to simulate the bounce example but with the ball being a rigid body sphere instead of a particle. This setup works well with the Semi-Implicit Euler integrator, and the gradients seem to be well conditioned but I run into several issues with the Featherstone integrator. The only difference in the setup is that I add a free_joint to the ball to be able to simulate it in generalized coordinates.
Bug Description
The output states for both integrators differ with the featherstone states consisting of an angular rotation not present in the Euler state. There is no friction and no reason for the ball to incur any kind of angular velocity or change in orientation. In addition to this, the gradients differ significantly with the Featherstone gradients blowing up to an order of 10e25
This might explain #570
Semi Implicit Euler final states after simulation
body_q = [1.0218604, 1.9264, 0.0000, 0.0000, 0.0000, 0.0000, 1.0000]
body_qd = [ 0.0000, 0.0000, 0.0000, -5.0017724, 2.6256406, 0.0000]
Featherstone final states after simulation
joint_q = [1.1582181, 2.086495, 0.0000, 0.0000, 0.0000, 0.6073468, 0.7944369]
joint_qd = [ 0.0000, 0.0000, 3.5609848, 2.5887039, -0.48731723, 0.0000]
Semi Implicit Euler gradient of initial velocity state wrt loss
[-2.4142773 -0.31994405]
Featherstone gradient of initial velocity state wrt loss
[2.3695136e+25 -2.5367573e+25]
Attempts at Solving
I have observed that when I increase the default update_mass_matrix_every=1 parameter of the Featherstone integrator, the gradients become smaller (but still very large), indicating a potential bug in the way the mass matrix is handled.
In case the following information is useful: If I comment out the #following lines in integrator_featherstone.py in the adjoint computation for the dense_solve function, the gradient seems to be more reasonable for the first iteration, but on multiple training iteration runs, the ball goes through the floor.
https://github.com/NVIDIA/warp/blob/23eb9ea711c5494ddeb2f51950a3d830a5864ed6/warp/sim/integrator_featherstone.py#L1428-L1430
I was also wondering if in the free-joint integration of jcalc_integrate, is it correct to compute the translation velocity as https://github.com/NVIDIA/warp/blob/23eb9ea711c5494ddeb2f51950a3d830a5864ed6/warp/sim/integrator_featherstone.py#L604
Since v_s already represents the linear velocity of the joint-frame origin (and after quaternion rotation becomes the world-frame velocity of that same point), shouldn’t we simply use
dpdt_s = v_s;
to avoid “sweeping” the origin around as if it were offset from the center of rotation?
Reproducing
Click to expand — Code
# Copyright (c) 2022 NVIDIA CORPORATION. All rights reserved.
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto. Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.
import numpy as np
import matplotlib.pyplot as plt
import warp as wp
import warp.sim
import warp.sim.render
import torch
import time
# from utils.plotter import Plotter
# from utils.create_terrain import generate_terrain
@wp.kernel
def loss_kernel(pos: wp.array(dtype=wp.vec3), target: wp.vec3, loss: wp.array(dtype=float)):
# distance to target
delta = pos[0] - target
loss[0] = wp.dot(delta, delta)
@wp.kernel
def loss_kernel_rigid(pos: wp.array(dtype=wp.transform), target: wp.vec3, loss: wp.array(dtype=float)):
# distance to target
p = wp.transform_get_translation(pos[0])
delta = p - target
loss[0] = wp.dot(delta, delta)
@wp.kernel
def loss_kernel_rigid_featherstone(pos: wp.array(dtype=float), target: wp.vec3, loss: wp.array(dtype=float)):
# distance to target
p = wp.vec3(pos[0], pos[1], pos[2])
delta = p - target
loss[0] = wp.dot(delta, delta)
@wp.kernel
def step_kernel(x: wp.array(dtype=wp.vec3), grad: wp.array(dtype=wp.vec3), alpha: float):
# gradient descent step
tid = wp.tid()
x[tid] = x[tid] - grad[tid] * alpha
@wp.kernel
def step_kernel_rigid(x: wp.array(dtype=wp.spatial_vector), grad: wp.array(dtype=wp.spatial_vector), alpha: float):
tid = wp.tid()
g = wp.spatial_bottom(grad[tid])
# Clamp each element of the gradient to the range [-1, 1]
# g = wp.vec3(wp.clamp(g[0], -1.0, 1.0), wp.clamp(g[1], -1.0, 1.0), wp.clamp(g[2], -1.0, 1.0))
lin_vel = wp.spatial_bottom(x[tid])
ang_vel = wp.spatial_top(x[tid])
# gradient descent step
lin_vel = lin_vel - g * alpha
x[tid] = wp.spatial_vector(ang_vel, lin_vel)
@wp.kernel
def step_kernel_rigid_featherstone(x: wp.array(dtype=float), grad: wp.array(dtype=float), alpha: float):
tid = wp.tid()
# wp.print(x[tid])
# wp.print(grad[tid])
xd = wp.vec3(x[tid + 3], x[tid + 4], x[tid + 5])
gd = wp.vec3(grad[tid + 3], grad[tid + 4], grad[tid + 5])
# gradient descent step
xd = xd - gd * alpha
x[tid + 3] = xd[0]
x[tid + 4] = xd[1]
x[tid + 5] = xd[2]
@wp.kernel
def record_contact_kernel(contact_count: wp.array(dtype=int),
contact_out: wp.array(dtype=int),
idx: int):
# Assume soft_contact_count[0] contains the number of contacts detected.
contact_out[idx] = contact_count[0]
class Bounce:
sim_name = "Bounce-Sim"
def __init__(
self,
mode,
integrator="semi-implicit",
verbose=False,
render=True,
dynamics_params=None,
profile=True,
wave_wall=False,
bundle=False,
):
self.verbose = verbose
self.mode = mode
self.integrator_type = integrator
self.profile = profile
self.bundle = bundle
assert self.mode in ["particle", "body"], "mode should be either 'particle' or 'body'"
assert self.integrator_type in ["xpbd", "featherstone", "semi-implicit"], "integrator should be either 'xpbd', 'featherstone' or 'semi-implicit'"
# seconds
sim_duration = 0.6
# control frequency
fps = 60
self.frame_dt = 1.0 / fps
# ep
self.ep_len = int(sim_duration / self.frame_dt)
# sim frequency
self.sim_substeps = 20
self.sim_steps = self.ep_len * self.sim_substeps
self.sim_dt = self.frame_dt / self.sim_substeps
if self.verbose:
print('sim steps:', self.sim_steps)
print('sim dt:', self.sim_dt)
self.iter = 0
self.render_time = 0.0
self.render_traj = True
# self.plotter = Plotter(self.sim_steps, sim_duration)
self.train_rate = 0.02
if dynamics_params is None:
rigid_ke = 1e4
# rigid_ke = 1e6
rigid_kf = 0.0
rigid_kd = 0.0e2
# rigid_kd = 1.0e3
rigid_mu = 0.0
# rigid_mu = 0.75
rigid_restitution = 1.0
soft_ke = 1.0e4
soft_kf = 0.0
soft_kd = 0.0e2
soft_mu = 0.0
box_ke = 1.0e4
box_kf = 0.0
box_kd = 0.0e2
box_mu = 0.0
box_restitution = rigid_restitution
ground_ke = rigid_ke
ground_kf = rigid_kf
ground_kd = rigid_kd
ground_mu = rigid_mu
ground_restitution = rigid_restitution
else:
rigid_ke = dynamics_params['rigid_ke']
rigid_kf = dynamics_params['rigid_kf']
rigid_kd = dynamics_params['rigid_kd']
rigid_mu = dynamics_params['rigid_mu']
rigid_restitution = dynamics_params['rigid_restitution']
soft_ke = dynamics_params['soft_ke']
soft_kf = dynamics_params['soft_kf']
soft_kd = dynamics_params['soft_kd']
soft_mu = dynamics_params['soft_mu']
box_ke = dynamics_params['box_ke']
box_kf = dynamics_params['box_kf']
box_kd = dynamics_params['box_kd']
box_mu = dynamics_params['box_mu']
box_restitution = dynamics_params['box_restitution']
ground_ke = dynamics_params['ground_ke']
ground_kf = dynamics_params['ground_kf']
ground_kd = dynamics_params['ground_kd']
ground_mu = dynamics_params['ground_mu']
ground_restitution = dynamics_params['ground_restitution']
builder = wp.sim.ModelBuilder()
# Add a particle or a body
if mode == "particle":
builder.add_particle(pos=wp.vec3(-0.5, 1.0, 0.0), vel=wp.vec3(5.0, -5.0, 0.0), mass=1.0)
else:
sphere_mass = 1.0
sphere_volume = 4.0 / 3.0 * np.pi * 0.1 ** 3
sphere_density = sphere_mass / sphere_volume
if self.verbose:
print('Sphere density:', sphere_density)
# # Dummy shape due to rendering bug
# builder.add_shape_box(body=-1, pos=wp.vec3(400.0, 1.0, 0.0), hx=0.25, hy=1.0, hz=1.0)
sphere = builder.add_body(origin=wp.transform((-0.5, 1.0, 0.0), wp.quat_identity()), m=1.0)
builder.add_shape_sphere(body=sphere, radius=0.1, density=sphere_density, has_ground_collision=True,
ke=rigid_ke, kf=rigid_kf, kd=rigid_kd, mu=rigid_mu, restitution=rigid_restitution,
is_solid=True)
# Featherstone needs a floating joint
if self.integrator_type == "featherstone":
builder.add_joint_free(child=sphere, armature=0.5)
# Add initial velocity to the body
if self.integrator_type != "featherstone":
builder.body_qd[0] = (0.0, 0.0, 0.0, 5.0, -5.0, 0.0)
else:
builder.joint_qd[3:6] = [5.0, -5.0, 0.0]
builder.joint_q[0:3] = [-0.5, 1.0, 0.0]
builder.body_qd[0] = (0.0, 0.0, 0.0, 5.0, -5.0, 0.0)
# Add wall or rough wall
if wave_wall:
pass
# # Add rough wall
# rough_wall = generate_terrain(num_waves=2, amplitude=0.1, width=256, length=256,
# vertical_scale=2 / 256, horizontal_scale=2 / 256)
# builder.add_shape_mesh(body=-1, pos=wp.vec3(2.0, 0.0, -1.0), rot=wp.quat_rpy(0.0, -np.pi / 2, 0.0),
# mesh=rough_wall, ke=box_ke, kf=box_kf, kd=box_kd, mu=box_mu,
# restitution=box_restitution)
else:
# Add Box
builder.add_shape_box(body=-1, pos=wp.vec3(2.0, 1.0, 0.0), hx=0.25, hy=1.0, hz=1.0, ke=box_ke, kf=box_kf, kd=box_kd, mu=box_mu, restitution=box_restitution)
# Ground
builder.set_ground_plane(ke=ground_ke, kd=ground_kd, kf=ground_kf, mu=ground_mu, restitution=ground_restitution)
# use `requires_grad=True` to create a model for differentiable simulation
self.model = builder.finalize(requires_grad=True)
self.model.ground = True
self.model.soft_contact_ke = soft_ke
self.model.soft_contact_kf = soft_kf
self.model.soft_contact_kd = soft_kd
self.model.soft_contact_mu = soft_mu
self.model.soft_contact_margin = 10.0
self.model.soft_contact_restitution = 1.0
if self.integrator_type == "xpbd":
self.integrator = wp.sim.XPBDIntegrator(enable_restitution=True)
elif self.integrator_type == "featherstone":
self.integrator = wp.sim.FeatherstoneIntegrator(self.model, angular_damping=0.05, friction_smoothing=1.0)
else:
self.integrator = wp.sim.SemiImplicitIntegrator(**dict(angular_damping=0.05, friction_smoothing=1.0))
self.target = (-2.0, 1.5, 0.0)
self.loss = wp.zeros(1, dtype=wp.float32, requires_grad=True)
# allocate sim states for trajectory
self.states = []
for _i in range(self.sim_steps + 1):
self.states.append(self.model.state(requires_grad=True))
# Allocate array for contact count
self.contact_states = wp.zeros((self.sim_steps + 1,), dtype=wp.int32)
if self.verbose:
if self.integrator_type == "featherstone":
print('Initial state position:', self.states[0].joint_q)
print('Initial state velocity:', self.states[0].joint_qd)
else:
print('Initial state position:', self.states[0].body_q)
print('Initial state velocity:', self.states[0].body_qd)
# one-shot contact creation (valid if we're doing simple collision against a constant normal plane)
# wp.sim.collide(self.model, self.states[0])
# Update state based on joint info
wp.sim.eval_fk(self.model, self.states[0].joint_q, self.states[0].joint_qd, None, self.states[0])
if render:
# self.renderer = wp.sim.render.SimRenderer(self.model, stage_path, scaling=1.0)
self.renderer = wp.sim.render.SimRendererOpenGL(self.model, self.sim_name, scaling=1.0)
else:
# self.renderer = Renderer(self.model, uses_generalized_coordinates=uses_generalized_coordinates, render_mode=RenderMode.OPENGL)
self.renderer = None
# capture forward/backward passes
self.use_cuda_graph = wp.get_device().is_cuda
# self.use_cuda_graph = False
self.device = wp.get_device()
if self.use_cuda_graph:
with wp.ScopedCapture() as capture:
self.tape = wp.Tape()
with self.tape:
self.forward()
self.tape.backward(self.loss)
self.graph = capture.graph
def forward(self):
# run control loop
for i in range(self.sim_steps):
self.states[i].clear_forces()
if self.mode == "body":
wp.sim.collide(self.model, self.states[i])
# Record the number of contacts at this simulation step.
# wp.launch(record_contact_kernel, dim=1,
# inputs=[self.model.rigid_contact_count, self.contact_states, i])
self.integrator.simulate(self.model, self.states[i], self.states[i + 1], self.sim_dt)
# compute loss on final state
if self.mode == "particle":
wp.launch(loss_kernel, dim=1, inputs=[self.states[-1].particle_q, self.target, self.loss])
elif self.integrator_type == "featherstone":
wp.launch(loss_kernel_rigid_featherstone, dim=1, inputs=[self.states[-1].joint_q, self.target, self.loss])
else:
wp.launch(loss_kernel_rigid, dim=1, inputs=[self.states[-1].body_q, self.target, self.loss])
return self.loss
def step(self):
with wp.ScopedTimer("step", active=self.profile):
if self.use_cuda_graph:
wp.capture_launch(self.graph)
else:
self.tape = wp.Tape()
with self.tape:
self.forward()
self.tape.backward(self.loss)
# gradient descent step
if self.mode == "particle":
x = self.states[0].particle_qd
wp.launch(step_kernel, dim=len(x), inputs=[x, x.grad, self.train_rate])
x_grad = self.tape.gradients[self.states[0].particle_qd]
elif self.integrator_type == "featherstone":
final_x = self.states[-1].joint_q
final_xd = self.states[-1].joint_qd
# x = self.states[0].body_qd
x = self.states[0].joint_qd
wp.launch(step_kernel_rigid_featherstone, dim=len(x), inputs=[x, x.grad, self.train_rate])
# x_grad = self.tape.gradients[self.states[0].body_qd]
x_grad = self.tape.gradients[self.states[0].joint_qd]
else:
final_x = self.states[-1].body_q
final_xd = self.states[-1].body_qd
x = self.states[0].body_qd
wp.launch(step_kernel_rigid, dim=len(x), inputs=[x, x.grad, self.train_rate])
x_grad = self.tape.gradients[self.states[0].body_qd]
if self.verbose:
print(f"Iter: {self.iter} Loss: {self.loss}")
print(f" final x: {final_x} \n final xd: {final_xd} \n grad initial xd: {x_grad}")
# # Store the first two elements of the gradient, the third is zero
# if self.mode == "particle":
# self.plotter.gradients.append(x_grad.numpy().flatten()[:2])
# self.plotter.velocities.append(x.numpy().flatten()[:2])
# else:
# self.plotter.gradients.append(x_grad.numpy().flatten()[3:5])
# self.plotter.velocities.append(x.numpy().flatten()[3:5])
# self.plotter.losses.append(self.loss.numpy()[0])
#
# # store all the gradients for the first step
# if self.iter == 0:
# for i in range(self.sim_steps + 1):
# if self.mode == "particle":
# x = self.states[i].particle_q
# xd = self.states[i].particle_qd
# # wp.launch(step_kernel, dim=len(x), inputs=[xd, xd.grad, self.train_rate])
# x_grad = self.tape.gradients[self.states[i].particle_qd]
# self.plotter.first_step_gradients.append(x_grad.numpy().flatten()[:2])
# self.plotter.first_step_velocities.append(xd.numpy().flatten()[:2])
# self.plotter.first_step_positions.append(x.numpy().flatten()[:2])
# elif self.integrator_type == "featherstone":
# # x = self.states[i].body_q
# # xd = self.states[i].body_qd
# x = self.states[i].joint_q
# xd = self.states[i].joint_qd
# # wp.launch(step_kernel_rigid_featherstone, dim=len(x), inputs=[xd, xd.grad, self.train_rate])
# # x_grad = self.tape.gradients[self.states[i].body_qd]
# x_grad = self.tape.gradients[self.states[i].joint_qd]
# self.plotter.first_step_gradients.append(x_grad.numpy().flatten()[3:5])
# self.plotter.first_step_velocities.append(xd.numpy().flatten()[3:5])
# self.plotter.first_step_positions.append(x.numpy().flatten()[:2])
# else:
# x = self.states[i].body_q
# xd = self.states[i].body_qd
# # wp.launch(step_kernel_rigid, dim=len(x), inputs=[xd, xd.grad, self.train_rate])
# x_grad = self.tape.gradients[self.states[i].body_qd]
# self.plotter.add("first_step_gradients", x_grad.numpy().flatten()[3:5])
# self.plotter.add("first_step_velocities", xd.numpy().flatten()[3:5])
# self.plotter.add("first_step_positions", x.numpy().flatten()[:2])
self.iter = self.iter + 1
def zero_grad(self):
self.tape.zero()
@staticmethod
def synchronize():
wp.synchronize()
def render(self):
if self.renderer is None:
return
with wp.ScopedTimer("render", active=self.profile):
self.render_traj = self.iter % 16 == 0
# if self.render_traj:
# draw trajectory
# if self.render_traj:
if self.mode == "particle":
traj_verts = [self.states[0].particle_q.numpy()[0].tolist()]
elif self.integrator_type == "featherstone":
traj_verts = [self.states[0].joint_q.numpy().tolist()[0:3]]
else:
traj_verts = [self.states[0].body_q.numpy()[0].tolist()[0:3]]
for i in range(0, self.sim_steps, self.sim_substeps):
# if self.render_traj:
if self.mode == "particle":
traj_verts.append(self.states[i].particle_q.numpy()[0].tolist())
elif self.integrator_type == "featherstone":
traj_verts.append(self.states[i].joint_q.numpy().tolist()[0:3])
else:
traj_verts.append(self.states[i].body_q.numpy()[0].tolist()[0:3])
self.renderer.begin_frame(self.render_time)
self.renderer.render(self.states[i])
self.renderer.render_box(
pos=self.target,
rot=wp.quat_identity(),
extents=(0.1, 0.1, 0.1),
name="target",
color=(0.0, 0.0, 0.0),
)
self.renderer.render_line_strip(
vertices=traj_verts,
color=wp.render.bourke_color_map(0.0, 7.0, self.loss.numpy()[0]),
radius=0.02,
name=f"traj_{self.iter - 1}",
)
self.renderer.end_frame()
self.render_time += self.frame_dt
# else:
# for i in range(0, self.sim_steps, self.sim_substeps):
# self.renderer.begin_frame(self.render_time)
# self.renderer.render(self.states[i])
# self.renderer.render_box(
# pos=self.target,
# rot=wp.quat_identity(),
# extents=(0.1, 0.1, 0.1),
# name="target",
# color=(0.0, 0.0, 0.0),
# )
# self.renderer.end_frame()
# # self.render_time += self.frame_dt
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument("--device", type=str, default=None, help="Override the default Warp device.")
parser.add_argument("--train_iters", type=int, default=250, help="Total number of training iterations.")
parser.add_argument("--verbose", action="store_true", help="Print out additional status messages during execution.")
parser.add_argument("--mode", type=str, default="body", help="particle or body")
parser.add_argument("--plot", action='store_true', help="Plot gradient and velocity evolution over time")
parser.add_argument("--integrator", type=str, default="semi-implicit", help="Type of integrator to use")
parser.add_argument("--wave_wall", action='store_true', default=False, help="Use a wave wall or flat wall")
parser.add_argument("--bundle", action='store_true', default=False, help="Use bundle optimization")
args = parser.parse_known_args()[0]
with wp.ScopedDevice(args.device):
bounce_env = Bounce(integrator=args.integrator, verbose=args.verbose, mode=args.mode, render=True, profile=False, wave_wall=args.wave_wall,
bundle=args.bundle)
# replay and optimize
for i in range(args.train_iters):
bounce_env.step()
bounce_env.render()
bounce_env.zero_grad()
# Plot collected gradients if flag is enabled
if args.plot:
bounce_env.plotter.plot()
To run with Semi-implicit Euler
python bounce_warp --train_iters 1 --integrator semi-implicit --verbose
To run with Featherstone
python bounce_warp --train_iters 1 --integrator featherstone --verbose
Any help on this would be great, Thanks!
System Information
warp-lang: 1.7.0 OS: Ubuntu 24.04 GPU: RTX 4070
Hi @dyumanaditya, I truly appreciate the level of detail in this bug report and I would point others to it as a gold standard for how to file issues, but I don't expect that we will have time to look into this issue for a while. @momo-van posted an announcement yesterday about the in-progress migration of warp.sim into https://github.com/newton-physics/newton here https://github.com/NVIDIA/warp/discussions/735 The good news is that you should be able to file this issue in the Newton repository, where a fix would be delivered if one is identified.
FYI this might be helpful, here's how I usually create a free-joint rigid object for Featherstone in Rewarped:
def create_cube(self, builder):
p = wp.vec3(0.5, 0.0, 0.055)
q = wp.quat_identity()
object_start_pose = wp.transform(p, q)
box_size = dict(
hx=0.026,
hy=0.026,
hz=0.026,
)
builder.add_articulation()
b = builder.add_body(name="body_object")
builder.add_shape_box(
body=b,
**box_size,
mu=0.5,
restitution=0.0,
is_solid=True,
has_ground_collision=True,
has_shape_collision=True,
)
M = len(builder.joint_q)
builder.add_joint_free(parent=-1, child=b, name="joint_object")
builder.joint_q[M : M + 3] = object_start_pose.p
builder.joint_q[M + 3 : M + 7] = object_start_pose.q
In particular, use joint_q to hold global coordinates, and don't add the xform to body_q. Might work for your case or in place of particles like in rewarped/envs/warp_examples/bounce.py
Thanks for the code snippet @etaoxing although this doesn't seem to solve the problem. I think there is something wrong with the underlying implementation and Featherstone seems to produce incorrect gradients even when a ball is attached to a prismatic joint.
I'm not too sure how to fix that. Please let me know if you ran into anything similar when you worked on rewarped
Unfortunately, I can't contribute a solution, but this matches my experience of having problems with autodifferentiating through the Featherstone integrator, as reported in #650 and #655. While I mostly encountered problems when computing gradients with respect to inertial parameters, your information that gradients seem to (initially) be more reasonable when setting update_mass_matrix_every=1 could indicate that there is a shared common issue somewhere in how inertial properties of the bodies are computed/used.
Sorry we weren't able to address this issue before removing warp.sim from the codebase. As of Warp v1.10, the warp.sim module has been removed. It was deprecated back in v1.8 and has been superseded by the Newton physics engine, which is now maintained as an independent Linux Foundation project with a redesigned API focused on robotics and robot learning.
If this issue is still relevant to your work, we'd encourage you to check out Newton. There's a migration guide available that should help with the transition. You can open issues in the Newton repository or ask questions in the Newton Discussions section. Many of us from the Warp team are also involved with Newton and happy to help there.