[BUG FIX] Fix per-environment and PD control commands in the differentiable rigid-body simulation.

genesis/engine/entities/rigid_entity/rigid_entity.py

@@ -59,23 +59,25 @@ def wrapper(self, *args, **kwargs):
             bound = sig.bind(self, *args, **kwargs)
             bound.apply_defaults()
             args_dict = dict(tuple(bound.arguments.items())[1:])
-            # Key the slot by (method, dofs subset) so same-step calls on distinct subsets (e.g. arm and gripper
-            # force control) each keep their own entry and gradient path; keyed by method alone, the second call
-            # would evict the first from the tape. Slices key directly when hashable (Python 3.12 onward) and
-            # resolve against the entity dof count otherwise.
-            dofs_idx_local = args_dict.get("dofs_idx_local")
-            if dofs_idx_local is None:
-                subset = None
-            elif isinstance(dofs_idx_local, slice):
-                if isinstance(dofs_idx_local, Hashable):
-                    subset = dofs_idx_local
+            # Key the slot by (method, dofs subset, envs subset) so same-step calls on distinct subsets (e.g. arm
+            # and gripper force control, or per-environment-group commands) each keep their own entry and gradient
+            # path; a coarser key would let the second call evict the first from the tape. Slices key directly when
+            # hashable (Python 3.12 onward) and resolve against their dimension size otherwise.
+            key = [fun.__name__]
+            for indices, n in (
+                (args_dict.get("dofs_idx_local"), self.n_dofs),
+                (args_dict.get("envs_idx"), self._solver._B),
+            ):
+                if indices is None:
+                    subset = None
+                elif isinstance(indices, slice):
+                    subset = indices if isinstance(indices, Hashable) else tuple(range(*indices.indices(n)))
+                elif isinstance(indices, torch.Tensor):
+                    subset = tuple(tensor_to_array(indices).reshape(-1).tolist())
                 else:
-                    subset = tuple(range(*dofs_idx_local.indices(self.n_dofs)))
-            elif isinstance(dofs_idx_local, torch.Tensor):
-                subset = tuple(tensor_to_array(dofs_idx_local).reshape(-1).tolist())
-            else:
-                subset = tuple(np.asarray(dofs_idx_local).reshape(-1).tolist())
-            self._update_tgt((fun.__name__, subset), args_dict)
+                    subset = tuple(np.asarray(indices).reshape(-1).tolist())
+                key.append(subset)
+            self._update_tgt(tuple(key), args_dict)
         return fun(self, *args, **kwargs)
 
     return wrapper
@@ -159,7 +161,15 @@ def __init__(
         self._load_model()
 
         # Initialize target variables and checkpoint
-        self._tgt_keys = ("set_pos", "set_quat", "set_dofs_velocity", "control_dofs_force")
+        self._tgt_keys = (
+            "set_pos",
+            "set_quat",
+            "set_dofs_velocity",
+            "control_dofs_force",
+            "control_dofs_velocity",
+            "control_dofs_position",
+            "control_dofs_position_velocity",
+        )
         self._tgt = dict()
         self._tgt_buffer = list()
         self._ckpt = dict()
@@ -1588,17 +1598,11 @@ def process_input(self, in_backward=False):
             # Do not update [tgt], as input information is finalized at this point
             self._update_tgt_while_set = False
 
-            match key[0]:
-                case "set_pos":
-                    self.set_pos(**data_kwargs)
-                case "set_quat":
-                    self.set_quat(**data_kwargs)
-                case "set_dofs_velocity":
-                    self.set_dofs_velocity(**data_kwargs)
-                case "control_dofs_force":
-                    self.control_dofs_force(**data_kwargs)
-                case _:
-                    gs.raise_exception(f"Invalid target key: {key[0]} not in {self._tgt_keys}")
+            # Every tracked setter replays uniformly from its taped kwargs, so dispatch by name (a rare legitimate
+            # use of getattr, on our own vetted key set).
+            if key[0] not in self._tgt_keys:
+                gs.raise_exception(f"Invalid target key: {key[0]} not in {self._tgt_keys}")
+            getattr(self, key[0])(**data_kwargs)
 
         self._tgt = dict()
         self._update_tgt_while_set = update_tgt_while_set
@@ -1637,6 +1641,15 @@ def process_input_grad(self):
                             self.set_dofs_force_grad, data_kwargs["dofs_idx_local"], data_kwargs["envs_idx"]
                         )
 
+                case "control_dofs_velocity" | "control_dofs_position" | "control_dofs_position_velocity":
+                    # PD control targets are replayed for primal correctness but have no input-gradient path.
+                    for target in (data_kwargs.get("position"), data_kwargs.get("velocity")):
+                        if isinstance(target, torch.Tensor) and target.requires_grad:
+                            gs.raise_exception(
+                                "Gradients with respect to PD control targets are not supported yet. Use "
+                                "'control_dofs_force' for differentiable control inputs."
+                            )
+
                 case _:
                     gs.raise_exception(f"Invalid target key: {key[0]} not in {self._tgt_keys}")
 
@@ -4011,6 +4024,7 @@ def control_dofs_force(self, force, dofs_idx_local=None, envs_idx=None):
         self._solver.control_dofs_force(force, dofs_idx, envs_idx)
 
     @gs.assert_built
+    @tracked
     def control_dofs_velocity(self, velocity, dofs_idx_local=None, envs_idx=None):
         """
         Set the PD controller's target velocity for the entity's dofs. This is used for velocity control.
@@ -4033,6 +4047,7 @@ def control_dofs_velocity(self, velocity, dofs_idx_local=None, envs_idx=None):
         self._solver.control_dofs_velocity(velocity, dofs_idx, envs_idx)
 
     @gs.assert_built
+    @tracked
     def control_dofs_position(self, position, dofs_idx_local=None, envs_idx=None):
         """
         Set the position controller's target position for the entity's dofs. The controller is a proportional term
@@ -4056,6 +4071,7 @@ def control_dofs_position(self, position, dofs_idx_local=None, envs_idx=None):
         self._solver.control_dofs_position(position, dofs_idx, envs_idx)
 
     @gs.assert_built
+    @tracked
     def control_dofs_position_velocity(self, position, velocity, dofs_idx_local=None, envs_idx=None):
         """
         Set a PD controller's target position and velocity for the entity's dofs. This is used for position control.

tests/grad/conftest.py

@@ -3,10 +3,12 @@
 import pytest
 
 
-def _add_hinge_arm(parent, body_name, pos, axis="0 1 0", **joint_kwargs):
+def _add_hinge_arm(parent, body_name, pos, axis="0 1 0", joint_pos=None, **joint_kwargs):
     """Add a 1-DOF hinge arm link (y-axis hinge by default, capsule geom, explicit inertia) and return its body
     element."""
     body = ET.SubElement(parent, "body", name=body_name, pos=pos)
+    if joint_pos is not None:
+        joint_kwargs["pos"] = joint_pos
     ET.SubElement(body, "joint", type="hinge", axis=axis, **joint_kwargs)
     ET.SubElement(body, "inertial", mass="0.5", pos="0.1 0 0", diaginertia="0.01 0.01 0.01")
     ET.SubElement(body, "geom", type="capsule", fromto="0 0 0 0.2 0 0", size="0.02", contype="0", conaffinity="0")
@@ -170,12 +172,14 @@ def grad_weld_pair():
     arm1 = _add_hinge_arm(worldbody, "arm1", "0 0 0", name="j1")
     _add_hinge_arm(arm1, "arm2", "0.2 0 0", name="j2", axis="1 0 0")
     equality = ET.SubElement(mjcf, "equality")
+    # relpose is the pose of body2 in body1's frame at the welded configuration; a wrong sign leaves the weld
+    # violated at rest and turns the scenario into a violent snap instead of a hold.
     ET.SubElement(
         equality,
         "weld",
         body1="arm2",
         body2="arm1",
-        relpose="0.2 0 0 1 0 0 0",
+        relpose="-0.2 0 0 1 0 0 0",
         solimp="0.95 0.99 0.001",
         solref="0.005 1",
     )
@@ -185,28 +189,37 @@ def grad_weld_pair():
 @pytest.fixture(scope="session")
 def grad_all_eq_fric():
     # Integration scene exercising every differentiated constraint group: frictionloss on j1, equality JOINT between
-    # j1 and j2, equality CONNECT between arm3 and arm4, equality WELD between arm5 and arm6. Each group acts on a
-    # disjoint pair of links so the constraint solver faces a well-posed system within every pair.
+    # j1 and j2, equality CONNECT between arm3 and arm4, equality WELD between arm6's nested child and arm5. Each
+    # group acts on a disjoint pair of links so the constraint solver faces a well-posed system within every pair.
+    # The weld chain nests a skew-axis hinge (see grad_weld_pair) so its rotation rows carry a nonzero velocity bias
+    # across two separate trees.
     mjcf = ET.Element("mujoco", model="all_eq_fric")
     worldbody = ET.SubElement(mjcf, "worldbody")
     _add_hinge_arm(worldbody, "arm1", "0 0 0", name="j1", frictionloss="0.5")
     for i_arm in range(2, 7):
-        _add_hinge_arm(worldbody, f"arm{i_arm}", f"0 {0.2 * (i_arm - 1):.1f} 0", name=f"j{i_arm}")
+        arm = _add_hinge_arm(worldbody, f"arm{i_arm}", f"0 {0.2 * (i_arm - 1):.1f} 0", name=f"j{i_arm}")
+        if i_arm == 6:
+            # Skew axis for a nonzero rotation-row velocity bias; the joint position offset keeps the hinge axis
+            # off the welded point so the weld position rows engage (an axis through it leaves the spin resisted
+            # only by the weak rotation rows and the chain winds up).
+            _add_hinge_arm(arm, "arm6b", "0.2 0 0", name="j6b", axis="1 0 0", joint_pos="0.05 0.02 0.06")
     equality = ET.SubElement(mjcf, "equality")
     ET.SubElement(
         equality, "joint", joint1="j1", joint2="j2", polycoef="0 1 0 0 0", solimp="0.95 0.99 0.001", solref="0.005 1"
     )
     ET.SubElement(
         equality, "connect", body1="arm3", body2="arm4", anchor="0.2 0 0", solimp="0.95 0.99 0.001", solref="0.005 1"
     )
+    # The softer solref keeps the 6-row weld on this three-dof chain clear of the stabilization stability
+    # boundary (the shared 0.005 timeconst is clamped to 2 * substep_dt, which is marginal here).
     ET.SubElement(
         equality,
         "weld",
-        body1="arm5",
-        body2="arm6",
-        relpose="0 -0.2 0 1 0 0 0",
+        body1="arm6b",
+        body2="arm5",
+        relpose="-0.2 -0.2 0 1 0 0 0",
         solimp="0.95 0.99 0.001",
-        solref="0.005 1",
+        solref="0.02 1",
     )
     return ET.tostring(mjcf, encoding="unicode")
 

tests/grad/test_rigid_collision.py

@@ -92,13 +92,17 @@ def _n_contacts(scene):
     # only guards the degenerate-FD case.
     fd_rtol = 1e-6 if precision == "64" else 2e-3
     fd_atol = 1e-14 if precision == "64" else 1e-6
-    base_force = np.array([0.0, 0.0, -8.0, 0.0, 0.0, 0.0])
+    # The small tangential preload keeps the contact manifold away from a gain / loss transition, which the
+    # tangential finite-difference perturbation would otherwise cross (the in-test contact-count guard trips).
+    base_force = np.array([0.5, 0.0, -8.0, 0.0, 0.0, 0.0])
     init_force = np.broadcast_to(base_force, (n_steps, 6)).copy()
 
     def settle(scene, obj):
+        # Settling under the same base force as the grad window keeps the contact manifold in its loaded
+        # configuration throughout, so the in-window contact-count guards hold.
         obj.set_dofs_position(gs.tensor(rest_dofs, dtype=gs.tc_float).sceneless())
         for _ in range(n_settle):
-            obj.control_dofs_force(0.0)
+            obj.control_dofs_force(base_force)
             scene.step()
 
     scene_ana, obj_ana = _build_contact_scene(shape, grad_capsule, requires_grad=True, show_viewer=show_viewer)
@@ -124,18 +128,23 @@ def loss_at(perturbed):
         scene_fd.reset()
         settle(scene_fd, obj_fd)
         for t in range(n_steps):
-            obj_fd.control_dofs_force(gs.tensor(perturbed[t], dtype=gs.tc_float))
+            obj_fd.control_dofs_force(perturbed[t])
             scene_fd.step()
             assert _n_contacts(scene_fd) == nc, "contact set changed under FD perturbation"
         return float((scene_fd.rigid_solver.get_state().qpos[0, :3] ** 2).sum().detach())
 
+    # Both the normal (z) and a tangential (x) force component: the tangential adjoint flows through the friction
+    # rows and their contact-frame chains, which the normal component alone leaves untested.
     for t in range(n_steps):
-        plus = init_force.copy()
-        plus[t, 2] += eps
-        minus = init_force.copy()
-        minus[t, 2] -= eps
-        fd_z = (loss_at(plus) - loss_at(minus)) / (2 * eps)
-        assert_allclose(ana[t, 2], fd_z, rtol=fd_rtol, atol=fd_atol, err_msg=f"contact force.grad mismatch at t={t}")
+        for i_d in (2, 0):
+            plus = init_force.copy()
+            plus[t, i_d] += eps
+            minus = init_force.copy()
+            minus[t, i_d] -= eps
+            fd = (loss_at(plus) - loss_at(minus)) / (2 * eps)
+            assert_allclose(
+                ana[t, i_d], fd, rtol=fd_rtol, atol=fd_atol, err_msg=f"contact force.grad mismatch at t={t} dof={i_d}"
+            )
 
 
 @pytest.mark.required
@@ -341,7 +350,9 @@ def test_constraint_solver_backward_matches_fd(monkeypatch):
     from genesis.engine.solvers.rigid.constraint.solver import func_solve_body, func_solve_init
     from genesis.engine.solvers.rigid.rigid_solver import kernel_step_1
 
-    # fp64 is required: the FD perturbation must be small enough for a reliable estimate, which fp32 cannot resolve
+    # fp64 is required: the parameter-side inputs (aref, efc_D, jac, mass) move the loss by far less than its own
+    # fp32 resolution at any usable step size, so their finite-difference deltas drown at fp32
+
     scene = gs.Scene(
         sim_options=gs.options.SimOptions(
             requires_grad=True,
@@ -413,7 +424,9 @@ def compute_loss(input_mass, input_jac, input_aref, input_efc_D, input_force):
         updated_acc_smooth = np.linalg.solve(input_mass[..., 0], input_force[..., 0])
         rigid_solver.dyn_state.dofs.acc_smooth.from_numpy(updated_acc_smooth[..., None])
         constraint_solver.resolve()
-        return ((qd_to_torch(constraint_solver.qacc) - target_qacc) ** 2).mean()
+        # Collapse to a Python scalar right away: the reduction reads a zero-copy view of qacc, and a deferred
+        # evaluation (MPS is lazy) would otherwise run only after the next call overwrites the buffer in place.
+        return float(((qd_to_torch(constraint_solver.qacc) - target_qacc) ** 2).mean())
 
     init_input_mass = qd_to_numpy(rigid_solver.rigid_info.mass_mat, copy=True)
     init_input_jac = qd_to_numpy(constraint_solver.constraint_state.jac, copy=True)
@@ -475,5 +488,5 @@ def compute_loss(input_mass, input_jac, input_aref, input_efc_D, input_force):
             loss_p = compute_loss(**{**inputs, key: init_x + rand_dx * fd_eps})
             loss_m = compute_loss(**{**inputs, key: init_x - rand_dx * fd_eps})
             dL_fd = (loss_p - loss_m) / (2 * fd_eps)
-            error += (dL - dL_fd).abs() / max(abs(dL), abs(dL_fd), gs.EPS)
+            error += abs(dL - dL_fd) / max(abs(dL), abs(dL_fd), gs.EPS)
         assert_allclose(error / trials, 0.0, atol=1e-4)