Unitree go2 initial work (#293)

* work

* modify standing keyframe, add data source

mani_skill/agents/base_agent.py

@@ -33,7 +33,7 @@
 @dataclass
 class Keyframe:
     pose: sapien.Pose
-    qpos: Array
+    qpos: Optional[Array] = None
     qvel: Optional[Array] = None
 
 

mani_skill/agents/robots/__init__.py

@@ -5,6 +5,7 @@
 from .humanoid import Humanoid
 from .panda import Panda, PandaRealSensed435
 from .trifingerpro import TriFingerPro
+from .unitree_go import *
 from .unitree_h1 import UnitreeH1, UnitreeH1Simplified
 from .ur_e import UR10e
 from .xarm import XArm7Ability

mani_skill/agents/robots/unitree_go/__init__.py

@@ -0,0 +1 @@
+from .unitree_go2 import UnitreeGo2

mani_skill/agents/robots/unitree_go/unitree_go2.py

@@ -0,0 +1,59 @@
+import numpy as np
+import sapien
+
+from mani_skill import ASSET_DIR
+from mani_skill.agents.base_agent import BaseAgent, Keyframe
+from mani_skill.agents.controllers import *
+from mani_skill.agents.registration import register_agent
+from mani_skill.sensors.camera import CameraConfig
+
+
+@register_agent()  # uncomment this if you want to register the agent so you can instantiate it by ID when creating environments
+class UnitreeGo2(BaseAgent):
+    uid = "unitree_go2"
+    urdf_path = f"{ASSET_DIR}/robots/unitree_go2/urdf/go2_description.urdf"  # You can use f"{PACKAGE_ASSET_DIR}" to reference a urdf file in the mani_skill /assets package folder
+
+    # you may need to use this modify the friction values of some links in order to make it possible to e.g. grasp objects or avoid sliding on the floor
+    urdf_config = dict()
+
+    fix_root_link = False
+
+    keyframes = dict(
+        standing=Keyframe(
+            pose=sapien.Pose(p=[0, 0, 0.29]),
+            qpos=np.array(
+                [0.0, 0.0, 0.0, 0.0, 0.9, 0.9, 0.9, 0.9, -1.8, -1.8, -1.8, -1.8]
+            ),
+        )
+    )
+
+    @property
+    def _controller_configs(
+        self,
+    ):
+
+        return dict(
+            pd_joint_pos=dict(
+                body=PDJointPosControllerConfig(
+                    [x.name for x in self.robot.active_joints],
+                    lower=None,
+                    upper=None,
+                    stiffness=100,
+                    damping=10,
+                    normalize_action=False,
+                ),
+                balance_passive_force=False,
+            ),
+            pd_joint_delta_pos=dict(
+                body=PDJointPosControllerConfig(
+                    [x.name for x in self.robot.active_joints],
+                    lower=-0.1,
+                    upper=0.1,
+                    stiffness=20,
+                    damping=5,
+                    normalize_action=True,
+                    use_delta=True,
+                ),
+                balance_passive_force=False,
+            ),
+        )

mani_skill/examples/demo_robot.py

@@ -50,8 +50,10 @@ def parse_args(args=None):
             kf = env.agent.keyframes[args.keyframe]
         else:
             for kf in env.agent.keyframes.values():
+                # keep the first keyframe we find
                 break
-        env.agent.robot.set_qpos(kf.qpos)
+        if kf.qpos is not None:
+            env.agent.robot.set_qpos(kf.qpos)
         if kf.qvel is not None:
             env.agent.robot.set_qvel(kf.qvel)
         env.agent.robot.set_pose(kf.pose)

mani_skill/utils/download_asset.py

@@ -167,6 +167,10 @@ def initialize_extra_sources():
         url="https://github.com/haosulab/ManiSkill-UnitreeH1/archive/refs/tags/v0.1.0.zip",
         target_path="robots/unitree_h1",
     )
+    DATA_SOURCES["unitree_go2"] = DataSource(
+        url="https://github.com/haosulab/ManiSkill-UnitreeGo2/archive/refs/tags/v0.1.0.zip",
+        target_path="robots/unitree_go2",
+    )
 
     # ---------------------------------------------------------------------------- #
     # Visual backgrounds