Merge pull request #135 from ziyanx02/main

add apply_external_force/torque

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()

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)