feat(training-server): release V0.7.3 整合强化学习服务器
web-platform-ci / TypeScript, lint, unit, build (push) Has been cancelled
web-platform-ci / Playwright E2E (push) Has been cancelled
web-platform-ci / TypeScript, lint, unit, build (pull_request) Has been cancelled
web-platform-ci / Playwright E2E (pull_request) Has been cancelled
web-platform-ci / TypeScript, lint, unit, build (push) Has been cancelled
web-platform-ci / Playwright E2E (push) Has been cancelled
web-platform-ci / TypeScript, lint, unit, build (pull_request) Has been cancelled
web-platform-ci / Playwright E2E (pull_request) Has been cancelled
This commit is contained in:
@@ -0,0 +1,4 @@
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
SRC_PATH: Path = Path(__file__).parent
|
||||
@@ -0,0 +1,3 @@
|
||||
"""仓库内置训练器支持的机器人资产。"""
|
||||
|
||||
from .unitree_go2.go2_constants import get_go2_robot_cfg as get_go2_robot_cfg
|
||||
@@ -0,0 +1 @@
|
||||
"""Unitree Go2 quadruped."""
|
||||
@@ -0,0 +1,148 @@
|
||||
"""Unitree Go2 constants."""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import mujoco
|
||||
|
||||
from src import SRC_PATH
|
||||
from mjlab.actuator import BuiltinPositionActuatorCfg
|
||||
from mjlab.entity import EntityArticulationInfoCfg, EntityCfg
|
||||
from mjlab.utils.actuator import ElectricActuator, reflected_inertia
|
||||
from mjlab.utils.os import update_assets
|
||||
from mjlab.utils.spec_config import CollisionCfg
|
||||
|
||||
##
|
||||
# MJCF and assets.
|
||||
##
|
||||
|
||||
GO2_XML: Path = (
|
||||
SRC_PATH / "assets" / "robots" / "unitree_go2" / "xmls" / "go2.xml"
|
||||
)
|
||||
assert GO2_XML.exists()
|
||||
|
||||
|
||||
def get_assets(meshdir: str) -> dict[str, bytes]:
|
||||
assets: dict[str, bytes] = {}
|
||||
update_assets(assets, GO2_XML.parent / "assets", meshdir)
|
||||
return assets
|
||||
|
||||
|
||||
def get_spec() -> mujoco.MjSpec:
|
||||
spec = mujoco.MjSpec.from_file(str(GO2_XML))
|
||||
spec.assets = get_assets(spec.meshdir)
|
||||
return spec
|
||||
|
||||
|
||||
##
|
||||
# Actuator config.
|
||||
##
|
||||
|
||||
GO2_ACTUATOR_HIP = BuiltinPositionActuatorCfg(
|
||||
target_names_expr=(
|
||||
".*hip_.*",
|
||||
),
|
||||
stiffness=20.0,
|
||||
damping=1.0,
|
||||
effort_limit=23.5,
|
||||
armature=0.01,
|
||||
)
|
||||
GO2_ACTUATOR_THIGH = BuiltinPositionActuatorCfg(
|
||||
target_names_expr=(
|
||||
".*thigh_.*",
|
||||
),
|
||||
stiffness=20.0,
|
||||
damping=1.0,
|
||||
effort_limit=23.5,
|
||||
armature=0.01,
|
||||
)
|
||||
GO2_ACTUATOR_CALF = BuiltinPositionActuatorCfg(
|
||||
target_names_expr=(
|
||||
".*calf_.*",
|
||||
),
|
||||
stiffness=40.0,
|
||||
damping=2.0,
|
||||
effort_limit=45,
|
||||
armature=0.02,
|
||||
)
|
||||
|
||||
##
|
||||
# Keyframes.
|
||||
##
|
||||
|
||||
|
||||
INIT_STATE = EntityCfg.InitialStateCfg(
|
||||
pos=(0.0, 0.0, 0.32),
|
||||
joint_pos={
|
||||
".*thigh_joint": 0.9,
|
||||
".*calf_joint": -1.8,
|
||||
".*R_hip_joint": 0.1,
|
||||
".*L_hip_joint": -0.1,
|
||||
},
|
||||
joint_vel={".*": 0.0},
|
||||
)
|
||||
|
||||
##
|
||||
# Collision config.
|
||||
##
|
||||
|
||||
_foot_regex = "^[FR][LR]_foot_collision$"
|
||||
|
||||
# This disables all collisions except the feet.
|
||||
# Furthermore, feet self collisions are disabled.
|
||||
FEET_ONLY_COLLISION = CollisionCfg(
|
||||
geom_names_expr=(_foot_regex,),
|
||||
contype=0,
|
||||
conaffinity=1,
|
||||
condim=3,
|
||||
priority=1,
|
||||
friction=(0.6,),
|
||||
solimp=(0.9, 0.95, 0.023),
|
||||
)
|
||||
|
||||
# This enables all collisions, excluding self collisions.
|
||||
# Foot collisions are given custom condim, friction and solimp.
|
||||
FULL_COLLISION = CollisionCfg(
|
||||
geom_names_expr=(".*_collision",),
|
||||
condim={_foot_regex: 3, ".*_collision": 1},
|
||||
priority={_foot_regex: 1},
|
||||
friction={_foot_regex: (0.6,)},
|
||||
solimp={_foot_regex: (0.9, 0.95, 0.023)},
|
||||
contype=1,
|
||||
conaffinity=0,
|
||||
)
|
||||
|
||||
##
|
||||
# Final config.
|
||||
##
|
||||
|
||||
GO2_ARTICULATION = EntityArticulationInfoCfg(
|
||||
actuators=(
|
||||
GO2_ACTUATOR_HIP,
|
||||
GO2_ACTUATOR_THIGH,
|
||||
GO2_ACTUATOR_CALF,
|
||||
),
|
||||
soft_joint_pos_limit_factor=0.9,
|
||||
)
|
||||
|
||||
|
||||
def get_go2_robot_cfg() -> EntityCfg:
|
||||
"""Get a fresh Go2 robot configuration instance.
|
||||
|
||||
Returns a new EntityCfg instance each time to avoid mutation issues when
|
||||
the config is shared across multiple places.
|
||||
"""
|
||||
return EntityCfg(
|
||||
init_state=INIT_STATE,
|
||||
collisions=(FULL_COLLISION,),
|
||||
spec_fn=get_spec,
|
||||
articulation=GO2_ARTICULATION,
|
||||
)
|
||||
|
||||
if __name__ == "__main__":
|
||||
import mujoco.viewer as viewer
|
||||
|
||||
from mjlab.entity.entity import Entity
|
||||
|
||||
robot = Entity(get_go2_robot_cfg())
|
||||
|
||||
viewer.launch(robot.spec.compile())
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,161 @@
|
||||
<mujoco model="go2">
|
||||
<compiler angle="radian" meshdir="assets" autolimits="true" />
|
||||
|
||||
<default>
|
||||
<default class="go2">
|
||||
<default class="visual">
|
||||
<geom type="mesh" contype="0" conaffinity="0" density="0" group="2"/>
|
||||
</default>
|
||||
<default class="collision">
|
||||
<geom priority="1" condim="6" group="3"/>
|
||||
</default>
|
||||
<site rgba="1 0 0 1" group="5"/>
|
||||
</default>
|
||||
</default>
|
||||
|
||||
<asset>
|
||||
<material name="metal" rgba=".9 .95 .95 1" />
|
||||
<material name="black" rgba="0 0 0 1" />
|
||||
<material name="white" rgba="1 1 1 1" />
|
||||
<material name="gray" rgba="0.671705 0.692426 0.774270 1" />
|
||||
|
||||
<mesh file="base_0.obj" />
|
||||
<mesh file="base_1.obj" />
|
||||
<mesh file="base_2.obj" />
|
||||
<mesh file="base_3.obj" />
|
||||
<mesh file="base_4.obj" />
|
||||
<mesh file="hip_0.obj" />
|
||||
<mesh file="hip_1.obj" />
|
||||
<mesh file="thigh_0.obj" />
|
||||
<mesh file="thigh_1.obj" />
|
||||
<mesh file="thigh_mirror_0.obj" />
|
||||
<mesh file="thigh_mirror_1.obj" />
|
||||
<mesh file="calf_0.obj" />
|
||||
<mesh file="calf_1.obj" />
|
||||
<mesh file="calf_mirror_0.obj" />
|
||||
<mesh file="calf_mirror_1.obj" />
|
||||
<mesh file="foot.obj" />
|
||||
</asset>
|
||||
|
||||
<worldbody>
|
||||
<body name="base_link" pos="0 0 0.445" childclass="go2">
|
||||
<inertial pos="0.021112 0 -0.005366" quat="-0.000543471 0.713435 -0.00173769 0.700719" mass="6.921" diaginertia="0.107027 0.0980771 0.0244531" />
|
||||
<freejoint name="floating_base_joint"/>
|
||||
<geom mesh="base_0" material="black" class="visual" />
|
||||
<geom mesh="base_1" material="black" class="visual" />
|
||||
<geom mesh="base_2" material="black" class="visual" />
|
||||
<geom mesh="base_3" material="white" class="visual" />
|
||||
<geom mesh="base_4" material="gray" class="visual" />
|
||||
<geom name="base1_collision" size="0.1881 0.04675 0.057" type="box" class="collision" />
|
||||
<geom name="base2_collision" size="0.05 0.045" pos="0.285 0 0.01" type="cylinder" class="collision" />
|
||||
<geom name="base3_collision" size="0.047" pos="0.293 0 -0.06" type="sphere" class="collision" />
|
||||
<site name="imu" pos="-0.02557 0 0.04232" group="5"/>
|
||||
<body name="FL_hip" pos="0.1934 0.0465 0">
|
||||
<inertial pos="-0.0054 0.00194 -0.000105" quat="0.497014 0.499245 0.505462 0.498237" mass="0.678" diaginertia="0.00088403 0.000596003 0.000479967" />
|
||||
<joint name="FL_hip_joint" axis="1 0 0" range="-1.0472 1.0472"/>
|
||||
<geom mesh="hip_0" material="metal" class="visual" />
|
||||
<geom mesh="hip_1" material="gray" class="visual" />
|
||||
<geom name="FL_hip_collision" size="0.046 0.02" pos="0 0.08 0" quat="1 1 0 0" type="cylinder" class="collision" />
|
||||
<body name="FL_thigh" pos="0 0.0955 0">
|
||||
<inertial pos="-0.00374 -0.0223 -0.0327" quat="0.829533 0.0847635 -0.0200632 0.551623" mass="1.152" diaginertia="0.00594973 0.00584149 0.000878787" />
|
||||
<joint name="FL_thigh_joint" axis="0 1 0" range="-1.5708 3.4907"/>
|
||||
<geom mesh="thigh_0" material="metal" class="visual" />
|
||||
<geom mesh="thigh_1" material="gray" class="visual" />
|
||||
<geom name="FL_thigh_collision" size="0.1065 0.01225 0.017" pos="0 0 -0.1065" quat="0.707107 0 0.707107 0" type="box" class="collision" />
|
||||
<body name="FL_calf" pos="0 0 -0.213">
|
||||
<inertial pos="0.00629595 -0.000622121 -0.141417" quat="0.710672 0.00154099 -0.00450087 0.703508" mass="0.241352" diaginertia="0.0014901 0.00146356 5.31397e-05" />
|
||||
<joint name="FL_calf_joint" axis="0 1 0" range="-2.7227 -0.83776"/>
|
||||
<geom mesh="calf_0" material="gray" class="visual" />
|
||||
<geom mesh="calf_1" material="black" class="visual" />
|
||||
<geom name="FL_calf1_collision" size="0.012 0.06" pos="0.008 0 -0.06" quat="0.994493 0 -0.104807 0" type="cylinder" class="collision" />
|
||||
<geom name="FL_calf2_collision" size="0.011 0.0325" pos="0.02 0 -0.148" quat="0.999688 0 0.0249974 0" type="cylinder" class="collision" />
|
||||
<geom pos="0 0 -0.213" mesh="foot" class="visual" material="black" />
|
||||
<geom name="FL_foot_collision" pos="0 0 -0.213" type="sphere" size="0.022" class="collision"/>
|
||||
<site name="FL" pos="0 0 -0.213" type="sphere" size="0.022"/>
|
||||
</body>
|
||||
</body>
|
||||
</body>
|
||||
<body name="FR_hip" pos="0.1934 -0.0465 0">
|
||||
<inertial pos="-0.0054 -0.00194 -0.000105" quat="0.498237 0.505462 0.499245 0.497014" mass="0.678" diaginertia="0.00088403 0.000596003 0.000479967" />
|
||||
<joint name="FR_hip_joint" axis="1 0 0" range="-1.0472 1.0472"/>
|
||||
<geom mesh="hip_0" material="metal" class="visual" quat="4.63268e-05 1 0 0" />
|
||||
<geom mesh="hip_1" material="gray" class="visual" quat="4.63268e-05 1 0 0" />
|
||||
<geom name="FR_hip_collision" size="0.046 0.02" pos="0 -0.08 0" quat="0.707107 0.707107 0 0" type="cylinder" class="collision" />
|
||||
<body name="FR_thigh" pos="0 -0.0955 0">
|
||||
<inertial pos="-0.00374 0.0223 -0.0327" quat="0.551623 -0.0200632 0.0847635 0.829533" mass="1.152" diaginertia="0.00594973 0.00584149 0.000878787" />
|
||||
<joint name="FR_thigh_joint" axis="0 1 0" range="-1.5708 3.4907"/>
|
||||
<geom mesh="thigh_mirror_0" material="metal" class="visual" />
|
||||
<geom mesh="thigh_mirror_1" material="gray" class="visual" />
|
||||
<geom name="FR_thigh_collision" size="0.1065 0.01225 0.017" pos="0 0 -0.1065" quat="0.707107 0 0.707107 0" type="box" class="collision" />
|
||||
<body name="FR_calf" pos="0 0 -0.213">
|
||||
<inertial pos="0.00629595 0.000622121 -0.141417" quat="0.703508 -0.00450087 0.00154099 0.710672" mass="0.241352" diaginertia="0.0014901 0.00146356 5.31397e-05" />
|
||||
<joint name="FR_calf_joint" axis="0 1 0" range="-2.7227 -0.83776"/>
|
||||
<geom mesh="calf_mirror_0" material="gray" class="visual" />
|
||||
<geom mesh="calf_mirror_1" material="black" class="visual" />
|
||||
<geom name="FR_calf1_collision" size="0.013 0.06" pos="0.01 0 -0.06" quat="0.995004 0 -0.0998334 0" type="cylinder" class="collision" />
|
||||
<geom name="FR_calf2_collision" size="0.011 0.0325" pos="0.02 0 -0.148" quat="0.999688 0 0.0249974 0" type="cylinder" class="collision" />
|
||||
<geom pos="0 0 -0.213" mesh="foot" class="visual" material="black" />
|
||||
<geom name="FR_foot_collision" pos="0 0 -0.213" type="sphere" size="0.022" class="collision"/>
|
||||
<site name="FR" pos="0 0 -0.213" type="sphere" size="0.022"/>
|
||||
</body>
|
||||
</body>
|
||||
</body>
|
||||
<body name="RL_hip" pos="-0.1934 0.0465 0">
|
||||
<inertial pos="0.0054 0.00194 -0.000105" quat="0.505462 0.498237 0.497014 0.499245" mass="0.678" diaginertia="0.00088403 0.000596003 0.000479967" />
|
||||
<joint name="RL_hip_joint" axis="1 0 0" range="-1.0472 1.0472"/>
|
||||
<geom mesh="hip_0" material="metal" class="visual" quat="4.63268e-05 0 1 0" />
|
||||
<geom mesh="hip_1" material="gray" class="visual" quat="4.63268e-05 0 1 0" />
|
||||
<geom name="RL_hip_collision" size="0.046 0.02" pos="0 0.08 0" quat="0.707107 0.707107 0 0" type="cylinder" class="collision" />
|
||||
<body name="RL_thigh" pos="0 0.0955 0">
|
||||
<inertial pos="-0.00374 -0.0223 -0.0327" quat="0.829533 0.0847635 -0.0200632 0.551623" mass="1.152" diaginertia="0.00594973 0.00584149 0.000878787" />
|
||||
<joint name="RL_thigh_joint" axis="0 1 0" range="-0.5236 4.5379"/>
|
||||
<geom mesh="thigh_0" material="metal" class="visual" />
|
||||
<geom mesh="thigh_1" material="gray" class="visual" />
|
||||
<geom name="RL_thigh_collision" size="0.1065 0.01225 0.017" pos="0 0 -0.1065" quat="0.707107 0 0.707107 0" type="box" class="collision" />
|
||||
<body name="RL_calf" pos="0 0 -0.213">
|
||||
<inertial pos="0.00629595 -0.000622121 -0.141417" quat="0.710672 0.00154099 -0.00450087 0.703508" mass="0.241352" diaginertia="0.0014901 0.00146356 5.31397e-05" />
|
||||
<joint name="RL_calf_joint" axis="0 1 0" range="-2.7227 -0.83776"/>
|
||||
<geom mesh="calf_0" material="gray" class="visual" />
|
||||
<geom mesh="calf_1" material="black" class="visual" />
|
||||
<geom name="RL_calf1_collision" size="0.013 0.06" pos="0.01 0 -0.06" quat="0.995004 0 -0.0998334 0" type="cylinder" class="collision" />
|
||||
<geom name="RL_calf2_collision" size="0.011 0.0325" pos="0.02 0 -0.148" quat="0.999688 0 0.0249974 0" type="cylinder" class="collision" />
|
||||
<geom pos="0 0 -0.213" mesh="foot" class="visual" material="black" />
|
||||
<geom name="RL_foot_collision" pos="0 0 -0.213" type="sphere" size="0.022" class="collision"/>
|
||||
<site name="RL" pos="0 0 -0.213" type="sphere" size="0.022"/>
|
||||
</body>
|
||||
</body>
|
||||
</body>
|
||||
<body name="RR_hip" pos="-0.1934 -0.0465 0">
|
||||
<inertial pos="0.0054 -0.00194 -0.000105" quat="0.499245 0.497014 0.498237 0.505462" mass="0.678" diaginertia="0.00088403 0.000596003 0.000479967" />
|
||||
<joint name="RR_hip_joint" axis="1 0 0" range="-1.0472 1.0472"/>
|
||||
<geom mesh="hip_0" material="metal" class="visual" quat="2.14617e-09 4.63268e-05 4.63268e-05 -1" />
|
||||
<geom mesh="hip_1" material="gray" class="visual" quat="2.14617e-09 4.63268e-05 4.63268e-05 -1" />
|
||||
<geom name="RR_hip_collision" size="0.046 0.02" pos="0 -0.08 0" quat="0.707107 0.707107 0 0" type="cylinder" class="collision" />
|
||||
<body name="RR_thigh" pos="0 -0.0955 0">
|
||||
<inertial pos="-0.00374 0.0223 -0.0327" quat="0.551623 -0.0200632 0.0847635 0.829533" mass="1.152" diaginertia="0.00594973 0.00584149 0.000878787" />
|
||||
<joint name="RR_thigh_joint" axis="0 1 0" range="-0.5236 4.5379"/>
|
||||
<geom mesh="thigh_mirror_0" material="metal" class="visual" />
|
||||
<geom mesh="thigh_mirror_1" material="gray" class="visual" />
|
||||
<geom name="RR_thigh_collision" size="0.1065 0.01225 0.017" pos="0 0 -0.1065" quat="0.707107 0 0.707107 0" type="box" class="collision" />
|
||||
<body name="RR_calf" pos="0 0 -0.213">
|
||||
<inertial pos="0.00629595 0.000622121 -0.141417" quat="0.703508 -0.00450087 0.00154099 0.710672" mass="0.241352" diaginertia="0.0014901 0.00146356 5.31397e-05" />
|
||||
<joint name="RR_calf_joint" axis="0 1 0" range="-2.7227 -0.83776"/>
|
||||
<geom mesh="calf_mirror_0" material="gray" class="visual" />
|
||||
<geom mesh="calf_mirror_1" material="black" class="visual" />
|
||||
<geom name="RR_calf1_collision" size="0.013 0.06" pos="0.01 0 -0.06" quat="0.995004 0 -0.0998334 0" type="cylinder" class="collision" />
|
||||
<geom name="RR_calf2_collision" size="0.011 0.0325" pos="0.02 0 -0.148" quat="0.999688 0 0.0249974 0" type="cylinder" class="collision" />
|
||||
<geom pos="0 0 -0.213" mesh="foot" class="visual" material="black" />
|
||||
<geom name="RR_foot_collision" pos="0 0 -0.213" type="sphere" size="0.022" class="collision"/>
|
||||
<site name="RR" pos="0 0 -0.213" type="sphere" size="0.022"/>
|
||||
</body>
|
||||
</body>
|
||||
</body>
|
||||
</body>
|
||||
</worldbody>
|
||||
<sensor>
|
||||
<gyro name="imu_ang_vel" site="imu"/>
|
||||
<velocimeter name="imu_lin_vel" site="imu"/>
|
||||
<accelerometer name="imu_lin_acc" site="imu"/>
|
||||
<subtreeangmom name="root_angmom" body="base_link"/>
|
||||
</sensor>
|
||||
</mujoco>
|
||||
@@ -0,0 +1,306 @@
|
||||
<mujoco model="scene_go2">
|
||||
<compiler angle="radian" meshdir="assets" autolimits="true" />
|
||||
|
||||
<option cone="elliptic" impratio="100" />
|
||||
|
||||
<default>
|
||||
<default class="go2">
|
||||
<geom friction="0.4" margin="0.001" condim="1"/>
|
||||
<joint axis="0 1 0" damping="0.1" armature="0.01" frictionloss="0.2"/>
|
||||
<motor ctrlrange="-23.7 23.7"/>
|
||||
<default class="abduction">
|
||||
<joint axis="1 0 0" range="-1.0472 1.0472"/>
|
||||
</default>
|
||||
<default class="hip">
|
||||
<default class="front_hip">
|
||||
<joint range="-1.5708 3.4907"/>
|
||||
</default>
|
||||
<default class="back_hip">
|
||||
<joint range="-0.5236 4.5379"/>
|
||||
</default>
|
||||
</default>
|
||||
<default class="knee">
|
||||
<joint range="-2.7227 -0.83776"/>
|
||||
<motor ctrlrange="-45.43 45.43"/>
|
||||
</default>
|
||||
<default class="visual">
|
||||
<geom type="mesh" contype="0" conaffinity="0" group="2"/>
|
||||
</default>
|
||||
<default class="collision">
|
||||
<geom group="3"/>
|
||||
<default class="foot">
|
||||
<geom size="0.022" pos="-0.002 0 -0.213" priority="1" condim="6"
|
||||
friction="0.4 0.02 0.01"/>
|
||||
</default>
|
||||
</default>
|
||||
</default>
|
||||
</default>
|
||||
|
||||
<asset>
|
||||
<material name="metal" rgba=".9 .95 .95 1" />
|
||||
<material name="black" rgba="0 0 0 1" />
|
||||
<material name="white" rgba="1 1 1 1" />
|
||||
<material name="gray" rgba="0.671705 0.692426 0.774270 1" />
|
||||
|
||||
<mesh file="base_0.obj" />
|
||||
<mesh file="base_1.obj" />
|
||||
<mesh file="base_2.obj" />
|
||||
<mesh file="base_3.obj" />
|
||||
<mesh file="base_4.obj" />
|
||||
<mesh file="hip_0.obj" />
|
||||
<mesh file="hip_1.obj" />
|
||||
<mesh file="thigh_0.obj" />
|
||||
<mesh file="thigh_1.obj" />
|
||||
<mesh file="thigh_mirror_0.obj" />
|
||||
<mesh file="thigh_mirror_1.obj" />
|
||||
<mesh file="calf_0.obj" />
|
||||
<mesh file="calf_1.obj" />
|
||||
<mesh file="calf_mirror_0.obj" />
|
||||
<mesh file="calf_mirror_1.obj" />
|
||||
<mesh file="foot.obj" />
|
||||
</asset>
|
||||
|
||||
<worldbody>
|
||||
<body name="base_link" pos="0 0 0.445" childclass="go2">
|
||||
<inertial pos="0.021112 0 -0.005366" quat="-0.000543471 0.713435 -0.00173769 0.700719"
|
||||
mass="6.921"
|
||||
diaginertia="0.107027 0.0980771 0.0244531" />
|
||||
<freejoint />
|
||||
<geom mesh="base_0" material="black" class="visual" />
|
||||
<geom mesh="base_1" material="black" class="visual" />
|
||||
<geom mesh="base_2" material="black" class="visual" />
|
||||
<geom mesh="base_3" material="white" class="visual" />
|
||||
<geom mesh="base_4" material="gray" class="visual" />
|
||||
<geom size="0.1881 0.04675 0.057" type="box" class="collision" />
|
||||
<geom size="0.05 0.045" pos="0.285 0 0.01" type="cylinder" class="collision" />
|
||||
<geom size="0.047" pos="0.293 0 -0.06" class="collision" />
|
||||
<site name="imu" pos="-0.02557 0 0.04232" />
|
||||
<body name="FL_hip" pos="0.1934 0.0465 0">
|
||||
<inertial pos="-0.0054 0.00194 -0.000105" quat="0.497014 0.499245 0.505462 0.498237"
|
||||
mass="0.678"
|
||||
diaginertia="0.00088403 0.000596003 0.000479967" />
|
||||
<joint name="FL_hip_joint" class="abduction" />
|
||||
<geom mesh="hip_0" material="metal" class="visual" />
|
||||
<geom mesh="hip_1" material="gray" class="visual" />
|
||||
<geom size="0.046 0.02" pos="0 0.08 0" quat="1 1 0 0" type="cylinder" class="collision" />
|
||||
<body name="FL_thigh" pos="0 0.0955 0">
|
||||
<inertial pos="-0.00374 -0.0223 -0.0327" quat="0.829533 0.0847635 -0.0200632 0.551623"
|
||||
mass="1.152"
|
||||
diaginertia="0.00594973 0.00584149 0.000878787" />
|
||||
<joint name="FL_thigh_joint" class="front_hip" />
|
||||
<geom mesh="thigh_0" material="metal" class="visual" />
|
||||
<geom mesh="thigh_1" material="gray" class="visual" />
|
||||
<geom size="0.1065 0.01225 0.017" pos="0 0 -0.1065" quat="0.707107 0 0.707107 0"
|
||||
type="box" class="collision" />
|
||||
<body name="FL_calf" pos="0 0 -0.213">
|
||||
<inertial pos="0.00629595 -0.000622121 -0.141417"
|
||||
quat="0.710672 0.00154099 -0.00450087 0.703508"
|
||||
mass="0.241352" diaginertia="0.0014901 0.00146356 5.31397e-05" />
|
||||
<joint name="FL_calf_joint" class="knee" />
|
||||
<geom mesh="calf_0" material="gray" class="visual" />
|
||||
<geom mesh="calf_1" material="black" class="visual" />
|
||||
<geom size="0.012 0.06" pos="0.008 0 -0.06" quat="0.994493 0 -0.104807 0"
|
||||
type="cylinder" class="collision" />
|
||||
<geom size="0.011 0.0325" pos="0.02 0 -0.148" quat="0.999688 0 0.0249974 0"
|
||||
type="cylinder" class="collision" />
|
||||
<geom pos="0 0 -0.213" mesh="foot" class="visual" material="black" />
|
||||
<geom name="FL" class="foot" />
|
||||
<body name="FL_foot" pos="0 0 -0.213" />
|
||||
</body>
|
||||
</body>
|
||||
</body>
|
||||
<body name="FR_hip" pos="0.1934 -0.0465 0">
|
||||
<inertial pos="-0.0054 -0.00194 -0.000105" quat="0.498237 0.505462 0.499245 0.497014"
|
||||
mass="0.678"
|
||||
diaginertia="0.00088403 0.000596003 0.000479967" />
|
||||
<joint name="FR_hip_joint" class="abduction" />
|
||||
<geom mesh="hip_0" material="metal" class="visual" quat="4.63268e-05 1 0 0" />
|
||||
<geom mesh="hip_1" material="gray" class="visual" quat="4.63268e-05 1 0 0" />
|
||||
<geom size="0.046 0.02" pos="0 -0.08 0" quat="0.707107 0.707107 0 0" type="cylinder"
|
||||
class="collision" />
|
||||
<body name="FR_thigh" pos="0 -0.0955 0">
|
||||
<inertial pos="-0.00374 0.0223 -0.0327" quat="0.551623 -0.0200632 0.0847635 0.829533"
|
||||
mass="1.152"
|
||||
diaginertia="0.00594973 0.00584149 0.000878787" />
|
||||
<joint name="FR_thigh_joint" class="front_hip" />
|
||||
<geom mesh="thigh_mirror_0" material="metal" class="visual" />
|
||||
<geom mesh="thigh_mirror_1" material="gray" class="visual" />
|
||||
<geom size="0.1065 0.01225 0.017" pos="0 0 -0.1065" quat="0.707107 0 0.707107 0"
|
||||
type="box" class="collision" />
|
||||
<body name="FR_calf" pos="0 0 -0.213">
|
||||
<inertial pos="0.00629595 0.000622121 -0.141417"
|
||||
quat="0.703508 -0.00450087 0.00154099 0.710672"
|
||||
mass="0.241352" diaginertia="0.0014901 0.00146356 5.31397e-05" />
|
||||
<joint name="FR_calf_joint" class="knee" />
|
||||
<geom mesh="calf_mirror_0" material="gray" class="visual" />
|
||||
<geom mesh="calf_mirror_1" material="black" class="visual" />
|
||||
<geom size="0.013 0.06" pos="0.01 0 -0.06" quat="0.995004 0 -0.0998334 0"
|
||||
type="cylinder" class="collision" />
|
||||
<geom size="0.011 0.0325" pos="0.02 0 -0.148" quat="0.999688 0 0.0249974 0"
|
||||
type="cylinder" class="collision" />
|
||||
<geom pos="0 0 -0.213" mesh="foot" class="visual" material="black" />
|
||||
<geom name="FR" class="foot" />
|
||||
<body name="FR_foot" pos="0 0 -0.213" />
|
||||
</body>
|
||||
</body>
|
||||
</body>
|
||||
<body name="RL_hip" pos="-0.1934 0.0465 0">
|
||||
<inertial pos="0.0054 0.00194 -0.000105" quat="0.505462 0.498237 0.497014 0.499245"
|
||||
mass="0.678"
|
||||
diaginertia="0.00088403 0.000596003 0.000479967" />
|
||||
<joint name="RL_hip_joint" class="abduction" />
|
||||
<geom mesh="hip_0" material="metal" class="visual" quat="4.63268e-05 0 1 0" />
|
||||
<geom mesh="hip_1" material="gray" class="visual" quat="4.63268e-05 0 1 0" />
|
||||
<geom size="0.046 0.02" pos="0 0.08 0" quat="0.707107 0.707107 0 0" type="cylinder"
|
||||
class="collision" />
|
||||
<body name="RL_thigh" pos="0 0.0955 0">
|
||||
<inertial pos="-0.00374 -0.0223 -0.0327" quat="0.829533 0.0847635 -0.0200632 0.551623"
|
||||
mass="1.152"
|
||||
diaginertia="0.00594973 0.00584149 0.000878787" />
|
||||
<joint name="RL_thigh_joint" class="back_hip" />
|
||||
<geom mesh="thigh_0" material="metal" class="visual" />
|
||||
<geom mesh="thigh_1" material="gray" class="visual" />
|
||||
<geom size="0.1065 0.01225 0.017" pos="0 0 -0.1065" quat="0.707107 0 0.707107 0"
|
||||
type="box" class="collision" />
|
||||
<body name="RL_calf" pos="0 0 -0.213">
|
||||
<inertial pos="0.00629595 -0.000622121 -0.141417"
|
||||
quat="0.710672 0.00154099 -0.00450087 0.703508"
|
||||
mass="0.241352" diaginertia="0.0014901 0.00146356 5.31397e-05" />
|
||||
<joint name="RL_calf_joint" class="knee" />
|
||||
<geom mesh="calf_0" material="gray" class="visual" />
|
||||
<geom mesh="calf_1" material="black" class="visual" />
|
||||
<geom size="0.013 0.06" pos="0.01 0 -0.06" quat="0.995004 0 -0.0998334 0"
|
||||
type="cylinder" class="collision" />
|
||||
<geom size="0.011 0.0325" pos="0.02 0 -0.148" quat="0.999688 0 0.0249974 0"
|
||||
type="cylinder" class="collision" />
|
||||
<geom pos="0 0 -0.213" mesh="foot" class="visual" material="black" />
|
||||
<geom name="RL" class="foot" />
|
||||
<body name="RL_foot" pos="0 0 -0.213" />
|
||||
</body>
|
||||
</body>
|
||||
</body>
|
||||
<body name="RR_hip" pos="-0.1934 -0.0465 0">
|
||||
<inertial pos="0.0054 -0.00194 -0.000105" quat="0.499245 0.497014 0.498237 0.505462"
|
||||
mass="0.678"
|
||||
diaginertia="0.00088403 0.000596003 0.000479967" />
|
||||
<joint name="RR_hip_joint" class="abduction" />
|
||||
<geom mesh="hip_0" material="metal" class="visual"
|
||||
quat="2.14617e-09 4.63268e-05 4.63268e-05 -1" />
|
||||
<geom mesh="hip_1" material="gray" class="visual"
|
||||
quat="2.14617e-09 4.63268e-05 4.63268e-05 -1" />
|
||||
<geom size="0.046 0.02" pos="0 -0.08 0" quat="0.707107 0.707107 0 0" type="cylinder"
|
||||
class="collision" />
|
||||
<body name="RR_thigh" pos="0 -0.0955 0">
|
||||
<inertial pos="-0.00374 0.0223 -0.0327" quat="0.551623 -0.0200632 0.0847635 0.829533"
|
||||
mass="1.152"
|
||||
diaginertia="0.00594973 0.00584149 0.000878787" />
|
||||
<joint name="RR_thigh_joint" class="back_hip" />
|
||||
<geom mesh="thigh_mirror_0" material="metal" class="visual" />
|
||||
<geom mesh="thigh_mirror_1" material="gray" class="visual" />
|
||||
<geom size="0.1065 0.01225 0.017" pos="0 0 -0.1065" quat="0.707107 0 0.707107 0"
|
||||
type="box" class="collision" />
|
||||
<body name="RR_calf" pos="0 0 -0.213">
|
||||
<inertial pos="0.00629595 0.000622121 -0.141417"
|
||||
quat="0.703508 -0.00450087 0.00154099 0.710672"
|
||||
mass="0.241352" diaginertia="0.0014901 0.00146356 5.31397e-05" />
|
||||
<joint name="RR_calf_joint" class="knee" />
|
||||
<geom mesh="calf_mirror_0" material="gray" class="visual" />
|
||||
<geom mesh="calf_mirror_1" material="black" class="visual" />
|
||||
<geom size="0.013 0.06" pos="0.01 0 -0.06" quat="0.995004 0 -0.0998334 0"
|
||||
type="cylinder" class="collision" />
|
||||
<geom size="0.011 0.0325" pos="0.02 0 -0.148" quat="0.999688 0 0.0249974 0"
|
||||
type="cylinder" class="collision" />
|
||||
<geom pos="0 0 -0.213" mesh="foot" class="visual" material="black" />
|
||||
<geom name="RR" class="foot" />
|
||||
<body name="RR_foot" pos="0 0 -0.213" />
|
||||
</body>
|
||||
</body>
|
||||
</body>
|
||||
</body>
|
||||
</worldbody>
|
||||
|
||||
<actuator>
|
||||
<motor class="abduction" name="FR_hip" joint="FR_hip_joint" />
|
||||
<motor class="hip" name="FR_thigh" joint="FR_thigh_joint" />
|
||||
<motor class="knee" name="FR_calf" joint="FR_calf_joint" />
|
||||
<motor class="abduction" name="FL_hip" joint="FL_hip_joint" />
|
||||
<motor class="hip" name="FL_thigh" joint="FL_thigh_joint" />
|
||||
<motor class="knee" name="FL_calf" joint="FL_calf_joint" />
|
||||
<motor class="abduction" name="RR_hip" joint="RR_hip_joint" />
|
||||
<motor class="hip" name="RR_thigh" joint="RR_thigh_joint" />
|
||||
<motor class="knee" name="RR_calf" joint="RR_calf_joint" />
|
||||
<motor class="abduction" name="RL_hip" joint="RL_hip_joint" />
|
||||
<motor class="hip" name="RL_thigh" joint="RL_thigh_joint" />
|
||||
<motor class="knee" name="RL_calf" joint="RL_calf_joint" />
|
||||
</actuator>
|
||||
|
||||
<sensor>
|
||||
<jointpos name="FR_hip_pos" joint="FR_hip_joint" />
|
||||
<jointpos name="FR_thigh_pos" joint="FR_thigh_joint" />
|
||||
<jointpos name="FR_calf_pos" joint="FR_calf_joint" />
|
||||
<jointpos name="FL_hip_pos" joint="FL_hip_joint" />
|
||||
<jointpos name="FL_thigh_pos" joint="FL_thigh_joint" />
|
||||
<jointpos name="FL_calf_pos" joint="FL_calf_joint" />
|
||||
<jointpos name="RR_hip_pos" joint="RR_hip_joint" />
|
||||
<jointpos name="RR_thigh_pos" joint="RR_thigh_joint" />
|
||||
<jointpos name="RR_calf_pos" joint="RR_calf_joint" />
|
||||
<jointpos name="RL_hip_pos" joint="RL_hip_joint" />
|
||||
<jointpos name="RL_thigh_pos" joint="RL_thigh_joint" />
|
||||
<jointpos name="RL_calf_pos" joint="RL_calf_joint" />
|
||||
|
||||
<jointvel name="FR_hip_vel" joint="FR_hip_joint" />
|
||||
<jointvel name="FR_thigh_vel" joint="FR_thigh_joint" />
|
||||
<jointvel name="FR_calf_vel" joint="FR_calf_joint" />
|
||||
<jointvel name="FL_hip_vel" joint="FL_hip_joint" />
|
||||
<jointvel name="FL_thigh_vel" joint="FL_thigh_joint" />
|
||||
<jointvel name="FL_calf_vel" joint="FL_calf_joint" />
|
||||
<jointvel name="RR_hip_vel" joint="RR_hip_joint" />
|
||||
<jointvel name="RR_thigh_vel" joint="RR_thigh_joint" />
|
||||
<jointvel name="RR_calf_vel" joint="RR_calf_joint" />
|
||||
<jointvel name="RL_hip_vel" joint="RL_hip_joint" />
|
||||
<jointvel name="RL_thigh_vel" joint="RL_thigh_joint" />
|
||||
<jointvel name="RL_calf_vel" joint="RL_calf_joint" />
|
||||
|
||||
<jointactuatorfrc name="FR_hip_torque" joint="FR_hip_joint" noise="0.01" />
|
||||
<jointactuatorfrc name="FR_thigh_torque" joint="FR_thigh_joint" noise="0.01" />
|
||||
<jointactuatorfrc name="FR_calf_torque" joint="FR_calf_joint" noise="0.01" />
|
||||
<jointactuatorfrc name="FL_hip_torque" joint="FL_hip_joint" noise="0.01" />
|
||||
<jointactuatorfrc name="FL_thigh_torque" joint="FL_thigh_joint" noise="0.01" />
|
||||
<jointactuatorfrc name="FL_calf_torque" joint="FL_calf_joint" noise="0.01" />
|
||||
<jointactuatorfrc name="RR_hip_torque" joint="RR_hip_joint" noise="0.01" />
|
||||
<jointactuatorfrc name="RR_thigh_torque" joint="RR_thigh_joint" noise="0.01" />
|
||||
<jointactuatorfrc name="RR_calf_torque" joint="RR_calf_joint" noise="0.01" />
|
||||
<jointactuatorfrc name="RL_hip_torque" joint="RL_hip_joint" noise="0.01" />
|
||||
<jointactuatorfrc name="RL_thigh_torque" joint="RL_thigh_joint" noise="0.01" />
|
||||
<jointactuatorfrc name="RL_calf_torque" joint="RL_calf_joint" noise="0.01" />
|
||||
|
||||
<framequat name="imu_quat" objtype="site" objname="imu" />
|
||||
<gyro name="imu_gyro" site="imu" />
|
||||
<accelerometer name="imu_acc" site="imu" />
|
||||
|
||||
<framepos name="frame_pos" objtype="site" objname="imu" />
|
||||
<framelinvel name="frame_vel" objtype="site" objname="imu" />
|
||||
</sensor>
|
||||
|
||||
<keyframe>
|
||||
<key name="home" qpos="0 0 0.27 1 0 0 0 0 0.9 -1.8 0 0.9 -1.8 0 0.9 -1.8 0 0.9 -1.8"
|
||||
ctrl="0 0.9 -1.8 0 0.9 -1.8 0 0.9 -1.8 0 0.9 -1.8" />
|
||||
</keyframe>
|
||||
|
||||
<!-- setup scene -->
|
||||
<statistic center="1.0 1.0 1.0" extent="0.8"/>
|
||||
<visual>
|
||||
<headlight diffuse="0.6 0.6 0.6" ambient="0.1 0.1 0.1" specular="0.9 0.9 0.9"/>
|
||||
<rgba haze="0.15 0.25 0.35 1"/>
|
||||
<global azimuth="-140" elevation="-20"/>
|
||||
</visual>
|
||||
<asset>
|
||||
<texture type="skybox" builtin="flat" rgb1="0 0 0" rgb2="0 0 0" width="512" height="3072"/>
|
||||
<texture type="2d" name="groundplane" builtin="checker" mark="edge" rgb1="0.2 0.3 0.4" rgb2="0.1 0.2 0.3" markrgb="0.8 0.8 0.8" width="300" height="300"/>
|
||||
<material name="groundplane" texture="groundplane" texuniform="true" texrepeat="5 5" reflectance="0.2"/>
|
||||
</asset>
|
||||
<worldbody>
|
||||
<light pos="1 0 3.5" dir="0 0 -1" directional="true"/>
|
||||
<geom name="floor" size="0 0 0.05" type="plane" material="groundplane"/>
|
||||
</worldbody>
|
||||
</mujoco>
|
||||
@@ -0,0 +1,5 @@
|
||||
from mjlab.utils.lab_api.tasks.importer import import_packages
|
||||
|
||||
_BLACKLIST_PKGS = ["utils", ".mdp"]
|
||||
|
||||
import_packages(__name__, _BLACKLIST_PKGS)
|
||||
@@ -0,0 +1 @@
|
||||
"""Velocity tracking environments for legged robots."""
|
||||
@@ -0,0 +1,24 @@
|
||||
from mjlab.tasks.registry import register_mjlab_task
|
||||
from src.tasks.velocity.rl import VelocityOnPolicyRunner
|
||||
|
||||
from .env_cfgs import (
|
||||
unitree_go2_flat_env_cfg,
|
||||
unitree_go2_rough_env_cfg,
|
||||
)
|
||||
from .rl_cfg import unitree_go2_ppo_runner_cfg
|
||||
|
||||
register_mjlab_task(
|
||||
task_id="Unitree-Go2-Rough",
|
||||
env_cfg=unitree_go2_rough_env_cfg(),
|
||||
play_env_cfg=unitree_go2_rough_env_cfg(play=True),
|
||||
rl_cfg=unitree_go2_ppo_runner_cfg(),
|
||||
runner_cls=VelocityOnPolicyRunner,
|
||||
)
|
||||
|
||||
register_mjlab_task(
|
||||
task_id="Unitree-Go2-Flat",
|
||||
env_cfg=unitree_go2_flat_env_cfg(),
|
||||
play_env_cfg=unitree_go2_flat_env_cfg(play=True),
|
||||
rl_cfg=unitree_go2_ppo_runner_cfg(),
|
||||
runner_cls=VelocityOnPolicyRunner,
|
||||
)
|
||||
@@ -0,0 +1,170 @@
|
||||
"""Unitree Go2 velocity environment configurations."""
|
||||
|
||||
from typing import Literal
|
||||
|
||||
from src.assets.robots import (
|
||||
get_go2_robot_cfg,
|
||||
)
|
||||
from mjlab.envs import ManagerBasedRlEnvCfg
|
||||
from mjlab.envs import mdp as envs_mdp
|
||||
from mjlab.envs.mdp.actions import JointPositionActionCfg
|
||||
from mjlab.managers import TerminationTermCfg
|
||||
from mjlab.managers.event_manager import EventTermCfg
|
||||
from mjlab.sensor import ContactMatch, ContactSensorCfg, RayCastSensorCfg
|
||||
from mjlab.tasks.velocity import mdp
|
||||
from mjlab.tasks.velocity.mdp import UniformVelocityCommandCfg
|
||||
|
||||
from src.tasks.velocity.velocity_env_cfg import make_velocity_env_cfg
|
||||
|
||||
TerrainType = Literal["rough", "obstacles"]
|
||||
|
||||
|
||||
def unitree_go2_rough_env_cfg(
|
||||
play: bool = False,
|
||||
) -> ManagerBasedRlEnvCfg:
|
||||
"""Create Unitree Go2 rough terrain velocity configuration."""
|
||||
cfg = make_velocity_env_cfg()
|
||||
|
||||
cfg.sim.mujoco.ccd_iterations = 500
|
||||
cfg.sim.contact_sensor_maxmatch = 500
|
||||
|
||||
cfg.scene.entities = {"robot": get_go2_robot_cfg()}
|
||||
|
||||
# Set raycast sensor frame to Go2 base_link.
|
||||
for sensor in cfg.scene.sensors or ():
|
||||
if sensor.name == "terrain_scan":
|
||||
assert isinstance(sensor, RayCastSensorCfg)
|
||||
sensor.frame.name = "base_link"
|
||||
|
||||
foot_names = ("FR", "FL", "RR", "RL")
|
||||
site_names = ("FR", "FL", "RR", "RL")
|
||||
geom_names = tuple(f"{name}_foot_collision" for name in foot_names)
|
||||
|
||||
feet_ground_cfg = ContactSensorCfg(
|
||||
name="feet_ground_contact",
|
||||
primary=ContactMatch(mode="geom", pattern=geom_names, entity="robot"),
|
||||
secondary=ContactMatch(mode="body", pattern="terrain"),
|
||||
fields=("found", "force"),
|
||||
reduce="netforce",
|
||||
num_slots=1,
|
||||
track_air_time=True,
|
||||
)
|
||||
nonfoot_ground_cfg = ContactSensorCfg(
|
||||
name="nonfoot_ground_touch",
|
||||
primary=ContactMatch(
|
||||
mode="geom",
|
||||
entity="robot",
|
||||
# Grab all collision geoms...
|
||||
pattern=r".*_collision\d*$",
|
||||
# Except for the foot geoms.
|
||||
exclude=tuple(geom_names),
|
||||
),
|
||||
secondary=ContactMatch(mode="body", pattern="terrain"),
|
||||
fields=("found", "force"),
|
||||
reduce="none",
|
||||
num_slots=1,
|
||||
history_length=4,
|
||||
)
|
||||
cfg.scene.sensors = (cfg.scene.sensors or ()) + (
|
||||
feet_ground_cfg,
|
||||
nonfoot_ground_cfg,
|
||||
)
|
||||
|
||||
if cfg.scene.terrain is not None and cfg.scene.terrain.terrain_generator is not None:
|
||||
cfg.scene.terrain.terrain_generator.curriculum = True
|
||||
|
||||
joint_pos_action = cfg.actions["joint_pos"]
|
||||
assert isinstance(joint_pos_action, JointPositionActionCfg)
|
||||
|
||||
cfg.viewer.body_name = "base_link"
|
||||
cfg.viewer.distance = 1.5
|
||||
cfg.viewer.elevation = -10.0
|
||||
|
||||
cfg.observations["critic"].terms["foot_height"].params["asset_cfg"].site_names = site_names
|
||||
|
||||
cfg.events["foot_friction"].params["asset_cfg"].geom_names = geom_names
|
||||
cfg.events["base_com"].params["asset_cfg"].body_names = ("base_link",)
|
||||
|
||||
cfg.rewards["pose"].params["std_standing"] = {
|
||||
r".*(FR|FL|RR|RL)_hip_joint.*": 0.05,
|
||||
r".*(FR|FL|RR|RL)_thigh_joint.*": 0.1,
|
||||
r".*(FR|FL|RR|RL)_calf_joint.*": 0.15,
|
||||
}
|
||||
cfg.rewards["pose"].params["std_walking"] = {
|
||||
r".*(FR|FL|RR|RL)_hip_joint.*": 0.15,
|
||||
r".*(FR|FL|RR|RL)_thigh_joint.*": 0.35,
|
||||
r".*(FR|FL|RR|RL)_calf_joint.*": 0.5,
|
||||
}
|
||||
cfg.rewards["pose"].params["std_running"] = {
|
||||
r".*(FR|FL|RR|RL)_hip_joint.*": 0.15,
|
||||
r".*(FR|FL|RR|RL)_thigh_joint.*": 0.35,
|
||||
r".*(FR|FL|RR|RL)_calf_joint.*": 0.5,
|
||||
}
|
||||
|
||||
cfg.rewards["foot_gait"].params["offset"] = [0.0, 0.5, 0.5, 0.0]
|
||||
cfg.rewards["body_orientation_l2"].params["asset_cfg"].body_names = ("base_link",)
|
||||
cfg.rewards["body_ang_vel"].params["asset_cfg"].body_names = ("base_link",)
|
||||
cfg.rewards["foot_clearance"].params["asset_cfg"].site_names = site_names
|
||||
cfg.rewards["foot_slip"].params["asset_cfg"].site_names = site_names
|
||||
|
||||
cfg.terminations["illegal_contact"] = TerminationTermCfg(
|
||||
func=mdp.illegal_contact,
|
||||
params={"sensor_name": nonfoot_ground_cfg.name, "force_threshold": 10.0},
|
||||
)
|
||||
|
||||
# Apply play mode overrides.
|
||||
if play:
|
||||
# Effectively infinite episode length.
|
||||
cfg.episode_length_s = int(1e9)
|
||||
|
||||
cfg.observations["actor"].enable_corruption = False
|
||||
cfg.events.pop("push_robot", None)
|
||||
cfg.curriculum = {}
|
||||
cfg.events["randomize_terrain"] = EventTermCfg(
|
||||
func=envs_mdp.randomize_terrain,
|
||||
mode="reset",
|
||||
params={},
|
||||
)
|
||||
|
||||
if cfg.scene.terrain is not None:
|
||||
if cfg.scene.terrain.terrain_generator is not None:
|
||||
cfg.scene.terrain.terrain_generator.curriculum = False
|
||||
cfg.scene.terrain.terrain_generator.num_cols = 5
|
||||
cfg.scene.terrain.terrain_generator.num_rows = 5
|
||||
cfg.scene.terrain.terrain_generator.border_width = 10.0
|
||||
|
||||
return cfg
|
||||
|
||||
|
||||
def unitree_go2_flat_env_cfg(play: bool = False) -> ManagerBasedRlEnvCfg:
|
||||
"""Create Unitree Go2 flat terrain velocity configuration."""
|
||||
cfg = unitree_go2_rough_env_cfg(play=play)
|
||||
|
||||
cfg.sim.njmax = 300
|
||||
cfg.sim.mujoco.ccd_iterations = 50
|
||||
cfg.sim.contact_sensor_maxmatch = 64
|
||||
cfg.sim.nconmax = None
|
||||
|
||||
# Switch to flat terrain.
|
||||
assert cfg.scene.terrain is not None
|
||||
cfg.scene.terrain.terrain_type = "plane"
|
||||
cfg.scene.terrain.terrain_generator = None
|
||||
|
||||
# Remove raycast sensor and height scan (no terrain to scan).
|
||||
cfg.scene.sensors = tuple(
|
||||
s for s in (cfg.scene.sensors or ()) if s.name != "terrain_scan"
|
||||
)
|
||||
del cfg.observations["actor"].terms["height_scan"]
|
||||
del cfg.observations["critic"].terms["height_scan"]
|
||||
|
||||
# Disable terrain curriculum (not present in play mode since rough clears all).
|
||||
cfg.curriculum.pop("terrain_levels", None)
|
||||
|
||||
if play:
|
||||
twist_cmd = cfg.commands["twist"]
|
||||
assert isinstance(twist_cmd, UniformVelocityCommandCfg)
|
||||
twist_cmd.ranges.lin_vel_x = (-0.5, 1.0)
|
||||
twist_cmd.ranges.lin_vel_y = (-0.5, 0.5)
|
||||
twist_cmd.ranges.ang_vel_z = (-0.5, 0.5)
|
||||
|
||||
return cfg
|
||||
@@ -0,0 +1,46 @@
|
||||
"""RL configuration for Unitree Go2 velocity task."""
|
||||
|
||||
from mjlab.rl import (
|
||||
RslRlModelCfg,
|
||||
RslRlOnPolicyRunnerCfg,
|
||||
RslRlPpoAlgorithmCfg,
|
||||
)
|
||||
|
||||
|
||||
def unitree_go2_ppo_runner_cfg() -> RslRlOnPolicyRunnerCfg:
|
||||
"""Create RL runner configuration for Unitree Go2 velocity task."""
|
||||
return RslRlOnPolicyRunnerCfg(
|
||||
actor=RslRlModelCfg(
|
||||
hidden_dims=(512, 256, 128),
|
||||
activation="elu",
|
||||
obs_normalization=True,
|
||||
distribution_cfg={
|
||||
"class_name": "GaussianDistribution",
|
||||
"init_std": 1.0,
|
||||
"std_type": "scalar",
|
||||
},
|
||||
),
|
||||
critic=RslRlModelCfg(
|
||||
hidden_dims=(512, 256, 128),
|
||||
activation="elu",
|
||||
obs_normalization=True,
|
||||
),
|
||||
algorithm=RslRlPpoAlgorithmCfg(
|
||||
value_loss_coef=1.0,
|
||||
use_clipped_value_loss=True,
|
||||
clip_param=0.2,
|
||||
entropy_coef=0.01,
|
||||
num_learning_epochs=5,
|
||||
num_mini_batches=4,
|
||||
learning_rate=1.0e-3,
|
||||
schedule="adaptive",
|
||||
gamma=0.99,
|
||||
lam=0.95,
|
||||
desired_kl=0.01,
|
||||
max_grad_norm=1.0,
|
||||
),
|
||||
experiment_name="go2_velocity",
|
||||
save_interval=100,
|
||||
num_steps_per_env=24,
|
||||
max_iterations=10001,
|
||||
)
|
||||
@@ -0,0 +1,7 @@
|
||||
from mjlab.envs.mdp import * # noqa: F401, F403
|
||||
|
||||
from .curriculums import * # noqa: F403
|
||||
from .observations import * # noqa: F403
|
||||
from .rewards import * # noqa: F403
|
||||
from .terminations import * # noqa: F403
|
||||
from .velocity_command import * # noqa: F403
|
||||
@@ -0,0 +1,107 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, TypedDict, cast
|
||||
|
||||
import torch
|
||||
|
||||
from mjlab.entity import Entity
|
||||
from mjlab.managers.scene_entity_config import SceneEntityCfg
|
||||
|
||||
from .velocity_command import UniformVelocityCommandCfg
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from mjlab.envs import ManagerBasedRlEnv
|
||||
|
||||
_DEFAULT_SCENE_CFG = SceneEntityCfg("robot")
|
||||
|
||||
|
||||
class VelocityStage(TypedDict):
|
||||
step: int
|
||||
lin_vel_x: tuple[float, float] | None
|
||||
lin_vel_y: tuple[float, float] | None
|
||||
ang_vel_z: tuple[float, float] | None
|
||||
|
||||
|
||||
class RewardWeightStage(TypedDict):
|
||||
step: int
|
||||
weight: float
|
||||
|
||||
|
||||
def terrain_levels_vel(
|
||||
env: ManagerBasedRlEnv,
|
||||
env_ids: torch.Tensor,
|
||||
command_name: str,
|
||||
asset_cfg: SceneEntityCfg = _DEFAULT_SCENE_CFG,
|
||||
) -> torch.Tensor:
|
||||
asset: Entity = env.scene[asset_cfg.name]
|
||||
|
||||
terrain = env.scene.terrain
|
||||
assert terrain is not None
|
||||
terrain_generator = terrain.cfg.terrain_generator
|
||||
assert terrain_generator is not None
|
||||
|
||||
command = env.command_manager.get_command(command_name)
|
||||
assert command is not None
|
||||
|
||||
# Compute the distance the robot walked.
|
||||
distance = torch.norm(
|
||||
asset.data.root_link_pos_w[env_ids, :2] - env.scene.env_origins[env_ids, :2], dim=1
|
||||
)
|
||||
|
||||
# Robots that walked far enough progress to harder terrains.
|
||||
move_up = distance > terrain_generator.size[0] / 2
|
||||
|
||||
# Robots that walked less than half of their required distance go to simpler
|
||||
# terrains.
|
||||
move_down = (
|
||||
distance < torch.norm(command[env_ids, :2], dim=1) * env.max_episode_length_s * 0.5
|
||||
)
|
||||
move_down *= ~move_up
|
||||
|
||||
# Update terrain levels.
|
||||
terrain.update_env_origins(env_ids, move_up, move_down)
|
||||
|
||||
return torch.mean(terrain.terrain_levels.float())
|
||||
|
||||
|
||||
def commands_vel(
|
||||
env: ManagerBasedRlEnv,
|
||||
env_ids: torch.Tensor,
|
||||
command_name: str,
|
||||
velocity_stages: list[VelocityStage],
|
||||
) -> dict[str, torch.Tensor]:
|
||||
del env_ids # Unused.
|
||||
command_term = env.command_manager.get_term(command_name)
|
||||
assert command_term is not None
|
||||
cfg = cast(UniformVelocityCommandCfg, command_term.cfg)
|
||||
for stage in velocity_stages:
|
||||
if env.common_step_counter > stage["step"]:
|
||||
if "lin_vel_x" in stage and stage["lin_vel_x"] is not None:
|
||||
cfg.ranges.lin_vel_x = stage["lin_vel_x"]
|
||||
if "lin_vel_y" in stage and stage["lin_vel_y"] is not None:
|
||||
cfg.ranges.lin_vel_y = stage["lin_vel_y"]
|
||||
if "ang_vel_z" in stage and stage["ang_vel_z"] is not None:
|
||||
cfg.ranges.ang_vel_z = stage["ang_vel_z"]
|
||||
return {
|
||||
# "lin_vel_x_min": torch.tensor(cfg.ranges.lin_vel_x[0]),
|
||||
# "lin_vel_x_max": torch.tensor(cfg.ranges.lin_vel_x[1]),
|
||||
# "lin_vel_y_min": torch.tensor(cfg.ranges.lin_vel_y[0]),
|
||||
# "lin_vel_y_max": torch.tensor(cfg.ranges.lin_vel_y[1]),
|
||||
# "ang_vel_z_min": torch.tensor(cfg.ranges.ang_vel_z[0]),
|
||||
# "ang_vel_z_max": torch.tensor(cfg.ranges.ang_vel_z[1]),
|
||||
}
|
||||
|
||||
|
||||
def reward_weight(
|
||||
env: ManagerBasedRlEnv,
|
||||
env_ids: torch.Tensor,
|
||||
reward_name: str,
|
||||
weight_stages: list[RewardWeightStage],
|
||||
) -> torch.Tensor:
|
||||
"""Update a reward term's weight based on training step stages."""
|
||||
del env_ids # Unused.
|
||||
reward_term_cfg = env.reward_manager.get_term_cfg(reward_name)
|
||||
for stage in weight_stages:
|
||||
if env.common_step_counter > stage["step"]:
|
||||
reward_term_cfg.weight = stage["weight"]
|
||||
return torch.tensor([reward_term_cfg.weight])
|
||||
@@ -0,0 +1,54 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
from mjlab.entity import Entity
|
||||
from mjlab.managers.scene_entity_config import SceneEntityCfg
|
||||
from mjlab.sensor import ContactSensor
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from mjlab.envs import ManagerBasedRlEnv
|
||||
|
||||
_DEFAULT_ASSET_CFG = SceneEntityCfg("robot")
|
||||
|
||||
|
||||
def foot_height(
|
||||
env: ManagerBasedRlEnv, asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG
|
||||
) -> torch.Tensor:
|
||||
asset: Entity = env.scene[asset_cfg.name]
|
||||
return asset.data.site_pos_w[:, asset_cfg.site_ids, 2] # (num_envs, num_sites)
|
||||
|
||||
|
||||
def foot_air_time(env: ManagerBasedRlEnv, sensor_name: str) -> torch.Tensor:
|
||||
sensor: ContactSensor = env.scene[sensor_name]
|
||||
sensor_data = sensor.data
|
||||
current_air_time = sensor_data.current_air_time
|
||||
assert current_air_time is not None
|
||||
return current_air_time
|
||||
|
||||
|
||||
def foot_contact(env: ManagerBasedRlEnv, sensor_name: str) -> torch.Tensor:
|
||||
sensor: ContactSensor = env.scene[sensor_name]
|
||||
sensor_data = sensor.data
|
||||
assert sensor_data.found is not None
|
||||
return (sensor_data.found > 0).float()
|
||||
|
||||
|
||||
def foot_contact_forces(env: ManagerBasedRlEnv, sensor_name: str) -> torch.Tensor:
|
||||
sensor: ContactSensor = env.scene[sensor_name]
|
||||
sensor_data = sensor.data
|
||||
assert sensor_data.force is not None
|
||||
forces_flat = sensor_data.force.flatten(start_dim=1) # [B, N*3]
|
||||
return torch.sign(forces_flat) * torch.log1p(torch.abs(forces_flat))
|
||||
|
||||
|
||||
def phase(env: ManagerBasedRlEnv, period: float, command_name: str) -> torch.Tensor:
|
||||
global_phase = (env.episode_length_buf * env.step_dt) % period / period
|
||||
phase = torch.zeros(env.num_envs, 2, device=env.device)
|
||||
phase[:, 0] = torch.sin(global_phase * torch.pi * 2.0)
|
||||
phase[:, 1] = torch.cos(global_phase * torch.pi * 2.0)
|
||||
stand_mask = torch.linalg.norm(env.command_manager.get_command(command_name), dim=1) < 0.1
|
||||
phase = torch.where(stand_mask.unsqueeze(1), torch.zeros_like(phase), phase)
|
||||
return phase
|
||||
@@ -0,0 +1,427 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
from mjlab.entity import Entity
|
||||
from mjlab.managers.reward_manager import RewardTermCfg
|
||||
from mjlab.managers.scene_entity_config import SceneEntityCfg
|
||||
from mjlab.sensor import BuiltinSensor, ContactSensor
|
||||
from mjlab.utils.lab_api.math import quat_apply_inverse
|
||||
from mjlab.utils.lab_api.string import (
|
||||
resolve_matching_names_values,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from mjlab.envs import ManagerBasedRlEnv
|
||||
|
||||
|
||||
_DEFAULT_ASSET_CFG = SceneEntityCfg("robot")
|
||||
|
||||
|
||||
def track_linear_velocity(
|
||||
env: ManagerBasedRlEnv,
|
||||
std: float,
|
||||
command_name: str,
|
||||
asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG,
|
||||
) -> torch.Tensor:
|
||||
"""Reward for tracking the commanded base linear velocity.
|
||||
|
||||
The commanded z velocity is assumed to be zero.
|
||||
"""
|
||||
asset: Entity = env.scene[asset_cfg.name]
|
||||
command = env.command_manager.get_command(command_name)
|
||||
assert command is not None, f"Command '{command_name}' not found."
|
||||
actual = asset.data.root_link_lin_vel_b
|
||||
xy_error = torch.sum(torch.square(command[:, :2] - actual[:, :2]), dim=1)
|
||||
z_error = torch.square(actual[:, 2])
|
||||
lin_vel_error = xy_error + (2 * z_error)
|
||||
return torch.exp(-lin_vel_error / std**2)
|
||||
|
||||
|
||||
def track_angular_velocity(
|
||||
env: ManagerBasedRlEnv,
|
||||
std: float,
|
||||
command_name: str,
|
||||
asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG,
|
||||
) -> torch.Tensor:
|
||||
"""Reward heading error for heading-controlled envs, angular velocity for others.
|
||||
|
||||
The commanded xy angular velocities are assumed to be zero.
|
||||
"""
|
||||
asset: Entity = env.scene[asset_cfg.name]
|
||||
command = env.command_manager.get_command(command_name)
|
||||
assert command is not None, f"Command '{command_name}' not found."
|
||||
actual = asset.data.root_link_ang_vel_b
|
||||
z_error = torch.square(command[:, 2] - actual[:, 2])
|
||||
xy_error = torch.sum(torch.square(actual[:, :2]), dim=1)
|
||||
ang_vel_error = z_error + (0.05 * xy_error)
|
||||
return torch.exp(-ang_vel_error / std**2)
|
||||
|
||||
|
||||
def body_orientation_l2(
|
||||
env: ManagerBasedRlEnv,
|
||||
asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG,
|
||||
) -> torch.Tensor:
|
||||
"""Reward flat base orientation (robot being upright).
|
||||
|
||||
If asset_cfg has body_ids specified, computes the projected gravity
|
||||
for that specific body. Otherwise, uses the root link projected gravity.
|
||||
"""
|
||||
asset: Entity = env.scene[asset_cfg.name]
|
||||
|
||||
# If body_ids are specified, compute projected gravity for that body.
|
||||
if asset_cfg.body_ids:
|
||||
body_quat_w = asset.data.body_link_quat_w[:, asset_cfg.body_ids, :] # [B, N, 4]
|
||||
body_quat_w = body_quat_w.squeeze(1) # [B, 4]
|
||||
gravity_w = asset.data.gravity_vec_w # [3]
|
||||
projected_gravity_b = quat_apply_inverse(body_quat_w, gravity_w) # [B, 3]
|
||||
xy_squared = torch.sum(torch.square(projected_gravity_b[:, :2]), dim=1)
|
||||
else:
|
||||
# Use root link projected gravity.
|
||||
xy_squared = torch.sum(torch.square(asset.data.projected_gravity_b[:, :2]), dim=1)
|
||||
return xy_squared
|
||||
|
||||
|
||||
def self_collision_cost(
|
||||
env: ManagerBasedRlEnv,
|
||||
sensor_name: str,
|
||||
force_threshold: float = 10.0,
|
||||
) -> torch.Tensor:
|
||||
"""Penalize self-collisions.
|
||||
|
||||
When the sensor provides force history (from ``history_length > 0``),
|
||||
counts substeps where any contact force exceeds *force_threshold*.
|
||||
Falls back to the instantaneous ``found`` count otherwise.
|
||||
"""
|
||||
sensor: ContactSensor = env.scene[sensor_name]
|
||||
data = sensor.data
|
||||
if data.force_history is not None:
|
||||
# force_history: [B, N, H, 3]
|
||||
force_mag = torch.norm(data.force_history, dim=-1) # [B, N, H]
|
||||
hit = (force_mag > force_threshold).any(dim=1) # [B, H]
|
||||
return hit.sum(dim=-1).float() # [B]
|
||||
assert data.found is not None
|
||||
return data.found.squeeze(-1)
|
||||
|
||||
|
||||
def body_angular_velocity_penalty(
|
||||
env: ManagerBasedRlEnv,
|
||||
asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG,
|
||||
) -> torch.Tensor:
|
||||
"""Penalize excessive body angular velocities."""
|
||||
asset: Entity = env.scene[asset_cfg.name]
|
||||
ang_vel = asset.data.body_link_ang_vel_w[:, asset_cfg.body_ids, :]
|
||||
ang_vel = ang_vel.squeeze(1)
|
||||
ang_vel_xy = ang_vel[:, :2] # Don't penalize z-angular velocity.
|
||||
return torch.sum(torch.square(ang_vel_xy), dim=1)
|
||||
|
||||
|
||||
def angular_momentum_penalty(
|
||||
env: ManagerBasedRlEnv,
|
||||
sensor_name: str,
|
||||
) -> torch.Tensor:
|
||||
"""Penalize whole-body angular momentum to encourage natural arm swing."""
|
||||
angmom_sensor: BuiltinSensor = env.scene[sensor_name]
|
||||
angmom = angmom_sensor.data
|
||||
angmom_magnitude_sq = torch.sum(torch.square(angmom), dim=-1)
|
||||
angmom_magnitude = torch.sqrt(angmom_magnitude_sq)
|
||||
env.extras["log"]["Metrics/angular_momentum_mean"] = torch.mean(angmom_magnitude)
|
||||
return angmom_magnitude_sq
|
||||
|
||||
|
||||
def feet_air_time(
|
||||
env: ManagerBasedRlEnv,
|
||||
sensor_name: str,
|
||||
threshold: float = 0.4,
|
||||
command_name: str | None = None,
|
||||
command_threshold: float = 0.1,
|
||||
) -> torch.Tensor:
|
||||
"""Reward feet air time."""
|
||||
sensor: ContactSensor = env.scene[sensor_name]
|
||||
sensor_data = sensor.data
|
||||
air_time = sensor_data.current_air_time
|
||||
contact_time = sensor_data.current_contact_time
|
||||
in_contact = contact_time > 0.0
|
||||
in_mode_time = torch.where(in_contact, contact_time, air_time)
|
||||
single_stance = torch.mean(in_contact.float(), dim=1) == 0.5
|
||||
mode_time = torch.min(torch.where(single_stance.unsqueeze(-1), in_mode_time, 0.0), dim=1)[0]
|
||||
error = torch.abs(mode_time - threshold)
|
||||
reward = torch.clamp(threshold - error, min=0.0)
|
||||
if command_name is not None:
|
||||
command = env.command_manager.get_command(command_name)
|
||||
if command is not None:
|
||||
linear_norm = torch.norm(command[:, :2], dim=1)
|
||||
angular_norm = torch.abs(command[:, 2])
|
||||
total_command = linear_norm + angular_norm
|
||||
scale = (total_command > command_threshold).float()
|
||||
reward *= scale
|
||||
return reward
|
||||
|
||||
|
||||
def feet_clearance(
|
||||
env: ManagerBasedRlEnv,
|
||||
target_height: float,
|
||||
command_name: str | None = None,
|
||||
command_threshold: float = 0.1,
|
||||
asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG,
|
||||
) -> torch.Tensor:
|
||||
"""Penalize deviation from target clearance height, weighted by foot velocity."""
|
||||
asset: Entity = env.scene[asset_cfg.name]
|
||||
foot_z = asset.data.site_pos_w[:, asset_cfg.site_ids, 2] # [B, N]
|
||||
foot_vel_xy = asset.data.site_lin_vel_w[:, asset_cfg.site_ids, :2] # [B, N, 2]
|
||||
vel_norm = torch.norm(foot_vel_xy, dim=-1) # [B, N]
|
||||
delta = torch.abs(foot_z - target_height) # [B, N]
|
||||
cost = torch.sum(delta * vel_norm, dim=1) # [B]
|
||||
if command_name is not None:
|
||||
command = env.command_manager.get_command(command_name)
|
||||
if command is not None:
|
||||
linear_norm = torch.norm(command[:, :2], dim=1)
|
||||
angular_norm = torch.abs(command[:, 2])
|
||||
total_command = linear_norm + angular_norm
|
||||
active = (total_command > command_threshold).float()
|
||||
cost = cost * active
|
||||
return cost
|
||||
|
||||
|
||||
def feet_gait(
|
||||
env: ManagerBasedRlEnv,
|
||||
period: float,
|
||||
offset: list[float],
|
||||
threshold: float,
|
||||
command_threshold: float,
|
||||
command_name: str,
|
||||
sensor_name: str,
|
||||
) -> torch.Tensor:
|
||||
sensor: ContactSensor = env.scene[sensor_name]
|
||||
is_contact = sensor.data.current_contact_time > 0
|
||||
global_phase = ((env.episode_length_buf * env.step_dt) / period).unsqueeze(1)
|
||||
offsets = torch.as_tensor(offset, device=env.device, dtype=global_phase.dtype).view(1, -1)
|
||||
leg_phase = (global_phase + offsets) % 1.0
|
||||
is_stance = (leg_phase < threshold)
|
||||
reward = (is_stance == is_contact).float().mean(dim=1)
|
||||
if command_name is not None:
|
||||
command = env.command_manager.get_command(command_name)
|
||||
if command is not None:
|
||||
linear_norm = torch.norm(command[:, :2], dim=1)
|
||||
angular_norm = torch.abs(command[:, 2])
|
||||
total_command = linear_norm + angular_norm
|
||||
scale = (total_command > command_threshold).float()
|
||||
reward *= scale
|
||||
return reward
|
||||
|
||||
|
||||
class feet_swing_height:
|
||||
"""Penalize deviation from target swing height, evaluated at landing."""
|
||||
|
||||
def __init__(self, cfg: RewardTermCfg, env: ManagerBasedRlEnv):
|
||||
self.sensor_name = cfg.params["sensor_name"]
|
||||
self.site_names = cfg.params["asset_cfg"].site_names
|
||||
self.peak_heights = torch.zeros(
|
||||
(env.num_envs, len(self.site_names)), device=env.device, dtype=torch.float32
|
||||
)
|
||||
self.step_dt = env.step_dt
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
env: ManagerBasedRlEnv,
|
||||
sensor_name: str,
|
||||
target_height: float,
|
||||
command_name: str,
|
||||
command_threshold: float,
|
||||
asset_cfg: SceneEntityCfg,
|
||||
) -> torch.Tensor:
|
||||
asset: Entity = env.scene[asset_cfg.name]
|
||||
contact_sensor: ContactSensor = env.scene[sensor_name]
|
||||
command = env.command_manager.get_command(command_name)
|
||||
assert command is not None
|
||||
foot_heights = asset.data.site_pos_w[:, asset_cfg.site_ids, 2]
|
||||
in_air = contact_sensor.data.found == 0
|
||||
self.peak_heights = torch.where(
|
||||
in_air,
|
||||
torch.maximum(self.peak_heights, foot_heights),
|
||||
self.peak_heights,
|
||||
)
|
||||
first_contact = contact_sensor.compute_first_contact(dt=self.step_dt)
|
||||
linear_norm = torch.norm(command[:, :2], dim=1)
|
||||
angular_norm = torch.abs(command[:, 2])
|
||||
total_command = linear_norm + angular_norm
|
||||
active = (total_command > command_threshold).float()
|
||||
error = self.peak_heights / target_height - 1.0
|
||||
cost = torch.sum(torch.square(error) * first_contact.float(), dim=1) * active
|
||||
num_landings = torch.sum(first_contact.float())
|
||||
peak_heights_at_landing = self.peak_heights * first_contact.float()
|
||||
mean_peak_height = torch.sum(peak_heights_at_landing) / torch.clamp(
|
||||
num_landings, min=1
|
||||
)
|
||||
env.extras["log"]["Metrics/peak_height_mean"] = mean_peak_height
|
||||
self.peak_heights = torch.where(
|
||||
first_contact,
|
||||
torch.zeros_like(self.peak_heights),
|
||||
self.peak_heights,
|
||||
)
|
||||
return cost
|
||||
|
||||
|
||||
def feet_slip(
|
||||
env: ManagerBasedRlEnv,
|
||||
sensor_name: str,
|
||||
command_name: str,
|
||||
command_threshold: float = 0.01,
|
||||
asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG,
|
||||
) -> torch.Tensor:
|
||||
"""Penalize foot sliding (xy velocity while in contact)."""
|
||||
asset: Entity = env.scene[asset_cfg.name]
|
||||
contact_sensor: ContactSensor = env.scene[sensor_name]
|
||||
command = env.command_manager.get_command(command_name)
|
||||
assert command is not None
|
||||
linear_norm = torch.norm(command[:, :2], dim=1)
|
||||
angular_norm = torch.abs(command[:, 2])
|
||||
total_command = linear_norm + angular_norm
|
||||
active = (total_command > command_threshold).float()
|
||||
assert contact_sensor.data.found is not None
|
||||
in_contact = (contact_sensor.data.found > 0).float() # [B, N]
|
||||
foot_vel_xy = asset.data.site_lin_vel_w[:, asset_cfg.site_ids, :2] # [B, N, 2]
|
||||
vel_xy_norm = torch.norm(foot_vel_xy, dim=-1) # [B, N]
|
||||
vel_xy_norm_sq = torch.square(vel_xy_norm) # [B, N]
|
||||
cost = torch.sum(vel_xy_norm_sq * in_contact, dim=1) * active
|
||||
num_in_contact = torch.sum(in_contact)
|
||||
mean_slip_vel = torch.sum(vel_xy_norm * in_contact) / torch.clamp(
|
||||
num_in_contact, min=1
|
||||
)
|
||||
env.extras["log"]["Metrics/slip_velocity_mean"] = mean_slip_vel
|
||||
return cost
|
||||
|
||||
|
||||
def soft_landing(
|
||||
env: ManagerBasedRlEnv,
|
||||
sensor_name: str,
|
||||
command_name: str | None = None,
|
||||
command_threshold: float = 0.05,
|
||||
) -> torch.Tensor:
|
||||
"""Penalize high impact forces at landing to encourage soft footfalls."""
|
||||
contact_sensor: ContactSensor = env.scene[sensor_name]
|
||||
sensor_data = contact_sensor.data
|
||||
assert sensor_data.force is not None
|
||||
forces = sensor_data.force # [B, N, 3]
|
||||
force_magnitude = torch.norm(forces, dim=-1) # [B, N]
|
||||
first_contact = contact_sensor.compute_first_contact(dt=env.step_dt) # [B, N]
|
||||
landing_impact = force_magnitude * first_contact.float() # [B, N]
|
||||
cost = torch.sum(landing_impact, dim=1) # [B]
|
||||
num_landings = torch.sum(first_contact.float())
|
||||
mean_landing_force = torch.sum(landing_impact) / torch.clamp(num_landings, min=1)
|
||||
env.extras["log"]["Metrics/landing_force_mean"] = mean_landing_force
|
||||
if command_name is not None:
|
||||
command = env.command_manager.get_command(command_name)
|
||||
if command is not None:
|
||||
linear_norm = torch.norm(command[:, :2], dim=1)
|
||||
angular_norm = torch.abs(command[:, 2])
|
||||
total_command = linear_norm + angular_norm
|
||||
active = (total_command > command_threshold).float()
|
||||
cost = cost * active
|
||||
return cost
|
||||
|
||||
|
||||
class variable_posture:
|
||||
"""Penalize deviation from default pose with speed-dependent tolerance.
|
||||
|
||||
Uses per-joint standard deviations to control how much each joint can deviate
|
||||
from default pose. Smaller std = stricter (less deviation allowed), larger
|
||||
std = more forgiving. The reward is: exp(-mean(error² / std²))
|
||||
|
||||
Three speed regimes (based on linear + angular command velocity):
|
||||
- std_standing (speed < walking_threshold): Tight tolerance for holding pose.
|
||||
- std_walking (walking_threshold <= speed < running_threshold): Moderate.
|
||||
- std_running (speed >= running_threshold): Loose tolerance for large motion.
|
||||
|
||||
Tune std values per joint based on how much motion that joint needs at each
|
||||
speed. Map joint name patterns to std values, e.g. {".*knee.*": 0.35}.
|
||||
"""
|
||||
|
||||
def __init__(self, cfg: RewardTermCfg, env: ManagerBasedRlEnv):
|
||||
asset: Entity = env.scene[cfg.params["asset_cfg"].name]
|
||||
default_joint_pos = asset.data.default_joint_pos
|
||||
assert default_joint_pos is not None
|
||||
self.default_joint_pos = default_joint_pos
|
||||
|
||||
_, joint_names = asset.find_joints(cfg.params["asset_cfg"].joint_names)
|
||||
|
||||
_, _, std_standing = resolve_matching_names_values(
|
||||
data=cfg.params["std_standing"],
|
||||
list_of_strings=joint_names,
|
||||
)
|
||||
self.std_standing = torch.tensor(
|
||||
std_standing, device=env.device, dtype=torch.float32
|
||||
)
|
||||
|
||||
_, _, std_walking = resolve_matching_names_values(
|
||||
data=cfg.params["std_walking"],
|
||||
list_of_strings=joint_names,
|
||||
)
|
||||
self.std_walking = torch.tensor(std_walking, device=env.device, dtype=torch.float32)
|
||||
|
||||
_, _, std_running = resolve_matching_names_values(
|
||||
data=cfg.params["std_running"],
|
||||
list_of_strings=joint_names,
|
||||
)
|
||||
self.std_running = torch.tensor(std_running, device=env.device, dtype=torch.float32)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
env: ManagerBasedRlEnv,
|
||||
std_standing,
|
||||
std_walking,
|
||||
std_running,
|
||||
asset_cfg: SceneEntityCfg,
|
||||
command_name: str,
|
||||
walking_threshold: float = 0.5,
|
||||
running_threshold: float = 1.5,
|
||||
) -> torch.Tensor:
|
||||
del std_standing, std_walking, std_running # Unused.
|
||||
|
||||
asset: Entity = env.scene[asset_cfg.name]
|
||||
command = env.command_manager.get_command(command_name)
|
||||
assert command is not None
|
||||
|
||||
linear_speed = torch.norm(command[:, :2], dim=1)
|
||||
angular_speed = torch.abs(command[:, 2])
|
||||
total_speed = linear_speed + angular_speed
|
||||
|
||||
standing_mask = (total_speed < walking_threshold).float()
|
||||
walking_mask = (
|
||||
(total_speed >= walking_threshold) & (total_speed < running_threshold)
|
||||
).float()
|
||||
running_mask = (total_speed >= running_threshold).float()
|
||||
|
||||
std = (
|
||||
self.std_standing * standing_mask.unsqueeze(1)
|
||||
+ self.std_walking * walking_mask.unsqueeze(1)
|
||||
+ self.std_running * running_mask.unsqueeze(1)
|
||||
)
|
||||
|
||||
current_joint_pos = asset.data.joint_pos[:, asset_cfg.joint_ids]
|
||||
desired_joint_pos = self.default_joint_pos[:, asset_cfg.joint_ids]
|
||||
error_squared = torch.square(current_joint_pos - desired_joint_pos)
|
||||
|
||||
return torch.exp(-torch.mean(error_squared / (std**2), dim=1))
|
||||
|
||||
|
||||
def stand_still(
|
||||
env: ManagerBasedRlEnv,
|
||||
command_name: str,
|
||||
command_threshold: float = 0.1,
|
||||
asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG
|
||||
) -> torch.Tensor:
|
||||
asset: Entity = env.scene[asset_cfg.name]
|
||||
diff_angle = asset.data.joint_pos[:, asset_cfg.joint_ids] - asset.data.default_joint_pos[:, asset_cfg.joint_ids]
|
||||
reward = torch.sum(torch.square(diff_angle), dim=1)
|
||||
if command_name is not None:
|
||||
command = env.command_manager.get_command(command_name)
|
||||
if command is not None:
|
||||
linear_norm = torch.norm(command[:, :2], dim=1)
|
||||
angular_norm = torch.abs(command[:, 2])
|
||||
total_command = linear_norm + angular_norm
|
||||
scale = (total_command <= command_threshold).float()
|
||||
reward *= scale
|
||||
return reward
|
||||
@@ -0,0 +1,25 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
from mjlab.sensor import ContactSensor
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from mjlab.envs import ManagerBasedRlEnv
|
||||
|
||||
|
||||
def illegal_contact(
|
||||
env: ManagerBasedRlEnv,
|
||||
sensor_name: str,
|
||||
force_threshold: float = 10.0,
|
||||
) -> torch.Tensor:
|
||||
sensor: ContactSensor = env.scene[sensor_name]
|
||||
data = sensor.data
|
||||
if data.force_history is not None:
|
||||
# force_history: [B, N, H, 3]
|
||||
force_mag = torch.norm(data.force_history, dim=-1) # [B, N, H]
|
||||
return (force_mag > force_threshold).any(dim=-1).any(dim=-1) # [B]
|
||||
assert data.found is not None
|
||||
return torch.any(data.found, dim=-1)
|
||||
@@ -0,0 +1,282 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass, field
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from mjlab.entity import Entity
|
||||
from mjlab.managers.command_manager import CommandTerm, CommandTermCfg
|
||||
from mjlab.utils.lab_api.math import (
|
||||
matrix_from_quat,
|
||||
quat_apply,
|
||||
wrap_to_pi,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import viser
|
||||
|
||||
from mjlab.envs.manager_based_rl_env import ManagerBasedRlEnv
|
||||
from mjlab.viewer.debug_visualizer import DebugVisualizer
|
||||
|
||||
|
||||
class UniformVelocityCommand(CommandTerm):
|
||||
cfg: UniformVelocityCommandCfg
|
||||
|
||||
def __init__(self, cfg: UniformVelocityCommandCfg, env: ManagerBasedRlEnv):
|
||||
super().__init__(cfg, env)
|
||||
|
||||
if self.cfg.heading_command and self.cfg.ranges.heading is None:
|
||||
raise ValueError("heading_command=True but ranges.heading is set to None.")
|
||||
if self.cfg.ranges.heading and not self.cfg.heading_command:
|
||||
raise ValueError("ranges.heading is set but heading_command=False.")
|
||||
|
||||
self.robot: Entity = env.scene[cfg.entity_name]
|
||||
|
||||
self.vel_command_b = torch.zeros(self.num_envs, 3, device=self.device)
|
||||
self.heading_target = torch.zeros(self.num_envs, device=self.device)
|
||||
self.heading_error = torch.zeros(self.num_envs, device=self.device)
|
||||
self.is_heading_env = torch.zeros(
|
||||
self.num_envs, dtype=torch.bool, device=self.device
|
||||
)
|
||||
self.is_standing_env = torch.zeros_like(self.is_heading_env)
|
||||
|
||||
self.metrics["error_vel_xy"] = torch.zeros(self.num_envs, device=self.device)
|
||||
self.metrics["error_vel_yaw"] = torch.zeros(self.num_envs, device=self.device)
|
||||
|
||||
# Set by create_gui() when the viewer is active.
|
||||
self._joystick_enabled: viser.GuiCheckboxHandle | None = None
|
||||
self._joystick_sliders: list[viser.GuiSliderHandle] = []
|
||||
self._joystick_get_env_idx: Callable[[], int] | None = None
|
||||
|
||||
@property
|
||||
def command(self) -> torch.Tensor:
|
||||
return self.vel_command_b
|
||||
|
||||
def _update_metrics(self) -> None:
|
||||
max_command_time = self.cfg.resampling_time_range[1]
|
||||
max_command_step = max_command_time / self._env.step_dt
|
||||
self.metrics["error_vel_xy"] += (
|
||||
torch.norm(
|
||||
self.vel_command_b[:, :2] - self.robot.data.root_link_lin_vel_b[:, :2], dim=-1
|
||||
)
|
||||
/ max_command_step
|
||||
)
|
||||
self.metrics["error_vel_yaw"] += (
|
||||
torch.abs(self.vel_command_b[:, 2] - self.robot.data.root_link_ang_vel_b[:, 2])
|
||||
/ max_command_step
|
||||
)
|
||||
|
||||
def _resample_command(self, env_ids: torch.Tensor) -> None:
|
||||
r = torch.empty(len(env_ids), device=self.device)
|
||||
self.vel_command_b[env_ids, 0] = r.uniform_(*self.cfg.ranges.lin_vel_x)
|
||||
self.vel_command_b[env_ids, 1] = r.uniform_(*self.cfg.ranges.lin_vel_y)
|
||||
self.vel_command_b[env_ids, 2] = r.uniform_(*self.cfg.ranges.ang_vel_z)
|
||||
self.vel_command_b[env_ids, :] *= (torch.norm(self.vel_command_b[env_ids, :], dim=1) > 0.1).unsqueeze(1)
|
||||
if self.cfg.heading_command:
|
||||
assert self.cfg.ranges.heading is not None
|
||||
self.heading_target[env_ids] = r.uniform_(*self.cfg.ranges.heading)
|
||||
self.is_heading_env[env_ids] = r.uniform_(0.0, 1.0) <= self.cfg.rel_heading_envs
|
||||
self.is_standing_env[env_ids] = r.uniform_(0.0, 1.0) <= self.cfg.rel_standing_envs
|
||||
|
||||
init_vel_mask = r.uniform_(0.0, 1.0) < self.cfg.init_velocity_prob
|
||||
init_vel_env_ids = env_ids[init_vel_mask]
|
||||
if len(init_vel_env_ids) > 0:
|
||||
root_pos = self.robot.data.root_link_pos_w[init_vel_env_ids]
|
||||
root_quat = self.robot.data.root_link_quat_w[init_vel_env_ids]
|
||||
lin_vel_b = self.robot.data.root_link_lin_vel_b[init_vel_env_ids]
|
||||
lin_vel_b[:, :2] = self.vel_command_b[init_vel_env_ids, :2]
|
||||
root_lin_vel_w = quat_apply(root_quat, lin_vel_b)
|
||||
root_ang_vel_b = self.robot.data.root_link_ang_vel_b[init_vel_env_ids]
|
||||
root_ang_vel_b[:, 2] = self.vel_command_b[init_vel_env_ids, 2]
|
||||
root_state = torch.cat(
|
||||
[root_pos, root_quat, root_lin_vel_w, root_ang_vel_b], dim=-1
|
||||
)
|
||||
self.robot.write_root_state_to_sim(root_state, init_vel_env_ids)
|
||||
|
||||
def _update_command(self) -> None:
|
||||
if self.cfg.heading_command:
|
||||
self.heading_error = wrap_to_pi(self.heading_target - self.robot.data.heading_w)
|
||||
env_ids = self.is_heading_env.nonzero(as_tuple=False).flatten()
|
||||
self.vel_command_b[env_ids, 2] = torch.clip(
|
||||
self.cfg.heading_control_stiffness * self.heading_error[env_ids],
|
||||
min=self.cfg.ranges.ang_vel_z[0],
|
||||
max=self.cfg.ranges.ang_vel_z[1],
|
||||
)
|
||||
standing_env_ids = self.is_standing_env.nonzero(as_tuple=False).flatten()
|
||||
self.vel_command_b[standing_env_ids, :] = 0.0
|
||||
|
||||
# GUI.
|
||||
|
||||
def create_gui(
|
||||
self,
|
||||
name: str,
|
||||
server: "viser.ViserServer",
|
||||
get_env_idx: Callable[[], int],
|
||||
) -> None:
|
||||
"""Create velocity joystick sliders in the Viser viewer."""
|
||||
from viser import Icon
|
||||
|
||||
ranges = self.cfg.ranges
|
||||
|
||||
axes = [
|
||||
("lin_vel_x", ranges.lin_vel_x[1]),
|
||||
("lin_vel_y", ranges.lin_vel_y[1]),
|
||||
("ang_vel_z", ranges.ang_vel_z[1]),
|
||||
]
|
||||
sliders: list = []
|
||||
|
||||
with server.gui.add_folder(name.capitalize()):
|
||||
enabled = server.gui.add_checkbox("Enable", initial_value=False)
|
||||
|
||||
for label, max_val in axes:
|
||||
max_input = server.gui.add_slider(
|
||||
f"Max {label}",
|
||||
initial_value=max_val,
|
||||
step=0.1,
|
||||
min=0.1,
|
||||
max=10.0,
|
||||
)
|
||||
slider = server.gui.add_slider(
|
||||
label,
|
||||
min=-max_val,
|
||||
max=max_val,
|
||||
step=0.05,
|
||||
initial_value=0.0,
|
||||
)
|
||||
|
||||
@max_input.on_update
|
||||
def _(_ev, _s=slider, _m=max_input) -> None:
|
||||
_s.min = -_m.value
|
||||
_s.max = _m.value
|
||||
|
||||
sliders.append(slider)
|
||||
|
||||
zero_btn = server.gui.add_button("Zero", icon=Icon.SQUARE_X)
|
||||
|
||||
@zero_btn.on_click
|
||||
def _(_) -> None:
|
||||
for s in sliders:
|
||||
s.value = 0.0
|
||||
|
||||
# Store GUI state for compute() override.
|
||||
self._joystick_enabled = enabled
|
||||
self._joystick_sliders = sliders
|
||||
self._joystick_get_env_idx = get_env_idx
|
||||
|
||||
def compute(self, dt: float) -> None:
|
||||
super().compute(dt)
|
||||
if self._joystick_enabled is not None and self._joystick_enabled.value:
|
||||
assert self._joystick_get_env_idx is not None
|
||||
idx = self._joystick_get_env_idx()
|
||||
for i, s in enumerate(self._joystick_sliders):
|
||||
self.vel_command_b[idx, i] = s.value
|
||||
|
||||
# Visualization.
|
||||
|
||||
def _debug_vis_impl(self, visualizer: "DebugVisualizer") -> None:
|
||||
"""Draw velocity command and actual velocity arrows."""
|
||||
env_indices = visualizer.get_env_indices(self.num_envs)
|
||||
if not env_indices:
|
||||
return
|
||||
|
||||
cmds = self.command.cpu().numpy()
|
||||
base_pos_ws = self.robot.data.root_link_pos_w.cpu().numpy()
|
||||
base_quat_w = self.robot.data.root_link_quat_w
|
||||
base_mat_ws = matrix_from_quat(base_quat_w).cpu().numpy()
|
||||
lin_vel_bs = self.robot.data.root_link_lin_vel_b.cpu().numpy()
|
||||
ang_vel_bs = self.robot.data.root_link_ang_vel_b.cpu().numpy()
|
||||
|
||||
scale = self.cfg.viz.scale
|
||||
z_offset = self.cfg.viz.z_offset
|
||||
|
||||
for batch in env_indices:
|
||||
base_pos_w = base_pos_ws[batch]
|
||||
base_mat_w = base_mat_ws[batch]
|
||||
cmd = cmds[batch]
|
||||
lin_vel_b = lin_vel_bs[batch]
|
||||
ang_vel_b = ang_vel_bs[batch]
|
||||
|
||||
# Skip if robot appears uninitialized (at origin).
|
||||
if np.linalg.norm(base_pos_w) < 1e-6:
|
||||
continue
|
||||
|
||||
# Helper to transform local to world coordinates.
|
||||
def local_to_world(
|
||||
vec: np.ndarray, pos: np.ndarray = base_pos_w, mat: np.ndarray = base_mat_w
|
||||
) -> np.ndarray:
|
||||
return pos + mat @ vec
|
||||
|
||||
# Command linear velocity arrow (blue).
|
||||
cmd_lin_from = local_to_world(np.array([0, 0, z_offset]) * scale)
|
||||
cmd_lin_to = local_to_world(
|
||||
(np.array([0, 0, z_offset]) + np.array([cmd[0], cmd[1], 0])) * scale
|
||||
)
|
||||
visualizer.add_arrow(
|
||||
cmd_lin_from, cmd_lin_to, color=(0.2, 0.2, 0.6, 0.6), width=0.015
|
||||
)
|
||||
|
||||
# Command angular velocity arrow (green).
|
||||
cmd_ang_from = cmd_lin_from
|
||||
cmd_ang_to = local_to_world(
|
||||
(np.array([0, 0, z_offset]) + np.array([0, 0, cmd[2]])) * scale
|
||||
)
|
||||
visualizer.add_arrow(
|
||||
cmd_ang_from, cmd_ang_to, color=(0.2, 0.6, 0.2, 0.6), width=0.015
|
||||
)
|
||||
|
||||
# Actual linear velocity arrow (cyan).
|
||||
act_lin_from = local_to_world(np.array([0, 0, z_offset]) * scale)
|
||||
act_lin_to = local_to_world(
|
||||
(np.array([0, 0, z_offset]) + np.array([lin_vel_b[0], lin_vel_b[1], 0])) * scale
|
||||
)
|
||||
visualizer.add_arrow(
|
||||
act_lin_from, act_lin_to, color=(0.0, 0.6, 1.0, 0.7), width=0.015
|
||||
)
|
||||
|
||||
# Actual angular velocity arrow (light green).
|
||||
act_ang_from = act_lin_from
|
||||
act_ang_to = local_to_world(
|
||||
(np.array([0, 0, z_offset]) + np.array([0, 0, ang_vel_b[2]])) * scale
|
||||
)
|
||||
visualizer.add_arrow(
|
||||
act_ang_from, act_ang_to, color=(0.0, 1.0, 0.4, 0.7), width=0.015
|
||||
)
|
||||
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
class UniformVelocityCommandCfg(CommandTermCfg):
|
||||
entity_name: str
|
||||
heading_command: bool = False
|
||||
heading_control_stiffness: float = 1.0
|
||||
rel_standing_envs: float = 0.0
|
||||
rel_heading_envs: float = 1.0
|
||||
init_velocity_prob: float = 0.0
|
||||
|
||||
@dataclass
|
||||
class Ranges:
|
||||
lin_vel_x: tuple[float, float]
|
||||
lin_vel_y: tuple[float, float]
|
||||
ang_vel_z: tuple[float, float]
|
||||
heading: tuple[float, float] | None = None
|
||||
|
||||
ranges: Ranges
|
||||
|
||||
@dataclass
|
||||
class VizCfg:
|
||||
z_offset: float = 0.2
|
||||
scale: float = 0.5
|
||||
|
||||
viz: VizCfg = field(default_factory=VizCfg)
|
||||
|
||||
def build(self, env: ManagerBasedRlEnv) -> UniformVelocityCommand:
|
||||
return UniformVelocityCommand(self, env)
|
||||
|
||||
def __post_init__(self):
|
||||
if self.heading_command and self.ranges.heading is None:
|
||||
raise ValueError(
|
||||
"The velocity command has heading commands active (heading_command=True) but "
|
||||
"the `ranges.heading` parameter is set to None."
|
||||
)
|
||||
@@ -0,0 +1,3 @@
|
||||
from .runner import (
|
||||
VelocityOnPolicyRunner as VelocityOnPolicyRunner,
|
||||
)
|
||||
@@ -0,0 +1,28 @@
|
||||
import os
|
||||
|
||||
import wandb
|
||||
|
||||
from mjlab.rl import RslRlVecEnvWrapper
|
||||
from mjlab.rl.exporter_utils import (
|
||||
attach_metadata_to_onnx,
|
||||
get_base_metadata,
|
||||
)
|
||||
from mjlab.rl.runner import MjlabOnPolicyRunner
|
||||
|
||||
|
||||
class VelocityOnPolicyRunner(MjlabOnPolicyRunner):
|
||||
env: RslRlVecEnvWrapper
|
||||
|
||||
def save(self, path: str, infos=None):
|
||||
super().save(path, infos)
|
||||
policy_path = path.split("model")[0]
|
||||
filename = "policy.onnx"
|
||||
self.export_policy_to_onnx(policy_path, filename)
|
||||
run_name: str = (
|
||||
wandb.run.name if self.logger.logger_type == "wandb" and wandb.run else "local"
|
||||
) # type: ignore[assignment]
|
||||
onnx_path = os.path.join(policy_path, filename)
|
||||
metadata = get_base_metadata(self.env.unwrapped, run_name)
|
||||
attach_metadata_to_onnx(onnx_path, metadata)
|
||||
if self.logger.logger_type in ["wandb"]:
|
||||
wandb.save(policy_path + filename, base_path=os.path.dirname(policy_path))
|
||||
@@ -0,0 +1,431 @@
|
||||
"""Velocity task configuration.
|
||||
|
||||
This module provides a factory function to create a base velocity task config.
|
||||
Robot-specific configurations call the factory and customize as needed.
|
||||
"""
|
||||
|
||||
import math
|
||||
from dataclasses import replace
|
||||
|
||||
from mjlab.envs import ManagerBasedRlEnvCfg
|
||||
from mjlab.envs import mdp as envs_mdp
|
||||
from mjlab.envs.mdp import dr
|
||||
from mjlab.envs.mdp.actions import JointPositionActionCfg
|
||||
from mjlab.managers.action_manager import ActionTermCfg
|
||||
from mjlab.managers.command_manager import CommandTermCfg
|
||||
from mjlab.managers.curriculum_manager import CurriculumTermCfg
|
||||
from mjlab.managers.event_manager import EventTermCfg
|
||||
from mjlab.managers.metrics_manager import MetricsTermCfg
|
||||
from mjlab.managers.observation_manager import ObservationGroupCfg, ObservationTermCfg
|
||||
from mjlab.managers.reward_manager import RewardTermCfg
|
||||
from mjlab.managers.scene_entity_config import SceneEntityCfg
|
||||
from mjlab.managers.termination_manager import TerminationTermCfg
|
||||
from mjlab.scene import SceneCfg
|
||||
from mjlab.sensor import GridPatternCfg, ObjRef, RayCastSensorCfg
|
||||
from mjlab.sim import MujocoCfg, SimulationCfg
|
||||
from mjlab.tasks.velocity import mdp
|
||||
from mjlab.tasks.velocity.mdp import UniformVelocityCommandCfg
|
||||
from mjlab.terrains import TerrainEntityCfg
|
||||
from mjlab.terrains.config import ROUGH_TERRAINS_CFG
|
||||
from mjlab.utils.noise import UniformNoiseCfg as Unoise
|
||||
from mjlab.viewer import ViewerConfig
|
||||
|
||||
import src.tasks.velocity.mdp as mdp
|
||||
|
||||
|
||||
def make_velocity_env_cfg() -> ManagerBasedRlEnvCfg:
|
||||
"""Create base velocity tracking task configuration."""
|
||||
|
||||
##
|
||||
# Sensors
|
||||
##
|
||||
|
||||
terrain_scan = RayCastSensorCfg(
|
||||
name="terrain_scan",
|
||||
frame=ObjRef(type="body", name="", entity="robot"), # Set per-robot.
|
||||
ray_alignment="yaw",
|
||||
pattern=GridPatternCfg(size=(1.6, 1.0), resolution=0.1),
|
||||
max_distance=5.0,
|
||||
exclude_parent_body=True,
|
||||
debug_vis=True,
|
||||
viz=RayCastSensorCfg.VizCfg(show_normals=True),
|
||||
)
|
||||
|
||||
##
|
||||
# Observations
|
||||
##
|
||||
|
||||
actor_terms = {
|
||||
"base_ang_vel": ObservationTermCfg(
|
||||
func=mdp.builtin_sensor,
|
||||
params={"sensor_name": "robot/imu_ang_vel"},
|
||||
noise=Unoise(n_min=-0.2, n_max=0.2),
|
||||
),
|
||||
"projected_gravity": ObservationTermCfg(
|
||||
func=mdp.projected_gravity,
|
||||
noise=Unoise(n_min=-0.05, n_max=0.05),
|
||||
),
|
||||
"command": ObservationTermCfg(
|
||||
func=mdp.generated_commands,
|
||||
params={"command_name": "twist"},
|
||||
),
|
||||
"phase": ObservationTermCfg(
|
||||
func=mdp.phase,
|
||||
params={"period": 0.6, "command_name": "twist"},
|
||||
),
|
||||
"joint_pos": ObservationTermCfg(
|
||||
func=mdp.joint_pos_rel,
|
||||
noise=Unoise(n_min=-0.01, n_max=0.01),
|
||||
),
|
||||
"joint_vel": ObservationTermCfg(
|
||||
func=mdp.joint_vel_rel,
|
||||
noise=Unoise(n_min=-1.5, n_max=1.5),
|
||||
),
|
||||
"actions": ObservationTermCfg(func=mdp.last_action),
|
||||
"height_scan": ObservationTermCfg(
|
||||
func=envs_mdp.height_scan,
|
||||
params={"sensor_name": "terrain_scan"},
|
||||
noise=Unoise(n_min=-0.1, n_max=0.1),
|
||||
scale=1 / terrain_scan.max_distance,
|
||||
),
|
||||
}
|
||||
|
||||
critic_terms = {
|
||||
**actor_terms,
|
||||
"base_lin_vel": ObservationTermCfg(
|
||||
func=mdp.builtin_sensor,
|
||||
params={"sensor_name": "robot/imu_lin_vel"},
|
||||
noise=Unoise(n_min=-0.5, n_max=0.5),
|
||||
),
|
||||
"height_scan": ObservationTermCfg(
|
||||
func=envs_mdp.height_scan,
|
||||
params={"sensor_name": "terrain_scan"},
|
||||
scale=1 / terrain_scan.max_distance,
|
||||
),
|
||||
"foot_height": ObservationTermCfg(
|
||||
func=mdp.foot_height,
|
||||
params={"asset_cfg": SceneEntityCfg("robot", site_names=())}, # Set per-robot.
|
||||
),
|
||||
"foot_air_time": ObservationTermCfg(
|
||||
func=mdp.foot_air_time,
|
||||
params={"sensor_name": "feet_ground_contact"},
|
||||
),
|
||||
"foot_contact": ObservationTermCfg(
|
||||
func=mdp.foot_contact,
|
||||
params={"sensor_name": "feet_ground_contact"},
|
||||
),
|
||||
"foot_contact_forces": ObservationTermCfg(
|
||||
func=mdp.foot_contact_forces,
|
||||
params={"sensor_name": "feet_ground_contact"},
|
||||
),
|
||||
}
|
||||
|
||||
observations = {
|
||||
"actor": ObservationGroupCfg(
|
||||
terms=actor_terms,
|
||||
concatenate_terms=True,
|
||||
enable_corruption=True,
|
||||
history_length=1,
|
||||
),
|
||||
"critic": ObservationGroupCfg(
|
||||
terms=critic_terms,
|
||||
concatenate_terms=True,
|
||||
enable_corruption=False,
|
||||
history_length=1,
|
||||
),
|
||||
}
|
||||
|
||||
##
|
||||
# Metrics
|
||||
##
|
||||
|
||||
metrics = {
|
||||
"mean_action_acc": MetricsTermCfg(
|
||||
func=mdp.mean_action_acc,
|
||||
),
|
||||
}
|
||||
|
||||
##
|
||||
# Actions
|
||||
##
|
||||
|
||||
actions: dict[str, ActionTermCfg] = {
|
||||
"joint_pos": JointPositionActionCfg(
|
||||
entity_name="robot",
|
||||
actuator_names=(".*",),
|
||||
scale=0.25, # Override per-robot.
|
||||
use_default_offset=True,
|
||||
)
|
||||
}
|
||||
|
||||
##
|
||||
# Commands
|
||||
##
|
||||
|
||||
commands: dict[str, CommandTermCfg] = {
|
||||
"twist": UniformVelocityCommandCfg(
|
||||
entity_name="robot",
|
||||
resampling_time_range=(3.0, 8.0),
|
||||
rel_standing_envs=0.05,
|
||||
heading_command=True,
|
||||
heading_control_stiffness=0.5,
|
||||
debug_vis=True,
|
||||
ranges=UniformVelocityCommandCfg.Ranges(
|
||||
lin_vel_x=(-1.0, 2.0),
|
||||
lin_vel_y=(-1.0, 1.0),
|
||||
ang_vel_z=(-1.0, 1.0),
|
||||
heading=(-math.pi, math.pi),
|
||||
),
|
||||
)
|
||||
}
|
||||
|
||||
##
|
||||
# Events
|
||||
##
|
||||
|
||||
events = {
|
||||
"reset_base": EventTermCfg(
|
||||
func=mdp.reset_root_state_uniform,
|
||||
mode="reset",
|
||||
params={
|
||||
"pose_range": {
|
||||
"x": (-0.5, 0.5),
|
||||
"y": (-0.5, 0.5),
|
||||
"z": (0.0, 0.0),
|
||||
"yaw": (-3.14, 3.14),
|
||||
},
|
||||
"velocity_range": {},
|
||||
},
|
||||
),
|
||||
"reset_robot_joints": EventTermCfg(
|
||||
func=mdp.reset_joints_by_offset,
|
||||
mode="reset",
|
||||
params={
|
||||
"position_range": (-0.0, 0.0),
|
||||
"velocity_range": (-0.0, 0.0),
|
||||
"asset_cfg": SceneEntityCfg("robot", joint_names=(".*",)),
|
||||
},
|
||||
),
|
||||
"push_robot": EventTermCfg(
|
||||
func=mdp.push_by_setting_velocity,
|
||||
mode="interval",
|
||||
interval_range_s=(5.0, 6.0),
|
||||
params={
|
||||
"velocity_range": {
|
||||
"x": (-0.5, 0.5),
|
||||
"y": (-0.5, 0.5),
|
||||
"z": (-0.4, 0.4),
|
||||
"roll": (-0.52, 0.52),
|
||||
"pitch": (-0.52, 0.52),
|
||||
"yaw": (-0.78, 0.78),
|
||||
},
|
||||
},
|
||||
),
|
||||
"foot_friction": EventTermCfg(
|
||||
mode="startup",
|
||||
func=dr.geom_friction,
|
||||
params={
|
||||
"asset_cfg": SceneEntityCfg("robot", geom_names=()), # Set per-robot.
|
||||
"operation": "abs",
|
||||
"ranges": (0.3, 1.6),
|
||||
"shared_random": True, # All foot geoms share the same friction.
|
||||
},
|
||||
),
|
||||
"encoder_bias": EventTermCfg(
|
||||
mode="startup",
|
||||
func=dr.encoder_bias,
|
||||
params={
|
||||
"asset_cfg": SceneEntityCfg("robot"),
|
||||
"bias_range": (-0.015, 0.015),
|
||||
},
|
||||
),
|
||||
"base_com": EventTermCfg(
|
||||
mode="startup",
|
||||
func=dr.body_com_offset,
|
||||
params={
|
||||
"asset_cfg": SceneEntityCfg("robot", body_names=()), # Set per-robot.
|
||||
"operation": "add",
|
||||
"ranges": {
|
||||
0: (-0.05, 0.05),
|
||||
1: (-0.05, 0.05),
|
||||
2: (-0.05, 0.05),
|
||||
},
|
||||
},
|
||||
),
|
||||
}
|
||||
|
||||
##
|
||||
# Rewards
|
||||
##
|
||||
|
||||
rewards = {
|
||||
"track_linear_velocity": RewardTermCfg(
|
||||
func=mdp.track_linear_velocity,
|
||||
weight=1.0,
|
||||
params={"command_name": "twist", "std": math.sqrt(0.25)},
|
||||
),
|
||||
"track_angular_velocity": RewardTermCfg(
|
||||
func=mdp.track_angular_velocity,
|
||||
weight=1.0,
|
||||
params={"command_name": "twist", "std": math.sqrt(0.5)},
|
||||
),
|
||||
"body_orientation_l2": RewardTermCfg(
|
||||
func=mdp.body_orientation_l2,
|
||||
weight=-1.0,
|
||||
params={"asset_cfg": SceneEntityCfg("robot", body_names=())}, # Set per-robot.
|
||||
),
|
||||
"pose": RewardTermCfg(
|
||||
func=mdp.variable_posture,
|
||||
weight=1.0,
|
||||
params={
|
||||
"asset_cfg": SceneEntityCfg("robot", joint_names=".*"),
|
||||
"command_name": "twist",
|
||||
"std_standing": {}, # Set per-robot.
|
||||
"std_walking": {}, # Set per-robot.
|
||||
"std_running": {}, # Set per-robot.
|
||||
"walking_threshold": 0.1,
|
||||
"running_threshold": 1.5,
|
||||
},
|
||||
),
|
||||
"body_ang_vel": RewardTermCfg(
|
||||
func=mdp.body_angular_velocity_penalty,
|
||||
weight=-0.05, # Override per-robot
|
||||
params={"asset_cfg": SceneEntityCfg("robot", body_names=())}, # Set per-robot.
|
||||
),
|
||||
"angular_momentum": RewardTermCfg(
|
||||
func=mdp.angular_momentum_penalty,
|
||||
weight=-0.025, # Override per-robot
|
||||
params={"sensor_name": "robot/root_angmom"},
|
||||
),
|
||||
"is_terminated": RewardTermCfg(func=mdp.is_terminated, weight=-200.0),
|
||||
"joint_acc_l2": RewardTermCfg(func=mdp.joint_acc_l2, weight=-2.5e-7),
|
||||
"joint_pos_limits": RewardTermCfg(func=mdp.joint_pos_limits, weight=-10.0),
|
||||
"action_rate_l2": RewardTermCfg(func=mdp.action_rate_l2, weight=-0.05),
|
||||
"foot_gait": RewardTermCfg(
|
||||
func=mdp.feet_gait,
|
||||
weight=0.5,
|
||||
params={
|
||||
"period": 0.6,
|
||||
"offset": [0.0, 0.5],
|
||||
"threshold": 0.56,
|
||||
"command_threshold": 0.1,
|
||||
"command_name": "twist",
|
||||
"sensor_name": "feet_ground_contact",
|
||||
}
|
||||
),
|
||||
"foot_clearance": RewardTermCfg(
|
||||
func=mdp.feet_clearance,
|
||||
weight=-1.0,
|
||||
params={
|
||||
"target_height": 0.10,
|
||||
"command_name": "twist",
|
||||
"command_threshold": 0.1,
|
||||
"asset_cfg": SceneEntityCfg("robot", site_names=()), # Set per-robot.
|
||||
},
|
||||
),
|
||||
"foot_slip": RewardTermCfg(
|
||||
func=mdp.feet_slip,
|
||||
weight=-0.25,
|
||||
params={
|
||||
"sensor_name": "feet_ground_contact",
|
||||
"command_name": "twist",
|
||||
"command_threshold": 0.1,
|
||||
"asset_cfg": SceneEntityCfg("robot", site_names=()), # Set per-robot.
|
||||
},
|
||||
),
|
||||
"soft_landing": RewardTermCfg(
|
||||
func=mdp.soft_landing,
|
||||
weight=-1e-3,
|
||||
params={
|
||||
"sensor_name": "feet_ground_contact",
|
||||
"command_name": "twist",
|
||||
"command_threshold": 0.1,
|
||||
},
|
||||
),
|
||||
"stand_still": RewardTermCfg(
|
||||
func=mdp.stand_still,
|
||||
weight=-1.0,
|
||||
params={
|
||||
"command_name": "twist",
|
||||
"command_threshold": 0.1,
|
||||
"asset_cfg": SceneEntityCfg("robot", joint_names=".*"),
|
||||
},
|
||||
),
|
||||
}
|
||||
|
||||
##
|
||||
# Terminations
|
||||
##
|
||||
|
||||
terminations = {
|
||||
"time_out": TerminationTermCfg(func=mdp.time_out, time_out=True),
|
||||
"fell_over": TerminationTermCfg(
|
||||
func=mdp.bad_orientation,
|
||||
params={"limit_angle": math.radians(70.0)},
|
||||
),
|
||||
}
|
||||
|
||||
##
|
||||
# Curriculum
|
||||
##
|
||||
|
||||
curriculum = {
|
||||
"terrain_levels": CurriculumTermCfg(
|
||||
func=mdp.terrain_levels_vel,
|
||||
params={"command_name": "twist"},
|
||||
),
|
||||
"command_vel": CurriculumTermCfg(
|
||||
func=mdp.commands_vel,
|
||||
params={
|
||||
"command_name": "twist",
|
||||
"velocity_stages": [
|
||||
{"step": 0, "lin_vel_x": (-0.5, 1.0), "lin_vel_y": (-0.5, 0.5), "ang_vel_z": (-1.0, 1.0)},
|
||||
{"step": 5000 * 24, "lin_vel_x": (-1.0, 2.0), "lin_vel_y": (-1.0, 1.0)},
|
||||
],
|
||||
},
|
||||
),
|
||||
}
|
||||
|
||||
##
|
||||
# Assemble and return
|
||||
##
|
||||
|
||||
return ManagerBasedRlEnvCfg(
|
||||
scene=SceneCfg(
|
||||
terrain=TerrainEntityCfg(
|
||||
terrain_type="generator",
|
||||
terrain_generator=replace(ROUGH_TERRAINS_CFG),
|
||||
max_init_terrain_level=5,
|
||||
),
|
||||
sensors=(terrain_scan,),
|
||||
num_envs=1,
|
||||
extent=2.0,
|
||||
),
|
||||
observations=observations,
|
||||
actions=actions,
|
||||
commands=commands,
|
||||
events=events,
|
||||
rewards=rewards,
|
||||
terminations=terminations,
|
||||
curriculum=curriculum,
|
||||
metrics=metrics,
|
||||
viewer=ViewerConfig(
|
||||
origin_type=ViewerConfig.OriginType.ASSET_BODY,
|
||||
entity_name="robot",
|
||||
body_name="", # Set per-robot.
|
||||
distance=3.0,
|
||||
elevation=-5.0,
|
||||
azimuth=90.0,
|
||||
),
|
||||
sim=SimulationCfg(
|
||||
nconmax=35,
|
||||
njmax=1500,
|
||||
mujoco=MujocoCfg(
|
||||
timestep=0.005,
|
||||
iterations=10,
|
||||
ls_iterations=20,
|
||||
),
|
||||
),
|
||||
decimation=4,
|
||||
episode_length_s=20.0,
|
||||
)
|
||||
Reference in New Issue
Block a user