gs and fix the simpublisher spin bug

Author: xinkai-jiang

DepthCapture_2024-12-10-11-16-51.png

@@ -0,0 +1,107 @@
+import genesis as gs
+import numpy as np
+
+from simpub.sim.genesis_publisher import GenesisPublisher
+
+if __name__ == "__main__":
+    gs.init(backend=gs.gpu)
+    viewer_options = gs.options.ViewerOptions(
+        camera_pos=(0, -3.5, 2.5),
+        camera_lookat=(0.0, 0.0, 0.5),
+        camera_fov=40,
+        max_FPS=60,
+    )
+
+    scene = gs.Scene(
+        viewer_options=viewer_options,
+        sim_options=gs.options.SimOptions(
+            dt=0.01,
+        ),
+        show_viewer=False,
+    )
+    plane = scene.add_entity(
+        gs.morphs.Plane(),
+    )
+    franka = scene.add_entity(
+        gs.morphs.MJCF(file="xml/franka_emika_panda/panda.xml"),
+    )
+
+    publisher = GenesisPublisher(scene)
+    scene.build()
+
+    jnt_names = [
+        "joint1",
+        "joint2",
+        "joint3",
+        "joint4",
+        "joint5",
+        "joint6",
+        "joint7",
+        "finger_joint1",
+        "finger_joint2",
+    ]
+    dofs_idx = [franka.get_joint(name).dof_idx_local for name in jnt_names]
+
+    # Optional: set control gains
+    franka.set_dofs_kp(
+        np.array([4500, 4500, 3500, 3500, 2000, 2000, 2000, 100, 100]),
+        dofs_idx,
+    )
+    franka.set_dofs_kv(
+        np.array([450, 450, 350, 350, 200, 200, 200, 10, 10]),
+        dofs_idx,
+    )
+    franka.set_dofs_force_range(
+        np.array([-87, -87, -87, -87, -12, -12, -12, -100, -100]),
+        np.array([87, 87, 87, 87, 12, 12, 12, 100, 100]),
+        dofs_idx,
+    )
+    # Hard reset
+    for i in range(150):
+        if i < 50:
+            franka.set_dofs_position(np.array([1, 1, 0, 0, 0, 0, 0, 0.04, 0.04]), dofs_idx)
+        elif i < 100:
+            franka.set_dofs_position(np.array([-1, 0.8, 1, -2, 1, 0.5, -0.5, 0.04, 0.04]), dofs_idx)
+        else:
+            franka.set_dofs_position(np.array([0, 0, 0, 0, 0, 0, 0, 0, 0]), dofs_idx)
+
+        scene.step()
+
+    # PD control
+    for i in range(1250):
+        if i == 0:
+            franka.control_dofs_position(
+                np.array([1, 1, 0, 0, 0, 0, 0, 0.04, 0.04]),
+                dofs_idx,
+            )
+        elif i == 250:
+            franka.control_dofs_position(
+                np.array([-1, 0.8, 1, -2, 1, 0.5, -0.5, 0.04, 0.04]),
+                dofs_idx,
+            )
+        elif i == 500:
+            franka.control_dofs_position(
+                np.array([0, 0, 0, 0, 0, 0, 0, 0, 0]),
+                dofs_idx,
+            )
+        elif i == 750:
+            # control first dof with velocity, and the rest with position
+            franka.control_dofs_position(
+                np.array([0, 0, 0, 0, 0, 0, 0, 0, 0])[1:],
+                dofs_idx[1:],
+            )
+            franka.control_dofs_velocity(
+                np.array([1.0, 0, 0, 0, 0, 0, 0, 0, 0])[:1],
+                dofs_idx[:1],
+            )
+        elif i == 1000:
+            franka.control_dofs_force(
+                np.array([0, 0, 0, 0, 0, 0, 0, 0, 0]),
+                dofs_idx,
+            )
+        # This is the internal control force computed based on the given control command
+        # If using force control, it's the same as the given control command
+        print("control force:", franka.get_dofs_control_force(dofs_idx))
+
+        scene.step()
+    publisher.spin()

DepthCapture_2024-12-10-11-16-53.png

@@ -103,6 +103,7 @@ async def update_loop(self):
                 )
             except Exception as e:
                 logger.error(f"Error when streaming {self.topic_name}: {e}")
+                traceback.print_exc()
         logger.info(f"Streamer for topic {self.topic_name} is stopped")
 
 

demos/genesis/franka_example.py

@@ -185,10 +185,16 @@ def stop_node(self):
                 self.loop.call_soon_threadsafe(self.loop.stop)
         except RuntimeError as e:
             logger.error(f"One error occurred when stop server: {e}")
-        self.spin(False)
-
-    def spin(self, wait=True):
-        self.executor.shutdown(wait=wait)
+        self.executor.shutdown(wait=False)
+
+    def spin(self):
+        while True:
+            try:
+                time.sleep(0.01)
+            except KeyboardInterrupt:
+                break
+        self.stop_node()
+        logger.info("The node has been stopped")
 
     def submit_task(
         self,
@@ -217,6 +223,7 @@ async def service_loop(self):
                         f"One error occurred when processing the Service "
                         f'"{service_name}": {e}'
                     )
+                    traceback.print_exc()
                     await service_socket.send(MSG.SERVICE_ERROR.value)
             await async_sleep(0.01)
         logger.info("Service loop has been stopped")

simpub/core/net_component.py

@@ -44,8 +44,10 @@ def __init__(
         ip_addr: str = "127.0.0.1",
         no_rendered_objects: Optional[List[str]] = None,
         no_tracked_objects: Optional[List[str]] = None,
+        fps: int = 60
     ) -> None:
         self.sim_scene = sim_scene
+        self.fps = fps
         if no_rendered_objects is None:
             self.no_rendered_objects = []
         else:
@@ -60,7 +62,7 @@ def initialize(self) -> None:
         self.scene_update_streamer = Streamer(
             topic_name="SceneUpdate",
             update_func=self.get_update,
-            fps=60,
+            fps=self.fps,
             start_streaming=True)
         self.asset_service = StrBytesService("Asset", self._on_asset_request)
         self.xr_device_set: Set[HashIdentifier] = set()

simpub/core/net_manager.py

@@ -0,0 +1,117 @@
+from typing import Tuple, List, Union, Dict
+import genesis as gs
+from genesis.engine.entities import RigidEntity
+from genesis.engine.entities.rigid_entity import RigidLink, RigidGeom
+
+from ..simdata import SimScene, SimObject, SimTransform, SimVisual, VisualType, SimMaterial, SimMesh
+
+
+def gs2unity_pos(pos: List[float]) -> List[float]:
+    return [-pos[1], pos[2], pos[0]]
+
+
+def gs2unity_quat(quat: List[float]) -> List[float]:
+    return [quat[2], -quat[3], -quat[1], quat[0]]
+
+
+class GenesisSceneParser:
+    def __init__(self, gs_scene: gs.Scene):
+        self.gs_scene = gs_scene
+        self.sim_scene, self.update_dict = self.parse_gs_scene(gs_scene)
+
+    def parse_gs_scene(self, gs_scene: gs.Scene) -> Tuple[SimScene, Dict[str, Union[RigidEntity, RigidLink]]]:
+        sim_scene = SimScene()
+        update_dict: Dict[str, Union[RigidEntity, RigidLink]] = {}
+        body_hierarchy = {}
+        sim_scene.root = SimObject(name="genesis_scene")
+        for gs_entity in gs_scene.entities:
+            sim_object, idx = self.process_rigid_entity(gs_entity, sim_scene)
+            body_hierarchy[idx] = {
+                "parent_id": -1,
+                "sim_object": sim_object,
+            }
+            update_dict[sim_object.name] = gs_entity
+            for link in gs_entity.links:
+                sim_object, parent_id, idx = self.process_link(
+                    link, sim_scene
+                )
+                # update the body hierarchy dictionary
+                body_hierarchy[idx] = {
+                    "parent_id": parent_id,
+                    "sim_object": sim_object,
+                }
+                update_dict[sim_object.name] = link
+        for body_info in body_hierarchy.values():
+            parent_id = body_info["parent_id"]
+            if parent_id == -1:
+                sim_scene.root.children.append(body_info["sim_object"])
+            if parent_id in body_hierarchy:
+                parent_info = body_hierarchy[parent_id]
+                parent_object: SimObject = parent_info["sim_object"]
+                parent_object.children.append(body_info["sim_object"])
+        return sim_scene, update_dict
+
+    def process_rigid_entity(
+        self, gs_rigid_entity: RigidEntity, sim_scene: SimScene
+    ) -> Tuple[SimObject, int]:
+        sim_object = SimObject(name=str(gs_rigid_entity.uid))
+        pos: List[float] = [0.0, 0.0, 0.0]
+        rot: List[float] = [0.0, 0.0, 0.0, 1.0]
+        if self.gs_scene.is_built:
+            pos = gs_rigid_entity.get_pos()
+            rot = gs_rigid_entity.get_quat()
+            assert type(pos) is List[float]
+            assert type(rot) is List[float]
+        sim_object.trans = SimTransform(
+            gs2unity_pos(pos),
+            gs2unity_quat(rot),
+        )
+        return sim_object, gs_rigid_entity.base_link_idx
+
+    def process_link(
+        self, gs_link: RigidLink, sim_scene: SimScene
+    ) -> Tuple[SimObject, int, int]:
+        sim_object = SimObject(name=gs_link.name)
+        sim_object.trans = SimTransform(
+            gs2unity_pos(gs_link.pos),
+            gs2unity_quat(gs_link.quat),
+        )
+        for gs_geom in gs_link.vgeoms:
+            sim_visual = self.process_vgeom(gs_geom, sim_scene)
+            sim_object.visuals.append(sim_visual)
+        return sim_object, gs_link.parent_idx, gs_link.idx
+
+    def process_vgeom(
+        self, gs_vgeom: RigidGeom, sim_scene: SimScene
+    ) -> SimVisual:
+        sim_visual = SimVisual(
+            str(gs_vgeom.uid),
+            type=VisualType.MESH,
+            trans=SimTransform(
+                gs2unity_pos(gs_vgeom.init_pos),
+                gs2unity_quat(gs_vgeom.init_quat),
+            ),
+        )
+        sim_visual.mesh = self.process_mesh(gs_vgeom, sim_scene)
+        sim_visual.material = SimMaterial(
+            color=[1.0, 1.0, 1.0, 1.0],
+        )
+        return sim_visual
+
+    def process_mesh(
+        self, gs_vgeom: RigidGeom, sim_scene: SimScene
+    ) -> SimMesh:
+        gs_trimesh_obj = gs_vgeom.get_trimesh()
+        return SimMesh.create_mesh(
+            sim_scene,
+            gs_trimesh_obj.vertices,
+            gs_trimesh_obj.faces,
+        )
+
+    # TODO: Implement the material and texture from trimesh
+    # def parse_material(self, gs_material: gs.Material):
+    #     pass
+    
+    # def parse_mesh(self, gs_mesh: gs.Mesh):
+    #     pass
+

simpub/core/simpub_server.py

@@ -0,0 +1,56 @@
+from typing import List, Optional, Dict, Tuple
+import numpy as np
+import taichi as ti
+
+from genesis.engine.scene import Scene as GSScene
+from genesis.engine.entities import RigidEntity
+from genesis.engine.entities.rigid_entity import RigidLink, RigidGeom
+
+from ..parser.gs import GenesisSceneParser
+from ..core.simpub_server import SimPublisher
+
+
+@ti.data_oriented
+class GenesisPublisher(SimPublisher):
+
+    def __init__(
+        self,
+        gs_scene: GSScene,
+        host: str = "127.0.0.1",
+        no_rendered_objects: Optional[List[str]] = None,
+        no_tracked_objects: Optional[List[str]] = None,
+    ) -> None:
+        self.parser = GenesisSceneParser(gs_scene)
+        sim_scene = self.parser.sim_scene
+        self.update_dict = self.parser.update_dict
+        self.tracked_obj_trans: Dict[str, Tuple[np.ndarray, np.ndarray]] = {}
+        super().__init__(
+            sim_scene,
+            host,
+            no_rendered_objects,
+            no_tracked_objects,
+            fps=60,
+        )
+
+    def get_update(self) -> Dict[str, List[float]]:
+        # TODO: Fix the problem with taiqi kernel issue
+        return {}
+        if not self.parser.gs_scene.is_built:
+            return {}
+        state = {}
+        for name, item in self.update_dict.items():
+            if name in self.no_tracked_objects:
+                continue
+            if isinstance(item, RigidEntity):
+                pos = item.get_pos().tolist()
+                print(pos, item.get_quat())
+                quat = item.get_quat().tolist()
+            elif isinstance(item, RigidLink):
+                pos = item.get_pos().tolist()
+                quat = item.get_quat().tolist()
+            else:
+                continue
+            state[name] = [
+                pos[0], pos[1], pos[2], quat[0], quat[1], quat[2], quat[3]
+            ]
+        return state