From 8fb0a216c56078345c4fb9cb173306413cf9ad6a Mon Sep 17 00:00:00 2001
From: ziyanx02 <ziyanx02@gmail.com>
Date: Fri, 20 Dec 2024 11:33:40 -0500
Subject: [PATCH] add apply links force/torque

---
 examples/rigid/apply_external_force_torque.py | 78 +++++++++++++++++++
 .../solvers/rigid/rigid_solver_decomp.py      | 36 +++++++++
 2 files changed, 114 insertions(+)
 create mode 100644 examples/rigid/apply_external_force_torque.py

diff --git a/examples/rigid/apply_external_force_torque.py b/examples/rigid/apply_external_force_torque.py
new file mode 100644
index 00000000..984baff7
--- /dev/null
+++ b/examples/rigid/apply_external_force_torque.py
@@ -0,0 +1,78 @@
+import argparse
+import numpy as np
+import genesis as gs
+
+from genesis.engine.solvers.rigid.rigid_solver_decomp import RigidSolver
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-v", "--vis", action="store_true", default=False)
+    args = parser.parse_args()
+
+    ########################## init ##########################
+    gs.init(backend=gs.gpu)
+
+    ########################## create a scene ##########################
+    viewer_options = gs.options.ViewerOptions(
+        camera_pos=(0, -3.5, 2.5),
+        camera_lookat=(0.0, 0.0, 1.0),
+        camera_fov=40,
+        max_FPS=60,
+    )
+
+    scene = gs.Scene(
+        viewer_options=viewer_options,
+        sim_options=gs.options.SimOptions(
+            dt=0.01,
+        ),
+        show_viewer=args.vis,
+    )
+
+    ########################## entities ##########################
+    plane = scene.add_entity(
+        gs.morphs.Plane(),
+    )
+    cube = scene.add_entity(
+        gs.morphs.Box(
+            pos=(0, 0, 1.0),
+            size=(0.2, 0.2, 0.2),
+        ),
+    )
+    ########################## build ##########################
+    scene.build()
+
+    for solver in scene.sim.solvers:
+        if not isinstance(solver, RigidSolver):
+            continue
+        rigid_solver = solver
+
+    link_idx = [
+        1,
+    ]
+    rotation_direction = 1
+    for i in range(1000):
+        cube_pos = rigid_solver.get_links_pos(link_idx)
+        cube_pos[:, 2] -= 1
+        force = -100 * cube_pos
+        rigid_solver.apply_links_external_force(
+            force=force,
+            links_idx=link_idx,
+        )
+
+        torque = [
+            [0, 0, rotation_direction * 5],
+        ]
+        rigid_solver.apply_links_external_torque(
+            torque=torque,
+            links_idx=link_idx,
+        )
+
+        scene.step()
+
+        if (i + 50) % 100 == 0:
+            rotation_direction *= -1
+
+
+if __name__ == "__main__":
+    main()
diff --git a/genesis/engine/solvers/rigid/rigid_solver_decomp.py b/genesis/engine/solvers/rigid/rigid_solver_decomp.py
index 49f89a59..0e0786b3 100644
--- a/genesis/engine/solvers/rigid/rigid_solver_decomp.py
+++ b/genesis/engine/solvers/rigid/rigid_solver_decomp.py
@@ -2256,6 +2256,42 @@ def _func_wakeup_entity(self, i_e, i_b):
                 n_awake_links = ti.atomic_add(self.n_awake_links[i_b], 1)
                 self.awake_links[n_awake_links, i_b] = i_l
 
+    def apply_links_external_force(self, force, links_idx, envs_idx=None):
+        force, links_idx, envs_idx = self._validate_2D_io_variables(force, links_idx, 3, envs_idx, idx_name="links_idx")
+
+        self._kernel_apply_links_external_force(force, links_idx, envs_idx)
+
+    @ti.kernel
+    def _kernel_apply_links_external_force(
+        self,
+        force: ti.types.ndarray(),
+        links_idx: ti.types.ndarray(),
+        envs_idx: ti.types.ndarray(),
+    ):
+        ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL)
+        for i_l_, i_b_ in ti.ndrange(links_idx.shape[0], envs_idx.shape[0]):
+            for i in ti.static(range(3)):
+                self.links_state[links_idx[i_l_], envs_idx[i_b_]].cfrc_ext_vel[i] -= force[i_b_, i_l_, i]
+
+    def apply_links_external_torque(self, torque, links_idx, envs_idx=None):
+        torque, links_idx, envs_idx = self._validate_2D_io_variables(
+            torque, links_idx, 3, envs_idx, idx_name="links_idx"
+        )
+
+        self._kernel_apply_links_external_torque(torque, links_idx, envs_idx)
+
+    @ti.kernel
+    def _kernel_apply_links_external_torque(
+        self,
+        torque: ti.types.ndarray(),
+        links_idx: ti.types.ndarray(),
+        envs_idx: ti.types.ndarray(),
+    ):
+        ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL)
+        for i_l_, i_b_ in ti.ndrange(links_idx.shape[0], envs_idx.shape[0]):
+            for i in ti.static(range(3)):
+                self.links_state[links_idx[i_l_], envs_idx[i_b_]].cfrc_ext_ang[i] -= torque[i_b_, i_l_, i]
+
     @ti.func
     def _func_apply_external_force(self, pos, force, link_idx, batch_idx):
         torque = (pos - self.links_state[link_idx, batch_idx].COM).cross(force)