[BUG FIX] Preserve external forces across sub-steps. (Genesis-Embodied-AI#1290)

Author: duburcqa

genesis/engine/simulator.py

@@ -263,6 +263,9 @@ def step(self, in_backward=False):
                 if self.cur_substep_local == 0 and not in_backward:
                     self.save_ckpt()
 
+        if self.rigid_solver.is_active():
+            self.rigid_solver._kernel_clear_external_force()
+
     def _step_grad(self):
         for _ in range(self._substeps - 1, -1, -1):
 

genesis/engine/solvers/rigid/rigid_solver_decomp.py

@@ -1695,7 +1695,9 @@ def _func_forward_dynamics(self):
         # self._func_actuation()
         self._func_bias_force()
         self._func_compute_qacc()
-        # FIXME: External forces should be cleared at the end of the step, not during substeps.
+
+    @ti.kernel
+    def _kernel_clear_external_force(self):
         self._func_clear_external_force()
 
     def substep(self):

tests/test_rigid_physics.py

@@ -2122,6 +2122,55 @@ def test_urdf_mimic(show_viewer, tol):
     assert_allclose(gs_qpos[-1], gs_qpos[-2], tol=tol)
 
 
+@pytest.mark.required
+@pytest.mark.parametrize("backend", [gs.cpu])
+def test_drone_hover_same_with_and_without_substeps(show_viewer, tol):
+    base_rpm = 15000
+    scene_ref = gs.Scene(
+        show_viewer=show_viewer,
+        sim_options=gs.options.SimOptions(
+            dt=0.002,
+            substeps=1,
+        ),
+    )
+    drone_ref = scene_ref.add_entity(
+        morph=gs.morphs.Drone(
+            file="urdf/drones/cf2x.urdf",
+            pos=(0, 0, 1.0),
+        ),
+    )
+    scene_ref.build()
+
+    for _ in range(2500):
+        drone_ref.set_propellels_rpm([base_rpm, base_rpm, base_rpm, base_rpm])
+        scene_ref.step()
+
+    pos_ref = drone_ref.get_dofs_position()
+
+    scene_test = gs.Scene(
+        show_viewer=show_viewer,
+        sim_options=gs.options.SimOptions(
+            dt=0.01,
+            substeps=5,
+        ),
+    )
+    drone_test = scene_test.add_entity(
+        morph=gs.morphs.Drone(
+            file="urdf/drones/cf2x.urdf",
+            pos=(0, 0, 1.0),
+        ),
+    )
+    scene_test.build()
+
+    for _ in range(500):
+        drone_test.set_propellels_rpm([base_rpm, base_rpm, base_rpm, base_rpm])
+        scene_test.step()
+
+    pos_test = drone_test.get_dofs_position()
+
+    assert_allclose(pos_ref, pos_test, tol=tol)
+
+
 @pytest.mark.required
 @pytest.mark.parametrize("backend", [gs.cpu])
 def test_drone_advanced(show_viewer):