merge

test precommit

add apply external force torque

update

.pre-commit-config.yaml

@@ -0,0 +1,6 @@
+repos:
+  - repo: https://github.com/psf/black
+    rev: 23.9.1  # Use the latest version or specify the version you prefer
+    hooks:
+      - id: black
+        args: ["--config=pyproject.toml"] 

README.md

@@ -15,7 +15,7 @@ Genesis is a physics platform designed for general purpose *Robotics/Embodied AI
 1. A **universal physics engine** re-built from the ground up, capable of simulating a wide range of materials and physical phenomena.
 2. A **lightweight**, **ultra-fast**, **pythonic**, and **user-friendly** robotics simulation platform.
 3. A powerful and fast **photo-realistic rendering system**.
-3. A **generative data engine** that transforms user-prompted natural language description into various modalities of data.
+4. A **generative data engine** that transforms user-prompted natural language description into various modalities of data.
 
 Powered by a universal physics engine re-designed and re-built from the ground up, Genesis integrates various physics solvers and their coupling into a unified framework. This core physics engine is further enhanced by a generative agent framework that operates at an upper level, aiming towards fully **automated data generation** for robotics and beyond. 
 

examples/locomotion/backflip/readme.md

@@ -8,7 +8,7 @@ To try Go2 backflips, follow these steps:
     ```
 2. Place the downloaded checkpoints into the `backflip` directory.
 
-For advanced training code, including features like domain randomization and reference rewards for backflip, visit the repository:
+For advanced training code, including features like domain randomization and reference rewards for backflip, visit this repository:
 
 ```
 https://github.com/ziyanx02/Genesis-backflip

examples/locomotion/go2_backflip.py

@@ -64,7 +64,9 @@ def get_cfgs():
             "dof_vel": 0.05,
         },
     }
-    reward_cfg = {"reward_scales": {},}
+    reward_cfg = {
+        "reward_scales": {},
+    }
     command_cfg = {
         "num_commands": 3,
         "lin_vel_x_range": [0, 0],
@@ -76,16 +78,14 @@ def get_cfgs():
 
 
 class BackflipEnv(Go2Env):
-
     def get_observations(self):
-
         phase = torch.pi * self.episode_length_buf[:, None] / self.max_episode_length
         self.obs_buf = torch.cat(
             [
-                self.base_ang_vel * self.obs_scales['ang_vel'],  # 3
+                self.base_ang_vel * self.obs_scales["ang_vel"],  # 3
                 self.projected_gravity,  # 3
-                (self.dof_pos - self.default_dof_pos) * self.obs_scales['dof_pos'],  # 12
-                self.dof_vel * self.obs_scales['dof_vel'],  # 12
+                (self.dof_pos - self.default_dof_pos) * self.obs_scales["dof_pos"],  # 12
+                self.dof_vel * self.obs_scales["dof_vel"],  # 12
                 self.actions,  # 12
                 self.last_actions,  # 12
                 torch.sin(phase),
@@ -105,6 +105,7 @@ def step(self, actions):
         self.get_observations()
         return self.obs_buf, None, self.rew_buf, self.reset_buf, self.extras
 
+
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument("-e", "--exp_name", type=str, default="single")

examples/locomotion/go2_env.py

@@ -154,7 +154,7 @@ def step(self, actions):
         time_out_idx = (self.episode_length_buf > self.max_episode_length).nonzero(as_tuple=False).flatten()
         self.extras["time_outs"] = torch.zeros_like(self.reset_buf, device=self.device, dtype=gs.tc_float)
         self.extras["time_outs"][time_out_idx] = 1.0
-        
+
         self.reset_idx(self.reset_buf.nonzero(as_tuple=False).flatten())
 
         # compute reward

examples/rigid/apply_external_force_torque.py

@@ -0,0 +1,68 @@
+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
+
+    rotation_direction = 1
+    for i in range(1000):
+
+        cube_pos = rigid_solver.get_links_pos([1,])
+        cube_pos[:, 2] -= 1
+        force = -200 * cube_pos
+        rigid_solver.apply_links_external_force(force=force, links_idx=[1,])
+
+        torque = [[0, 0, rotation_direction * 5],]
+        rigid_solver.apply_links_external_torque(torque=torque, links_idx=[1,])
+
+        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,40 @@ 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)

genesis/logging/time_elapser.py

@@ -39,7 +39,7 @@ def __exit__(self, exc_type, exc_value, traceback):
         pass
 
     def run(self):
-        # re-print logger msg from the begining of the previous line
+        # re-print logger msg from the beginning of the previous line
         self.logger.raw("\x1b[1F" + self.start_msg + " ")
         t_start = time.perf_counter()
         t_elapsed = time.perf_counter() - t_start