mjlab 0.0.0__py3-none-any.whl

mjlab/__init__.py

@@ -0,0 +1,17 @@
+import os
+from pathlib import Path
+
+import warp as wp
+
+MJLAB_SRC_PATH: Path = Path(__file__).parent
+
+
+def configure_warp() -> None:
+  """Configure Warp globally for mjlab."""
+  wp.config.enable_backward = False
+
+  quiet = os.environ.get("MJLAB_WARP_VERBOSE", "").lower() not in ("1", "true", "yes")
+  wp.config.quiet = quiet
+
+
+configure_warp()

mjlab/asset_zoo/README.md

@@ -0,0 +1,6 @@
+# Asset Zoo
+
+| Robot | Image |
+|-------|-------|
+| **Unitree G1** | <img src="img/g1.png" width="200" height="150" alt="Unitree G1 Humanoid"> |
+| **Unitree Go1** | <img src="img/go1.png" width="200" height="150" alt="Unitree Go1 Quadruped"> |

mjlab/asset_zoo/render_robots.py

@@ -0,0 +1,49 @@
+from dataclasses import replace
+from pathlib import Path
+
+import mujoco
+from PIL import Image
+
+from mjlab.entity import Entity, EntityCfg
+from mjlab.utils.spec_config import TextureCfg
+
+_HERE = Path(__file__).parent
+_HEIGHT = 300
+_WIDTH = 400
+
+
+def render_robots(name: str, cfg: EntityCfg):
+  skybox_tex = TextureCfg(
+    name="skybox",
+    type="skybox",
+    builtin="gradient",
+    rgb1=(1, 1, 1),
+    rgb2=(1, 1, 1),
+    width=512,
+    height=3072,
+  )
+  robot = Entity(replace(cfg, textures=(skybox_tex,)))
+  model = robot.compile()
+  data = mujoco.MjData(model)
+  mujoco.mj_resetDataKeyframe(model, data, 0)
+  data.qpos[:3] = 0.0
+  mujoco.mj_forward(model, data)
+  cam = mujoco.MjvCamera()
+  mujoco.mjv_defaultCamera(cam)
+  cam.elevation = -10
+  cam.distance = 1.15 * model.stat.extent
+  cam.azimuth = -160
+  cam.lookat[2] -= 0.1
+  with mujoco.Renderer(model, height=_HEIGHT, width=_WIDTH) as renderer:
+    renderer.update_scene(data, cam)
+    img = renderer.render()
+    img = Image.fromarray(img)
+    img.save(_HERE / "img" / f"{name}.png")
+
+
+if __name__ == "__main__":
+  from mjlab.asset_zoo.robots.unitree_g1.g1_constants import G1_ROBOT_CFG
+  from mjlab.asset_zoo.robots.unitree_go1.go1_constants import GO1_ROBOT_CFG
+
+  render_robots("g1", G1_ROBOT_CFG)
+  render_robots("go1", GO1_ROBOT_CFG)

mjlab/asset_zoo/robots/__init__.py

@@ -0,0 +1,7 @@
+from mjlab.asset_zoo.robots.unitree_g1.g1_constants import G1_ROBOT_CFG
+from mjlab.asset_zoo.robots.unitree_go1.go1_constants import GO1_ROBOT_CFG
+
+__all__ = (
+  "G1_ROBOT_CFG",
+  "GO1_ROBOT_CFG",
+)

mjlab/asset_zoo/robots/unitree_g1/__init__.py

@@ -0,0 +1 @@
+"""Unitree G1 humanoid."""

mjlab/asset_zoo/robots/unitree_g1/g1_constants.py

@@ -0,0 +1,296 @@
+"""Unitree G1 constants."""
+
+from pathlib import Path
+
+import mujoco
+
+from mjlab import MJLAB_SRC_PATH
+from mjlab.entity import EntityArticulationInfoCfg, EntityCfg
+from mjlab.utils.actuator import (
+  ElectricActuator,
+  reflected_inertia_from_two_stage_planetary,
+)
+from mjlab.utils.os import update_assets
+from mjlab.utils.spec_config import ActuatorCfg, CollisionCfg
+
+##
+# MJCF and assets.
+##
+
+G1_XML: Path = (
+  MJLAB_SRC_PATH / "asset_zoo" / "robots" / "unitree_g1" / "xmls" / "g1.xml"
+)
+assert G1_XML.exists()
+
+
+def get_assets(meshdir: str) -> dict[str, bytes]:
+  assets: dict[str, bytes] = {}
+  update_assets(assets, G1_XML.parent / "assets", meshdir)
+  return assets
+
+
+def get_spec() -> mujoco.MjSpec:
+  spec = mujoco.MjSpec.from_file(str(G1_XML))
+  spec.assets = get_assets(spec.meshdir)
+  return spec
+
+
+##
+# Actuator config.
+##
+
+# Motor specs (from Unitree).
+ROTOR_INERTIAS_5020 = (
+  0.139e-4,
+  0.017e-4,
+  0.169e-4,
+)
+GEARS_5020 = (
+  1,
+  1 + (46 / 18),
+  1 + (56 / 16),
+)
+ARMATURE_5020 = reflected_inertia_from_two_stage_planetary(
+  ROTOR_INERTIAS_5020, GEARS_5020
+)
+
+ROTOR_INERTIAS_7520_14 = (
+  0.489e-4,
+  0.098e-4,
+  0.533e-4,
+)
+GEARS_7520_14 = (
+  1,
+  4.5,
+  1 + (48 / 22),
+)
+ARMATURE_7520_14 = reflected_inertia_from_two_stage_planetary(
+  ROTOR_INERTIAS_7520_14, GEARS_7520_14
+)
+
+ROTOR_INERTIAS_7520_22 = (
+  0.489e-4,
+  0.109e-4,
+  0.738e-4,
+)
+GEARS_7520_22 = (
+  1,
+  4.5,
+  5,
+)
+ARMATURE_7520_22 = reflected_inertia_from_two_stage_planetary(
+  ROTOR_INERTIAS_7520_22, GEARS_7520_22
+)
+
+ROTOR_INERTIAS_4010 = (
+  0.068e-4,
+  0.0,
+  0.0,
+)
+GEARS_4010 = (
+  1,
+  5,
+  5,
+)
+ARMATURE_4010 = reflected_inertia_from_two_stage_planetary(
+  ROTOR_INERTIAS_4010, GEARS_4010
+)
+
+ACTUATOR_5020 = ElectricActuator(
+  reflected_inertia=ARMATURE_5020,
+  velocity_limit=37.0,
+  effort_limit=25.0,
+)
+ACTUATOR_7520_14 = ElectricActuator(
+  reflected_inertia=ARMATURE_7520_14,
+  velocity_limit=32.0,
+  effort_limit=88.0,
+)
+ACTUATOR_7520_22 = ElectricActuator(
+  reflected_inertia=ARMATURE_7520_22,
+  velocity_limit=20.0,
+  effort_limit=139.0,
+)
+ACTUATOR_4010 = ElectricActuator(
+  reflected_inertia=ARMATURE_4010,
+  velocity_limit=22.0,
+  effort_limit=5.0,
+)
+
+NATURAL_FREQ = 10 * 2.0 * 3.1415926535  # 10Hz
+DAMPING_RATIO = 2.0
+
+STIFFNESS_5020 = ARMATURE_5020 * NATURAL_FREQ**2
+STIFFNESS_7520_14 = ARMATURE_7520_14 * NATURAL_FREQ**2
+STIFFNESS_7520_22 = ARMATURE_7520_22 * NATURAL_FREQ**2
+STIFFNESS_4010 = ARMATURE_4010 * NATURAL_FREQ**2
+
+DAMPING_5020 = 2.0 * DAMPING_RATIO * ARMATURE_5020 * NATURAL_FREQ
+DAMPING_7520_14 = 2.0 * DAMPING_RATIO * ARMATURE_7520_14 * NATURAL_FREQ
+DAMPING_7520_22 = 2.0 * DAMPING_RATIO * ARMATURE_7520_22 * NATURAL_FREQ
+DAMPING_4010 = 2.0 * DAMPING_RATIO * ARMATURE_4010 * NATURAL_FREQ
+
+G1_ACTUATOR_5020 = ActuatorCfg(
+  joint_names_expr=[
+    ".*_elbow_joint",
+    ".*_shoulder_pitch_joint",
+    ".*_shoulder_roll_joint",
+    ".*_shoulder_yaw_joint",
+    ".*_wrist_roll_joint",
+  ],
+  effort_limit=ACTUATOR_5020.effort_limit,
+  armature=ACTUATOR_5020.reflected_inertia,
+  stiffness=STIFFNESS_5020,
+  damping=DAMPING_5020,
+)
+G1_ACTUATOR_7520_14 = ActuatorCfg(
+  joint_names_expr=[".*_hip_pitch_joint", ".*_hip_yaw_joint", "waist_yaw_joint"],
+  effort_limit=ACTUATOR_7520_14.effort_limit,
+  armature=ACTUATOR_7520_14.reflected_inertia,
+  stiffness=STIFFNESS_7520_14,
+  damping=DAMPING_7520_14,
+)
+G1_ACTUATOR_7520_22 = ActuatorCfg(
+  joint_names_expr=[".*_hip_roll_joint", ".*_knee_joint"],
+  effort_limit=ACTUATOR_7520_22.effort_limit,
+  armature=ACTUATOR_7520_22.reflected_inertia,
+  stiffness=STIFFNESS_7520_22,
+  damping=DAMPING_7520_22,
+)
+G1_ACTUATOR_4010 = ActuatorCfg(
+  joint_names_expr=[".*_wrist_pitch_joint", ".*_wrist_yaw_joint"],
+  effort_limit=ACTUATOR_4010.effort_limit,
+  armature=ACTUATOR_4010.reflected_inertia,
+  stiffness=STIFFNESS_4010,
+  damping=DAMPING_4010,
+)
+
+# Waist pitch/roll and ankles are 4-bar linkages with 2 5020 actuators.
+# Due to the parallel linkage, the effective armature at the ankle and waist joints
+# is configuration dependent. Since the exact geometry of the linkage is unknown, we
+# assume a nominal 1:1 gear ratio. Under this assumption, the joint armature in the
+# nominal configuration is approximated as the sum of the 2 actuators' armatures.
+G1_ACTUATOR_WAIST = ActuatorCfg(
+  joint_names_expr=["waist_pitch_joint", "waist_roll_joint"],
+  effort_limit=ACTUATOR_5020.effort_limit * 2,
+  armature=ACTUATOR_5020.reflected_inertia * 2,
+  stiffness=STIFFNESS_5020 * 2,
+  damping=DAMPING_5020 * 2,
+)
+G1_ACTUATOR_ANKLE = ActuatorCfg(
+  joint_names_expr=[".*_ankle_pitch_joint", ".*_ankle_roll_joint"],
+  effort_limit=ACTUATOR_5020.effort_limit * 2,
+  armature=ACTUATOR_5020.reflected_inertia * 2,
+  stiffness=STIFFNESS_5020 * 2,
+  damping=DAMPING_5020 * 2,
+)
+
+##
+# Keyframe config.
+##
+
+HOME_KEYFRAME = EntityCfg.InitialStateCfg(
+  pos=(0, 0, 0.783675),
+  joint_pos={
+    ".*_hip_pitch_joint": -0.1,
+    ".*_knee_joint": 0.3,
+    ".*_ankle_pitch_joint": -0.2,
+    ".*_shoulder_pitch_joint": 0.2,
+    ".*_elbow_joint": 1.28,
+    "left_shoulder_roll_joint": 0.2,
+    "right_shoulder_roll_joint": -0.2,
+  },
+  joint_vel={".*": 0.0},
+)
+
+KNEES_BENT_KEYFRAME = EntityCfg.InitialStateCfg(
+  pos=(0, 0, 0.76),
+  joint_pos={
+    ".*_hip_pitch_joint": -0.312,
+    ".*_knee_joint": 0.669,
+    ".*_ankle_pitch_joint": -0.363,
+    ".*_elbow_joint": 0.6,
+    "left_shoulder_roll_joint": 0.2,
+    "left_shoulder_pitch_joint": 0.2,
+    "right_shoulder_roll_joint": -0.2,
+    "right_shoulder_pitch_joint": 0.2,
+  },
+  joint_vel={".*": 0.0},
+)
+
+##
+# Collision config.
+##
+
+# This enables all collisions, including self collisions.
+# Self-collisions are given condim=1 while foot collisions
+# are given condim=3 and custom friction and solimp.
+FULL_COLLISION = CollisionCfg(
+  geom_names_expr=[".*_collision"],
+  condim={r"^(left|right)_foot[1-7]_collision$": 3},
+  priority={r"^(left|right)_foot[1-7]_collision$": 1},
+  friction={r"^(left|right)_foot[1-7]_collision$": (0.6,)},
+)
+
+FULL_COLLISION_WITHOUT_SELF = CollisionCfg(
+  geom_names_expr=[".*_collision"],
+  contype=0,
+  conaffinity=1,
+  condim={r"^(left|right)_foot[1-7]_collision$": 3},
+  priority={r"^(left|right)_foot[1-7]_collision$": 1},
+  friction={r"^(left|right)_foot[1-7]_collision$": (0.6,)},
+)
+
+FEET_ONLY_COLLISION = CollisionCfg(
+  geom_names_expr=[r"^(left|right)_foot[1-7]_collision$"],
+  contype=0,
+  conaffinity=1,
+  condim=3,
+  priority=1,
+  friction=(0.6,),
+)
+
+##
+# Final config.
+##
+
+G1_ARTICULATION = EntityArticulationInfoCfg(
+  actuators=(
+    G1_ACTUATOR_5020,
+    G1_ACTUATOR_7520_14,
+    G1_ACTUATOR_7520_22,
+    G1_ACTUATOR_4010,
+    G1_ACTUATOR_WAIST,
+    G1_ACTUATOR_ANKLE,
+  ),
+  soft_joint_pos_limit_factor=0.9,
+)
+
+G1_ROBOT_CFG = EntityCfg(
+  init_state=KNEES_BENT_KEYFRAME,
+  collisions=(FULL_COLLISION,),
+  spec_fn=get_spec,
+  articulation=G1_ARTICULATION,
+)
+
+G1_ACTION_SCALE: dict[str, float] = {}
+for a in G1_ARTICULATION.actuators:
+  e = a.effort_limit
+  s = a.stiffness
+  names = a.joint_names_expr
+  if not isinstance(e, dict):
+    e = {n: e for n in names}
+  if not isinstance(s, dict):
+    s = {n: s for n in names}
+  for n in names:
+    if n in e and n in s and s[n]:
+      G1_ACTION_SCALE[n] = 0.25 * e[n] / s[n]
+
+if __name__ == "__main__":
+  import mujoco.viewer as viewer
+
+  from mjlab.entity.entity import Entity
+
+  robot = Entity(G1_ROBOT_CFG)
+
+  viewer.launch(robot.spec.compile())