[FEATURE] Improve robustness of differentiable rigid body simulation.

genesis/engine/entities/rigid_entity/rigid_entity.py

@@ -142,7 +142,7 @@ def __init__(
         self._load_model()
 
         # Initialize target variables and checkpoint
-        self._tgt_keys = ("pos", "quat", "qpos", "dofs_velocity")
+        self._tgt_keys = ("pos", "quat", "qpos", "dofs_velocity", "control_dofs_force")
         self._tgt = dict()
         self._tgt_buffer = list()
         self._ckpt = dict()
@@ -1536,6 +1536,8 @@ def process_input(self, in_backward=False):
                     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} not in {self._tgt_keys}")
 
@@ -1566,6 +1568,14 @@ def process_input_grad(self):
                         velocity._backward_from_qd(
                             self.set_dofs_velocity_grad, data_kwargs["dofs_idx_local"], data_kwargs["envs_idx"]
                         )
+
+                case "control_dofs_force":
+                    force = data_kwargs.pop("force")
+                    if force.requires_grad:
+                        force._backward_from_qd(
+                            self.set_dofs_force_grad, data_kwargs["dofs_idx_local"], data_kwargs["envs_idx"]
+                        )
+
                 case _:
                     gs.raise_exception(f"Invalid target key: {key} not in {self._tgt_keys}")
 
@@ -3879,6 +3889,11 @@ def set_dofs_velocity_grad(self, dofs_idx_local, envs_idx, velocity_grad):
         dofs_idx = self._get_global_idx(dofs_idx_local, self.n_dofs, self._dof_start, unsafe=True)
         self._solver.set_dofs_velocity_grad(dofs_idx, envs_idx, velocity_grad.data)
 
+    @gs.assert_built
+    def set_dofs_force_grad(self, dofs_idx_local, envs_idx, force_grad):
+        dofs_idx = self._get_global_idx(dofs_idx_local, self.n_dofs, self._dof_start, unsafe=True)
+        self._solver.set_dofs_force_grad(dofs_idx, envs_idx, force_grad.data)
+
     # ------------------------------------------------------------------------------------
     # ----------------------------- DOF property setters ---------------------------------
     # ------------------------------------------------------------------------------------
@@ -3912,6 +3927,7 @@ def set_dofs_position(self, position, dofs_idx_local=None, envs_idx=None, *, zer
     # ------------------------------------------------------------------------------------
 
     @gs.assert_built
+    @tracked
     def control_dofs_force(self, force, dofs_idx_local=None, envs_idx=None):
         """
         Control the entity's dofs' motor force. This is used for force/torque control.

genesis/engine/scene.py

@@ -995,13 +995,16 @@ def reset(self, state: SimState | None = None, envs_idx=None):
         self._reset(state, envs_idx=envs_idx)
         self._recorder_manager.reset(envs_idx)
 
-    def _reset(self, state: SimState | None = None, *, envs_idx=None):
+    def _reset(self, state: SimState | None = None, *, envs_idx=None, keep_init: bool = False):
         if self._is_built:
             if state is None:
                 state = self._init_state
             else:
                 assert isinstance(state, SimState), "state must be a SimState object"
-                self._init_state = state
+                # keep_init=True restores the state while leaving the registered init untouched, so a later bare
+                # reset() still rewinds to the true initial state.
+                if not keep_init:
+                    self._init_state = state
             self._sim.reset(state, envs_idx)
         else:
             self._init_state = self._get_state()
@@ -1022,6 +1025,40 @@ def _reset(self, state: SimState | None = None, *, envs_idx=None):
     def _reset_grad(self):
         self._backward_ready = True
 
+    @gs.assert_built
+    def backward(self, loss: torch.Tensor, *args, **kwargs):
+        """
+        Differentiates `loss` through the recorded rollout and restores the pre-backward physics state.
+
+        Unrolling the gradient tape rewinds the physics state to step 0, so this method snapshots the current state
+        first, runs the backward pass, and restores the snapshot afterwards. The scene then sits at the same physics
+        state as before the call, with gradients populated and forward / backward re-armed, ready to continue the
+        rollout or to be reset. The registered initial state (`reset()` with no argument) is preserved.
+
+        Parameters
+        ----------
+        loss : torch.Tensor
+            Scalar loss to differentiate. Extra positional and keyword arguments (e.g. `gradient`, `retain_graph`)
+            are forwarded to `torch.autograd.backward`.
+
+        Returns
+        -------
+        snapshot : SimState
+            The physics state the scene was restored to.
+        """
+        # Snapshot the current state before the gradient-tape unroll rewinds physics to step 0.
+        snapshot = self.get_state()
+        # The sim unroll (self._backward) re-enters the torch graph from each step's queried states, so the graph
+        # must survive the initial autograd pass.
+        kwargs.setdefault("retain_graph", True)
+        # The functional torch.autograd.backward fills torch and queried-state grads while leaving the sim unroll to
+        # the explicit self._backward call below, keeping gs.Tensor.backward's automatic scene._backward out of it.
+        torch.autograd.backward(loss, *args, **kwargs)
+        self._backward()
+        # keep_init=True preserves the registered initial state so a later bare reset() still rewinds to it.
+        self._reset(snapshot, keep_init=True)
+        return snapshot
+
     def _get_state(self):
         return self._sim.get_state()
 

genesis/engine/solvers/kinematic_solver.py

@@ -46,6 +46,7 @@
     kernel_set_dofs_position,
     kernel_set_dofs_velocity,
     kernel_set_dofs_velocity_grad,
+    kernel_set_dofs_force_grad,
     kernel_set_dofs_zero_velocity,
     kernel_set_links_pos,
     kernel_set_links_quat,
@@ -1012,6 +1013,14 @@ def set_dofs_velocity_grad(self, dofs_idx, envs_idx, velocity_grad):
             velocity_grad_ = velocity_grad_.unsqueeze(0)
         kernel_set_dofs_velocity_grad(dofs_idx, envs_idx, velocity_grad_, self.dyn_state, self.rigid_config)
 
+    def set_dofs_force_grad(self, dofs_idx, envs_idx, force_grad):
+        force_grad_, dofs_idx, envs_idx = self._sanitize_io_variables(
+            force_grad, dofs_idx, self.n_dofs, "dofs_idx", envs_idx, skip_allocation=True
+        )
+        if self.n_envs == 0:
+            force_grad_ = force_grad_.unsqueeze(0)
+        kernel_set_dofs_force_grad(dofs_idx, envs_idx, force_grad_, self.dyn_state, self.rigid_config)
+
     @mutates(StateChange.GEOMETRY, links=MutatedLinks.ARTICULATED)
     def set_dofs_position(self, position, dofs_idx=None, envs_idx=None):
         position, dofs_idx, envs_idx = self._sanitize_io_variables(

genesis/engine/solvers/rigid/abd/accessor.py

@@ -819,6 +819,20 @@ def kernel_set_dofs_velocity_grad(
         dyn_state.dofs.vel.grad[dofs_idx[i_d_], envs_idx[i_b_]] = 0.0
 
 
+@qd.kernel(fastcache=True)
+def kernel_set_dofs_force_grad(
+    dofs_idx: qd.types.ndarray(),
+    envs_idx: qd.types.ndarray(),
+    force_grad: qd.types.ndarray(),
+    dyn_state: array_class.DynState,
+    rigid_config: qd.template(),
+):
+    qd.loop_config(serialize=rigid_config.para_level < gs.PARA_LEVEL.ALL)
+    for i_d_, i_b_ in qd.ndrange(dofs_idx.shape[0], envs_idx.shape[0]):
+        force_grad[i_b_, i_d_] = dyn_state.dofs.ctrl_force.grad[dofs_idx[i_d_], envs_idx[i_b_]]
+        dyn_state.dofs.ctrl_force.grad[dofs_idx[i_d_], envs_idx[i_b_]] = 0.0
+
+
 @qd.kernel(fastcache=True)
 def kernel_set_dofs_zero_velocity(
     dofs_idx: qd.types.ndarray(),

genesis/engine/solvers/rigid/abd/diff.py

@@ -19,7 +19,7 @@
 import genesis as gs
 import genesis.utils.geom as gu
 import genesis.utils.array_class as array_class
-from .forward_kinematics import func_update_cartesian_space
+from .forward_kinematics import func_update_cartesian_space, func_forward_velocity
 
 
 @qd.func
@@ -150,6 +150,7 @@ def kernel_prepare_backward_substep(
         func_update_cartesian_space(
             dyn_state, dyn_info, rigid_info, rigid_config, force_update_fixed_geoms=False, is_backward=True
         )
+        func_forward_velocity(dyn_state, dyn_info, rigid_info, rigid_config, is_backward=True)
 
         # FIXME: Parameter pruning for ndarray is buggy for now and requires match variable and arg names.
         # Save results of [update_cartesian_space] to adjoint cache
@@ -171,9 +172,9 @@ def kernel_begin_backward_substep(
         func_copy_next_to_curr_grad(f, dyn_state, rigid_adjoint_cache, rigid_info, rigid_config)
 
         if not rigid_config.enable_mujoco_compatibility:
-            # FIXME: Parameter pruning for ndarray is buggy for now and requires match variable and arg names.
-            # Save results of [update_cartesian_space] to adjoint cache
-            func_copy_cartesian_space(dyn_state, dyn_state_adjoint_cache, rigid_config)
+            # Restore the cartesian space that was overwritten by the post-integrate forward replay in the backward
+            # substep (see _update_cartesian_grad in rigid_solver.py).
+            func_copy_cartesian_space(dyn_state_adjoint_cache, dyn_state, rigid_config)
 
     return is_grad_valid
 
@@ -264,6 +265,27 @@ def kernel_copy_acc(
         dyn_state.dofs.acc[i_d, i_b] = rigid_adjoint_cache.dofs_acc[f, i_d, i_b]
 
 
+@qd.kernel(fastcache=True)
+def kernel_copy_next_to_curr_no_check(
+    dyn_state: array_class.DynState,
+    rigid_info: array_class.RigidInfo,
+    rigid_config: qd.template(),
+):
+    # Unguarded copy of the _next slots to current, used in the backward substep right before the forward replay so
+    # the backward kernels see the post-integrate qpos / vel.
+    n_qs = rigid_info.qpos.shape[0]
+    n_dofs = dyn_state.dofs.vel.shape[0]
+    _B = dyn_state.dofs.vel.shape[1]
+
+    qd.loop_config(serialize=rigid_config.para_level < gs.PARA_LEVEL.ALL)
+    for i_q, i_b in qd.ndrange(n_qs, _B):
+        rigid_info.qpos[i_q, i_b] = rigid_info.qpos_next[i_q, i_b]
+
+    qd.loop_config(serialize=rigid_config.para_level < gs.PARA_LEVEL.ALL)
+    for i_d, i_b in qd.ndrange(n_dofs, _B):
+        dyn_state.dofs.vel[i_d, i_b] = dyn_state.dofs.vel_next[i_d, i_b]
+
+
 @qd.func
 def func_integrate_dq_entity(
     i_e,