[BUG FIX] fix set mass (#605)

* fix set mass

examples/rigid/domain_randomization.py

@@ -48,6 +48,19 @@ def main():
         friction_ratio=0.5 + torch.rand(scene.n_envs, robot.n_links),
         ls_idx_local=np.arange(0, robot.n_links),
     )
+    from IPython import embed
+
+    embed()
+
+    # set mass of a single link
+    link = robot.get_link("RR_thigh")
+    rigid = scene.sim.rigid_solver
+    ori_mass = rigid.links_info.inertial_mass.to_numpy()
+    print("original mass", link.get_mass(), ori_mass)
+    link.set_mass(1)
+    new_mass = rigid.links_info.inertial_mass.to_numpy()
+    print("diff mass", new_mass - ori_mass)
+
     robot.set_mass_shift(
         mass_shift=-0.5 + torch.rand(scene.n_envs, robot.n_links),
         ls_idx_local=np.arange(0, robot.n_links),

genesis/engine/solvers/rigid/rigid_solver_decomp.py

@@ -852,11 +852,18 @@ def _kernel_adjust_link_inertia(
         link_idx: ti.i32,
         ratio: ti.f32,
     ):
-        for I_l in ti.grouped(self.links_info):
-            self.links_info[I_l].invweight /= ratio
-            self.links_info[I_l].inertial_mass *= ratio
-            for j1, j2 in ti.ndrange(3, 3):
-                self.links_info[I_l].inertial_i[j1, j2] *= ratio
+        if ti.static(self._options.batch_links_info):
+            for i_b in range(self._B):
+                self.links_info[link_idx, i_b].invweight /= ratio
+                self.links_info[link_idx, i_b].inertial_mass *= ratio
+                for j1, j2 in ti.ndrange(3, 3):
+                    self.links_info[link_idx, i_b].inertial_i[j1, j2] *= ratio
+        else:
+            for i_b in range(self._B):
+                self.links_info[link_idx].invweight /= ratio
+                self.links_info[link_idx].inertial_mass *= ratio
+                for j1, j2 in ti.ndrange(3, 3):
+                    self.links_info[link_idx].inertial_i[j1, j2] *= ratio
 
     def _init_vgeom_fields(self):
         struct_vgeom_info = ti.types.struct(