From bf23a04e32288629c2293f17c0f4d756c349a3f9 Mon Sep 17 00:00:00 2001
From: ziyanx02 <ziyanx02@gmail.com>
Date: Sat, 21 Dec 2024 15:15:30 -0500
Subject: [PATCH] fix set_friction_ratio, add domain randomization example

---
 examples/rigid/domain_randomization.py        | 101 ++++++++++++++++++
 .../entities/rigid_entity/rigid_entity.py     |  40 ++++++-
 2 files changed, 138 insertions(+), 3 deletions(-)
 create mode 100644 examples/rigid/domain_randomization.py

diff --git a/examples/rigid/domain_randomization.py b/examples/rigid/domain_randomization.py
new file mode 100644
index 0000000..e8e191b
--- /dev/null
+++ b/examples/rigid/domain_randomization.py
@@ -0,0 +1,101 @@
+import argparse
+
+import numpy as np
+import torch
+
+import genesis as gs
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-v", "--vis", action="store_true", default=False)
+    args = parser.parse_args()
+
+    ########################## init ##########################
+    gs.init(seed=0, precision="32", logging_level="debug")
+
+    ########################## create a scene ##########################
+    scene = gs.Scene(
+        viewer_options=gs.options.ViewerOptions(
+            camera_pos=(0.0, -2, 1.5),
+            camera_lookat=(0.0, 0.0, 0.5),
+            camera_fov=40,
+            max_FPS=200,
+        ),
+        show_viewer=args.vis,
+        rigid_options=gs.options.RigidOptions(
+            dt=0.01,
+            constraint_solver=gs.constraint_solver.Newton,
+        ),
+    )
+
+    ########################## entities ##########################
+    scene.add_entity(
+        gs.morphs.Plane(),
+    )
+    robot = scene.add_entity(
+        gs.morphs.URDF(
+            file="urdf/go2/urdf/go2.urdf",
+            pos=(0, 0, 0.4),
+        ),
+    )
+    ########################## build ##########################
+    n_envs = 8
+    scene.build(n_envs=n_envs)
+
+    ########################## domain randomization ##########################
+    robot.set_friction_ratio(
+        friction_ratio=0.5 + torch.rand(scene.n_envs, robot.n_links),
+        link_indices=np.arange(0, robot.n_links),
+    )
+    robot.set_mass_shift(
+        mass_shift=-0.5 + torch.rand(scene.n_envs, robot.n_links),
+        link_indices=np.arange(0, robot.n_links),
+    )
+    robot.set_COM_shift(
+        com_shift=-0.05 + 0.1 * torch.rand(scene.n_envs, robot.n_links, 3),
+        link_indices=np.arange(0, robot.n_links),
+    )
+
+    joint_names = [
+        "FR_hip_joint",
+        "FR_thigh_joint",
+        "FR_calf_joint",
+        "FL_hip_joint",
+        "FL_thigh_joint",
+        "FL_calf_joint",
+        "RR_hip_joint",
+        "RR_thigh_joint",
+        "RR_calf_joint",
+        "RL_hip_joint",
+        "RL_thigh_joint",
+        "RL_calf_joint",
+    ]
+    motor_dofs = [robot.get_joint(name).dof_idx_local for name in joint_names]
+
+    robot.set_dofs_kp(np.full(12, 20), motor_dofs)
+    robot.set_dofs_kv(np.full(12, 1), motor_dofs)
+    default_dof_pos = np.array(
+        [
+            0.0,
+            0.8,
+            -1.5,
+            0.0,
+            0.8,
+            -1.5,
+            0.0,
+            1.0,
+            -1.5,
+            0.0,
+            1.0,
+            -1.5,
+        ]
+    )
+    robot.control_dofs_position(default_dof_pos, motor_dofs)
+
+    for i in range(1000):
+        scene.step()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/genesis/engine/entities/rigid_entity/rigid_entity.py b/genesis/engine/entities/rigid_entity/rigid_entity.py
index 83c8272..0517c80 100644
--- a/genesis/engine/entities/rigid_entity/rigid_entity.py
+++ b/genesis/engine/entities/rigid_entity/rigid_entity.py
@@ -2182,6 +2182,10 @@ def set_friction_ratio(self, friction_ratio, link_indices, envs_idx=None):
         envs_idx : None | array_like, optional
             The indices of the environments. If None, all environments will be considered. Defaults to None.
         """
+        geom_indices = []
+        for i in link_indices:
+            for j in range(self._links[i].n_geoms):
+                geom_indices.append(self._links[i]._geom_start + j)
         self._solver.set_geoms_friction_ratio(
             torch.cat(
                 [
@@ -2190,9 +2194,7 @@ def set_friction_ratio(self, friction_ratio, link_indices, envs_idx=None):
                 ],
                 dim=-1,
             ),
-            torch.tensor(
-                [[self._links[j]._geom_start + i for i in range(self._links[j].n_geoms)] for j in link_indices]
-            ).view(-1),
+            geom_indices,
             envs_idx,
         )
 
@@ -2217,6 +2219,38 @@ def set_friction(self, friction):
         for link in self._links:
             link.set_friction(friction)
 
+    def set_mass_shift(self, mass_shift, link_indices, envs_idx=None):
+        """
+        Set the mass shift of specified links.
+        Parameters
+        ----------
+        mass : torch.Tensor, shape (n_envs, n_links)
+            The mass shift
+        link_indices : array_like
+            The indices of the links to set mass shift.
+        envs_idx : None | array_like, optional
+            The indices of the environments. If None, all environments will be considered. Defaults to None.
+        """
+        for i in range(len(link_indices)):
+            link_indices[i] += self._link_start
+        self._solver.set_links_mass_shift(mass_shift, link_indices, envs_idx)
+
+    def set_COM_shift(self, com_shift, link_indices, envs_idx=None):
+        """
+        Set the center of mass (COM) shift of specified links.
+        Parameters
+        ----------
+        com : torch.Tensor, shape (n_envs, n_links, 3)
+            The COM shift
+        link_indices : array_like
+            The indices of the links to set COM shift.
+        envs_idx : None | array_like, optional
+            The indices of the environments. If None, all environments will be considered. Defaults to None.
+        """
+        for i in range(len(link_indices)):
+            link_indices[i] += self._link_start
+        self._solver.set_links_COM_shift(com_shift, link_indices, envs_idx)
+
     @gs.assert_built
     def get_mass(self):
         """