diff --git a/mani_skill/agents/base_agent.py b/mani_skill/agents/base_agent.py index 6c8f76e9c..23c2fa253 100644 --- a/mani_skill/agents/base_agent.py +++ b/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 diff --git a/mani_skill/agents/robots/__init__.py b/mani_skill/agents/robots/__init__.py index f1344ea26..3c5e78873 100644 --- a/mani_skill/agents/robots/__init__.py +++ b/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 diff --git a/mani_skill/agents/robots/unitree_go/__init__.py b/mani_skill/agents/robots/unitree_go/__init__.py new file mode 100644 index 000000000..0a040e280 --- /dev/null +++ b/mani_skill/agents/robots/unitree_go/__init__.py @@ -0,0 +1 @@ +from .unitree_go2 import UnitreeGo2 diff --git a/mani_skill/agents/robots/unitree_go/unitree_go2.py b/mani_skill/agents/robots/unitree_go/unitree_go2.py new file mode 100644 index 000000000..97b14d0cd --- /dev/null +++ b/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, + ), + ) diff --git a/mani_skill/examples/demo_robot.py b/mani_skill/examples/demo_robot.py index bc2542860..4f0076801 100644 --- a/mani_skill/examples/demo_robot.py +++ b/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) diff --git a/mani_skill/utils/download_asset.py b/mani_skill/utils/download_asset.py index f60aa95ef..b07459eee 100644 --- a/mani_skill/utils/download_asset.py +++ b/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