[FEATURE] Heterogeneous Simulation

examples/rigid/heterogeneous_simulation.py

@@ -0,0 +1,125 @@
+import argparse
+
+import numpy as np
+import genesis as gs
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-v", "--vis", action="store_true", default=False)
+    parser.add_argument("-n", "--n_envs", type=int, default=4)
+    args = parser.parse_args()
+
+    ########################## init ##########################
+    gs.init(backend=gs.gpu, precision="32")
+    ########################## create a scene ##########################
+    scene = gs.Scene(
+        viewer_options=gs.options.ViewerOptions(
+            camera_pos=(3, -1, 1.5),
+            camera_lookat=(0.0, 0.0, 0.5),
+            camera_fov=30,
+            max_FPS=60,
+        ),
+        sim_options=gs.options.SimOptions(
+            dt=0.01,
+            # gravity=(0, 0, -0),
+        ),
+        rigid_options=gs.options.RigidOptions(
+            box_box_detection=True,
+        ),
+        show_viewer=args.vis,
+    )
+
+    ########################## entities ##########################
+    plane = scene.add_entity(
+        gs.morphs.Plane(),
+    )
+    franka = scene.add_entity(
+        gs.morphs.MJCF(file="xml/franka_emika_panda/panda.xml"),
+    )
+
+    morph_heterogeneous = [
+        gs.morphs.Box(
+            size=(0.02, 0.02, 0.02),
+            pos=(0.65, 0.0, 0.02),
+        ),
+        gs.morphs.Sphere(
+            radius=0.015,
+            pos=(0.65, 0.0, 0.02),
+        ),
+        gs.morphs.Sphere(
+            radius=0.025,
+            pos=(0.65, 0.0, 0.02),
+        ),
+        # gs.morphs.Box(
+        #     size=(0.04, 0.04, 0.04),
+        #     pos=(0.65, 0.0, 0.02),
+        # ),
+        # gs.morphs.Box(
+        #     size=(0.04, 0.04, 0.04),
+        #     pos=(0.65, 0.0, 0.02),
+        # ),
+        # gs.morphs.Box(
+        #     size=(0.04, 0.04, 0.04),
+        #     pos=(0.65, 0.0, 0.02),
+        # ),
+    ]
+
+    grasping_object = scene.add_entity(
+        gs.morphs.Box(
+            size=(0.04, 0.04, 0.04),
+            pos=(0.65, 0.0, 0.02),
+        ),
+        # comment out this line turn off heterogeneous simulation
+        morph_heterogeneous=morph_heterogeneous,
+        vis_mode="collision",
+    )
+    ########################## build ##########################
+    scene.build(n_envs=args.n_envs, env_spacing=(1, 1))
+
+    motors_dof = np.arange(7)
+    fingers_dof = np.arange(7, 9)
+    l_qpos = [-1.0124, 1.5559, 1.3662, -1.6878, -1.5799, 1.7757, 1.4602, 0.04, 0.04]
+    if args.n_envs == 0:
+        franka.set_qpos(np.array(l_qpos))
+    else:
+        franka.set_qpos(np.array([l_qpos] * args.n_envs))
+    scene.step()
+
+    end_effector = franka.get_link("hand")
+    qpos = franka.inverse_kinematics(
+        link=end_effector,
+        pos=np.array([[0.65, 0.0, 0.135]] * args.n_envs),
+        quat=np.array([[0, 1, 0, 0]] * args.n_envs),
+    )
+    print(qpos.shape, motors_dof.shape)
+    franka.control_dofs_position(qpos[..., :-2], motors_dof)
+
+    # hold
+    for i in range(100):
+        print("hold", i)
+        scene.step()
+
+    # grasp
+    finder_pos = -0.0
+    for i in range(100):
+        print("grasp", i)
+        franka.control_dofs_position(qpos[..., :-2], motors_dof)
+        franka.control_dofs_position(np.array([[finder_pos, finder_pos]] * args.n_envs), fingers_dof)
+        scene.step()
+
+    # lift
+    qpos = franka.inverse_kinematics(
+        link=end_effector,
+        pos=np.array([[0.65, 0.0, 0.3]] * args.n_envs),
+        quat=np.array([[0, 1, 0, 0]] * args.n_envs),
+    )
+    for i in range(200):
+        print("lift", i)
+        franka.control_dofs_position(qpos[..., :-2], motors_dof)
+        franka.control_dofs_position(np.array([[finder_pos, finder_pos]] * args.n_envs), fingers_dof)
+        scene.step()
+
+
+if __name__ == "__main__":
+    main()

genesis/engine/entities/rigid_entity/rigid_entity.py

@@ -63,6 +63,7 @@ def __init__(
         vface_start=0,
         equality_start=0,
         visualize_contact: bool = False,
+        morph_heterogeneous: list[Morph] | None = None,
     ):
         super().__init__(idx, scene, morph, solver, material, surface)
 
@@ -89,25 +90,101 @@ def __init__(
         self._is_free: bool = morph.is_free
 
         self._is_built: bool = False
+        self._morph_heterogeneous = morph_heterogeneous
+        self._enable_heterogeneous = not (morph_heterogeneous is None or len(morph_heterogeneous) == 0)
 
         self._load_model()
 
+    def load_morph(self, morph: Morph):
+        if isinstance(morph, gs.morphs.Mesh):
+            self._load_mesh(morph, self._surface)
+        elif isinstance(morph, (gs.morphs.MJCF, gs.morphs.URDF, gs.morphs.Drone)):
+            self._load_scene(morph, self._surface)
+        elif isinstance(morph, gs.morphs.Primitive):
+            self._load_primitive(morph, self._surface)
+        elif isinstance(morph, gs.morphs.Terrain):
+            self._load_terrain(morph, self._surface)
+        else:
+            gs.raise_exception(f"Unsupported morph: {morph}.")
+
+    def _load_heterogeneous(self):
+        # TODO
+
+        # handle self._requires_jac_and_IK
+        # handle self._is_local_collision_mask
+        # handle contact pair validity
+        # handle max ndof, nlink, etc
+        # only support  primitive and mesh
+        # com can be off
+
+        self.list_het_link_start = gs.List()
+        self.list_het_link_end = gs.List()
+        self.list_het_n_links = gs.List()
+        self.list_het_geom_group_start = gs.List()
+        self.list_het_geom_group_end = gs.List()
+
+        self.list_het_link_start.append(self._link_start)
+        self.list_het_n_links.append(len(self._links))
+        self.list_het_link_end.append(self._link_start + len(self._links))
+        self.list_het_geom_group_start.append(self._geom_start)
+        self.list_het_geom_group_end.append(self._geom_start + len(self.geoms))
+
+        if self._enable_heterogeneous:
+            for morph in self._morph_heterogeneous:
+                if isinstance(morph, gs.morphs.Mesh):
+                    g_infos = self._load_mesh(morph, self._surface, load_geom_only_for_heterogeneous=True)
+                elif isinstance(morph, gs.morphs.Primitive):
+                    g_infos = self._load_primitive(morph, self._surface, load_geom_only_for_heterogeneous=True)
+                else:
+                    gs.raise_exception(
+                        f"morph_heterogeneous only supports primitive and mesh, Unsupported morph: {morph}."
+                    )
+
+                if len(self._links) != 1:
+                    gs.raise_exception("morph_heterogeneous only supports single link entity.")
+
+                link = self._links[0]
+                cg_infos, vg_infos = self._convert_g_infos_to_cg_infos_and_vg_infos(morph, g_infos, False)
+
+                # Add visual geometries
+                for g_info in vg_infos:
+                    link._add_vgeom(
+                        vmesh=g_info["vmesh"],
+                        init_pos=g_info.get("pos", gu.zero_pos()),
+                        init_quat=g_info.get("quat", gu.identity_quat()),
+                    )
+                # Add collision geometries
+                for g_info in cg_infos:
+                    friction = self.material.friction
+                    if friction is None:
+                        friction = g_info.get("friction", gu.default_friction())
+                    link._add_geom(
+                        mesh=g_info["mesh"],
+                        init_pos=g_info.get("pos", gu.zero_pos()),
+                        init_quat=g_info.get("quat", gu.identity_quat()),
+                        type=g_info["type"],
+                        friction=friction,
+                        sol_params=g_info["sol_params"],
+                        data=g_info.get("data"),
+                        needs_coup=self.material.needs_coup,
+                        contype=g_info["contype"],
+                        conaffinity=g_info["conaffinity"],
+                    )
+
+                self.list_het_link_start.append(self.list_het_link_end[-1])
+                self.list_het_link_end.append(self._link_start + len(self._links))
+                self.list_het_n_links.append(self.list_het_link_end[-1] - self.list_het_link_start[-1])
+                self.list_het_geom_group_start.append(self.list_het_geom_group_end[-1])
+                self.list_het_geom_group_end.append(self.list_het_geom_group_end[-1] + len(cg_infos))
+
     def _load_model(self):
         self._links = gs.List()
         self._joints = gs.List()
         self._equalities = gs.List()
 
-        if isinstance(self._morph, gs.morphs.Mesh):
-            self._load_mesh(self._morph, self._surface)
-        elif isinstance(self._morph, (gs.morphs.MJCF, gs.morphs.URDF, gs.morphs.Drone)):
-            self._load_scene(self._morph, self._surface)
-        elif isinstance(self._morph, gs.morphs.Primitive):
-            self._load_primitive(self._morph, self._surface)
-        elif isinstance(self._morph, gs.morphs.Terrain):
-            self._load_terrain(self._morph, self._surface)
-        else:
-            gs.raise_exception(f"Unsupported morph: {self._morph}.")
+        self.load_morph(self._morph)
 
+        self._load_heterogeneous()
         self._requires_jac_and_IK = self._morph.requires_jac_and_IK
         self._is_local_collision_mask = isinstance(self._morph, gs.morphs.MJCF)
 
@@ -120,7 +197,7 @@ def _update_child_idxs(self):
                 if link.idx not in parent_link.child_idxs:
                     parent_link.child_idxs.append(link.idx)
 
-    def _load_primitive(self, morph, surface):
+    def _load_primitive(self, morph, surface, load_geom_only_for_heterogeneous=False):
         if morph.fixed:
             joint_type = gs.JOINT_TYPE.FIXED
             n_qs = 0
@@ -185,6 +262,9 @@ def _load_primitive(self, morph, surface):
                 )
             )
 
+        if load_geom_only_for_heterogeneous:
+            return g_infos
+
         link, (joint,) = self._add_by_info(
             l_info=dict(
                 is_robot=False,
@@ -209,7 +289,7 @@ def _load_primitive(self, morph, surface):
             surface=surface,
         )
 
-    def _load_mesh(self, morph, surface):
+    def _load_mesh(self, morph, surface, load_geom_only_for_heterogeneous=False):
         if morph.fixed:
             joint_type = gs.JOINT_TYPE.FIXED
             n_qs = 0
@@ -255,6 +335,9 @@ def _load_mesh(self, morph, surface):
                     )
                 )
 
+        if load_geom_only_for_heterogeneous:
+            return g_infos
+
         link_name = morph.file.rsplit("/", 1)[-1].replace(".", "_")
 
         link, (joint,) = self._add_by_info(
@@ -715,6 +798,37 @@ def _add_by_info(self, l_info, j_infos, g_infos, morph, surface):
             )
             joints.append(joint)
 
+        cg_infos, vg_infos = self._convert_g_infos_to_cg_infos_and_vg_infos(morph, g_infos, l_info["is_robot"])
+
+        # Add visual geometries
+        for g_info in vg_infos:
+            link._add_vgeom(
+                vmesh=g_info["vmesh"],
+                init_pos=g_info.get("pos", gu.zero_pos()),
+                init_quat=g_info.get("quat", gu.identity_quat()),
+            )
+        # Add collision geometries
+        for g_info in cg_infos:
+            friction = self.material.friction
+            if friction is None:
+                friction = g_info.get("friction", gu.default_friction())
+            link._add_geom(
+                mesh=g_info["mesh"],
+                init_pos=g_info.get("pos", gu.zero_pos()),
+                init_quat=g_info.get("quat", gu.identity_quat()),
+                type=g_info["type"],
+                friction=friction,
+                sol_params=g_info["sol_params"],
+                data=g_info.get("data"),
+                needs_coup=self.material.needs_coup,
+                contype=g_info["contype"],
+                conaffinity=g_info["conaffinity"],
+            )
+
+        return link, joints
+
+    @staticmethod
+    def _convert_g_infos_to_cg_infos_and_vg_infos(morph, g_infos, is_robot):
         # Separate collision from visual geometry for post-processing
         cg_infos, vg_infos = [], []
         for g_info in g_infos:
@@ -740,7 +854,7 @@ def _add_by_info(self, l_info, j_infos, g_infos, morph, surface):
             # the non-physical part that is added to the original geometries to convexify them are generally inside
             # the mechanical structure and not interacting directly with the outer world. On top of that, not only
             # iy increases the memory footprint and compilation time, but also the simulation speed (marginally).
-            if l_info["is_robot"]:
+            if is_robot:
                 decompose_error_threshold = morph.decompose_robot_error_threshold
             else:
                 decompose_error_threshold = morph.decompose_object_error_threshold
@@ -760,33 +874,7 @@ def _add_by_info(self, l_info, j_infos, g_infos, morph, surface):
             mesh = g_info["mesh"]
             mesh.set_color((*np.random.rand(3), 0.7))
 
-        # Add visual geometries
-        for g_info in vg_infos:
-            link._add_vgeom(
-                vmesh=g_info["vmesh"],
-                init_pos=g_info.get("pos", gu.zero_pos()),
-                init_quat=g_info.get("quat", gu.identity_quat()),
-            )
-
-        # Add collision geometries
-        for g_info in cg_infos:
-            friction = self.material.friction
-            if friction is None:
-                friction = g_info.get("friction", gu.default_friction())
-            link._add_geom(
-                mesh=g_info["mesh"],
-                init_pos=g_info.get("pos", gu.zero_pos()),
-                init_quat=g_info.get("quat", gu.identity_quat()),
-                type=g_info["type"],
-                friction=friction,
-                sol_params=g_info["sol_params"],
-                data=g_info.get("data"),
-                needs_coup=self.material.needs_coup,
-                contype=g_info["contype"],
-                conaffinity=g_info["conaffinity"],
-            )
-
-        return link, joints
+        return cg_infos, vg_infos
 
     def _add_equality(self, name, type, objs_name, data, sol_params):
         objs_id = []