cyy/go2_rl_robotlab
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

init commit.

Browse Source
This commit is contained in:
wertyuilife
2026-03-23 03:17:59 +08:00
commit fce9166e87
99 changed files with 13778 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

35
source/robot_lab/config/extension.toml Normal file
View File

@@ -0,0 +1,35 @@
[package]
# Semantic Versioning is used: https://semver.org/
version = "2.3.0"
# Description
category = "isaaclab"
readme = "README.md"
title = "RL Extension Library for Robots, Based on IsaacLab."
author = "Ziqi Fan"
maintainer = "Ziqi Fan"
description="RL Extension Library for Robots, Based on IsaacLab."
repository = "https://github.com/fan-ziqi/robot_lab.git"
keywords = ["extension", "robot_lab", "isaaclab"]
[dependencies]
"isaaclab" = {}
"isaaclab_assets" = {}
"isaaclab_mimic" = {}
"isaaclab_rl" = {}
"isaaclab_tasks" = {}
# NOTE: Add additional dependencies here
[[python.module]]
name = "robot_lab"
[isaaclab_settings]
# TODO: Uncomment and list any apt dependencies here.
# If none, leave it commented out.
# apt_deps = ["example_package"]
# TODO: Uncomment and provide path to a ros_ws
# with rosdeps to be installed. If none,
# leave it commented out.
# ros_ws = "path/from/extension_root/to/ros_ws"

217
source/robot_lab/data/Robots/unitree/go2_description/meshes/base.dae Normal file
View File

File diff suppressed because one or more lines are too long

130
source/robot_lab/data/Robots/unitree/go2_description/meshes/calf.dae Normal file
View File

File diff suppressed because one or more lines are too long

130
source/robot_lab/data/Robots/unitree/go2_description/meshes/calf_mirror.dae Normal file
View File

File diff suppressed because one or more lines are too long

89
source/robot_lab/data/Robots/unitree/go2_description/meshes/foot.dae Normal file
View File

File diff suppressed because one or more lines are too long

130
source/robot_lab/data/Robots/unitree/go2_description/meshes/hip.dae Normal file
View File

File diff suppressed because one or more lines are too long

130
source/robot_lab/data/Robots/unitree/go2_description/meshes/thigh.dae Normal file
View File

File diff suppressed because one or more lines are too long

130
source/robot_lab/data/Robots/unitree/go2_description/meshes/thigh_mirror.dae Normal file
View File

File diff suppressed because one or more lines are too long

761
source/robot_lab/data/Robots/unitree/go2_description/urdf/go2_description.urdf Normal file
View File

@@ -0,0 +1,761 @@
<?xml version='1.0' encoding='utf-8'?>
<robot name="go2_description">
<link name="base">
<inertial>
<origin xyz="0.021112 0 -0.005366" rpy="0 0 0" />
<mass value="6.921" />
<inertia ixx="0.02448" ixy="0.00012166" ixz="0.0014849" iyy="0.098077" iyz="-3.12E-05" izz="0.107" />
</inertial>
<visual>
<origin xyz="0 0 0" rpy="0 0 0" />
<geometry>
<mesh filename="package://go2_description/meshes/base.dae" />
</geometry>
<material name="深色橡胶_001-effect"><color rgba="0.6717055 0.6924257 0.7742702 1" /></material><material name="白色logo_001-effect"><color rgba="1 1 1 1" /></material><material name="黑色贴纸_007-effect"><color rgba="0.009018 0.009018 0.009018 1" /></material><material name="黑色金属_007-effect"><color rgba="0.01341 0.01341 0.01341 1" /></material><material name="黑色塑料_003-effect"><color rgba="0.023455 0.023455 0.023455 1" /></material></visual>
<collision>
<origin rpy="0 0 0" xyz="0 0 0" />
<geometry>
<box size="0.3762 0.0935 0.114" />
</geometry>
</collision>
</link>
<link name="Head_upper">
<inertial>
<origin xyz="0 0 0" rpy="0 0 0" />
<mass value="0.001" />
<inertia ixx="9.6e-06" ixy="0" ixz="0" iyy="9.6e-06" iyz="0" izz="9.6e-06" />
</inertial>
<collision>
<origin xyz="0 0 0" rpy="0 0 0" />
<geometry>
<cylinder radius="0.05" length="0.09" />
</geometry>
</collision>
</link>
<joint name="Head_upper_joint" type="fixed" dont_collapse="true">
<origin xyz="0.285 0 0.01" rpy="0 0 0" />
<parent link="base" />
<child link="Head_upper" />
<axis xyz="0 0 0" />
</joint>
<link name="Head_lower">
<inertial>
<origin xyz="0 0 0" rpy="0 0 0" />
<mass value="0.001" />
<inertia ixx="9.6e-06" ixy="0" ixz="0" iyy="9.6e-06" iyz="0" izz="9.6e-06" />
</inertial>
<collision>
<origin xyz="0 0 0" rpy="0 0 0" />
<geometry>
<sphere radius="0.047" />
</geometry>
</collision>
</link>
<joint name="Head_lower_joint" type="fixed" dont_collapse="true">
<origin xyz="0.008 0 -0.07" rpy="0 0 0" />
<parent link="Head_upper" />
<child link="Head_lower" />
<axis xyz="0 0 0" />
</joint>
<link name="FL_hip">
<inertial>
<origin xyz="-0.0054 0.00194 -0.000105" rpy="0 0 0" />
<mass value="0.678" />
<inertia ixx="0.00048" ixy="-3.01E-06" ixz="1.11E-06" iyy="0.000884" iyz="-1.42E-06" izz="0.000596" />
</inertial>
<visual>
<origin xyz="0 0 0" rpy="0 0 0" />
<geometry>
<mesh filename="package://go2_description/meshes/hip.dae" />
</geometry>
<material name="深色橡胶_001-effect"><color rgba="0.6717055 0.6924257 0.7742702 1" /></material><material name="Material-effect"><color rgba="0.9 0.95 0.95 1" /></material></visual>
<collision>
<origin rpy="1.5707963267948966 0 0" xyz="0 0.08 0" />
<geometry>
<cylinder length="0.04" radius="0.046" />
</geometry>
</collision>
</link>
<joint name="FL_hip_joint" type="revolute">
<origin xyz="0.1934 0.0465 0" rpy="0 0 0" />
<parent link="base" />
<child link="FL_hip" />
<axis xyz="1 0 0" />
<limit lower="-1.0472" upper="1.0472" effort="23.7" velocity="30.1" />
</joint>
<link name="FL_hip_rotor">
<inertial>
<origin rpy="0 0 0" xyz="0.0 0.0 0.0" />
<mass value="0.089" />
<inertia ixx="0.000111842" ixy="0.0" ixz="0.0" iyy="0.000059647" iyz="0.0" izz="0.000059647" />
</inertial>
</link>
<joint name="FL_hip_rotor_joint" type="fixed">
<origin rpy="0 0 0" xyz="0.11215 0.04675 0" />
<parent link="base" />
<child link="FL_hip_rotor" />
</joint>
<link name="FL_thigh">
<inertial>
<origin xyz="-0.00374 -0.0223 -0.0327" rpy="0 0 0" />
<mass value="1.152" />
<inertia ixx="0.00584" ixy="8.72E-05" ixz="-0.000289" iyy="0.0058" iyz="0.000808" izz="0.00103" />
</inertial>
<visual>
<origin xyz="0 0 0" rpy="0 0 0" />
<geometry>
<mesh filename="package://go2_description/meshes/thigh.dae" />
</geometry>
<material name="深色橡胶_003-effect"><color rgba="0.6717055 0.6924257 0.7742702 1" /></material><material name="Material_009-effect"><color rgba="0.9 0.95 0.95 1" /></material></visual>
<collision>
<origin rpy="0 1.5707963267948966 0" xyz="0 0 -0.1065" />
<geometry>
<box size="0.213 0.0245 0.034" />
</geometry>
</collision>
</link>
<joint name="FL_thigh_joint" type="revolute">
<origin xyz="0 0.0955 0" rpy="0 0 0" />
<parent link="FL_hip" />
<child link="FL_thigh" />
<axis xyz="0 1 0" />
<limit lower="-1.5708" upper="3.4907" effort="23.7" velocity="30.1" />
</joint>
<link name="FL_thigh_rotor">
<inertial>
<origin rpy="0 0 0" xyz="0.0 0.0 0.0" />
<mass value="0.089" />
<inertia ixx="0.000059647" ixy="0.0" ixz="0.0" iyy="0.000111842" iyz="0.0" izz="0.000059647" />
</inertial>
</link>
<joint name="FL_thigh_rotor_joint" type="fixed">
<origin rpy="0 0 0" xyz="0 -0.00015 0" />
<parent link="FL_hip" />
<child link="FL_thigh_rotor" />
</joint>
<link name="FL_calf">
<inertial>
<origin xyz="0.00548 -0.000975 -0.115" rpy="0 0 0" />
<mass value="0.154" />
<inertia ixx="0.00108" ixy="3.4E-07" ixz="1.72E-05" iyy="0.0011" iyz="8.28E-06" izz="3.29E-05" />
</inertial>
<visual>
<origin xyz="0 0 0" rpy="0 0 0" />
<geometry>
<mesh filename="package://go2_description/meshes/calf.dae" />
</geometry>
<material name="黑色足端_005-effect"><color rgba="0 0 0 1" /></material><material name="深色橡胶_005-effect"><color rgba="0.6717055 0.6924257 0.7742702 1" /></material></visual>
<collision>
<origin rpy="0 -0.21 0" xyz="0.008 0 -0.06" />
<geometry>
<cylinder length="0.12" radius="0.012" />
</geometry>
</collision>
</link>
<joint name="FL_calf_joint" type="revolute">
<origin xyz="0 0 -0.213" rpy="0 0 0" />
<parent link="FL_thigh" />
<child link="FL_calf" />
<axis xyz="0 1 0" />
<limit lower="-2.7227" upper="-0.83776" effort="45.43" velocity="15.70" />
</joint>
<link name="FL_calflower">
<collision>
<origin rpy="0 0 0" xyz="0 0 0" />
<geometry>
<cylinder length="0.065" radius="0.011" />
</geometry>
</collision>
</link>
<joint name="FL_calflower_joint" type="fixed">
<origin xyz="0.020 0 -0.148" rpy="0 0.05 0" />
<parent link="FL_calf" />
<child link="FL_calflower" />
<axis xyz="0 0 0" />
</joint>
<link name="FL_calflower1">
<collision>
<origin rpy="0 0 0" xyz="0 0 0" />
<geometry>
<cylinder length="0.03" radius="0.0155" />
</geometry>
</collision>
</link>
<joint name="FL_calflower1_joint" type="fixed">
<origin xyz="-0.01 0 -0.04" rpy="0 0.48 0" />
<parent link="FL_calflower" />
<child link="FL_calflower1" />
<axis xyz="0 0 0" />
</joint>
<link name="FL_calf_rotor">
<inertial>
<origin rpy="0 0 0" xyz="0.0 0.0 0.0" />
<mass value="0.089" />
<inertia ixx="0.000059647" ixy="0.0" ixz="0.0" iyy="0.000111842" iyz="0.0" izz="0.000059647" />
</inertial>
</link>
<joint name="FL_calf_rotor_joint" type="fixed">
<origin rpy="0 0 0" xyz="0 -0.03235 0" />
<parent link="FL_thigh" />
<child link="FL_calf_rotor" />
</joint>
<link name="FL_foot">
<inertial>
<origin xyz="0 0 0" rpy="0 0 0" />
<mass value="0.04" />
<inertia ixx="9.6e-06" ixy="0" ixz="0" iyy="9.6e-06" iyz="0" izz="9.6e-06" />
</inertial>
<visual>
<origin xyz="0 0 0" rpy="0 0 0" />
<geometry>
<mesh filename="package://go2_description/meshes/foot.dae" />
</geometry>
<material name="黑色足端_007-effect"><color rgba="0 0 0 1" /></material></visual>
<collision>
<origin rpy="0 0 0" xyz="-0.002 0 0" />
<geometry>
<sphere radius="0.022" />
</geometry>
</collision>
</link>
<joint name="FL_foot_joint" type="fixed" dont_collapse="true">
<origin xyz="0 0 -0.213" rpy="0 0 0" />
<parent link="FL_calf" />
<child link="FL_foot" />
<axis xyz="0 0 0" />
</joint>
<link name="FR_hip">
<inertial>
<origin xyz="-0.0054 -0.00194 -0.000105" rpy="0 0 0" />
<mass value="0.678" />
<inertia ixx="0.00048" ixy="3.01E-06" ixz="1.11E-06" iyy="0.000884" iyz="1.42E-06" izz="0.000596" />
</inertial>
<visual>
<origin xyz="0 0 0" rpy="3.1415 0 0" />
<geometry>
<mesh filename="package://go2_description/meshes/hip.dae" />
</geometry>
<material name="深色橡胶_001-effect"><color rgba="0.6717055 0.6924257 0.7742702 1" /></material><material name="Material-effect"><color rgba="0.9 0.95 0.95 1" /></material></visual>
<collision>
<origin rpy="1.5707963267948966 0 0" xyz="0 -0.08 0" />
<geometry>
<cylinder length="0.04" radius="0.046" />
</geometry>
</collision>
</link>
<joint name="FR_hip_joint" type="revolute">
<origin xyz="0.1934 -0.0465 0" rpy="0 0 0" />
<parent link="base" />
<child link="FR_hip" />
<axis xyz="1 0 0" />
<limit lower="-1.0472" upper="1.0472" effort="23.7" velocity="30.1" />
</joint>
<link name="FR_hip_rotor">
<inertial>
<origin rpy="0 0 0" xyz="0.0 0.0 0.0" />
<mass value="0.089" />
<inertia ixx="0.000111842" ixy="0.0" ixz="0.0" iyy="0.000059647" iyz="0.0" izz="0.000059647" />
</inertial>
</link>
<joint name="FR_hip_rotor_joint" type="fixed">
<origin rpy="0 0 0" xyz="0.11215 -0.04675 0" />
<parent link="base" />
<child link="FR_hip_rotor" />
</joint>
<link name="FR_thigh">
<inertial>
<origin xyz="-0.00374 0.0223 -0.0327" rpy="0 0 0" />
<mass value="1.152" />
<inertia ixx="0.00584" ixy="-8.72E-05" ixz="-0.000289" iyy="0.0058" iyz="-0.000808" izz="0.00103" />
</inertial>
<visual>
<origin xyz="0 0 0" rpy="0 0 0" />
<geometry>
<mesh filename="package://go2_description/meshes/thigh_mirror.dae" />
</geometry>
<material name="深色橡胶_012-effect"><color rgba="0.6717055 0.6924257 0.7742702 1" /></material><material name="Material_011-effect"><color rgba="0.9 0.95 0.95 1" /></material></visual>
<collision>
<origin rpy="0 1.5707963267948966 0" xyz="0 0 -0.1065" />
<geometry>
<box size="0.213 0.0245 0.034" />
</geometry>
</collision>
</link>
<joint name="FR_thigh_joint" type="revolute">
<origin xyz="0 -0.0955 0" rpy="0 0 0" />
<parent link="FR_hip" />
<child link="FR_thigh" />
<axis xyz="0 1 0" />
<limit lower="-1.5708" upper="3.4907" effort="23.7" velocity="30.1" />
</joint>
<link name="FR_thigh_rotor">
<inertial>
<origin rpy="0 0 0" xyz="0.0 0.0 0.0" />
<mass value="0.089" />
<inertia ixx="0.000059647" ixy="0.0" ixz="0.0" iyy="0.000111842" iyz="0.0" izz="0.000059647" />
</inertial>
</link>
<joint name="FR_thigh_rotor_joint" type="fixed">
<origin rpy="0 0 0" xyz="0 0.00015 0" />
<parent link="FR_hip" />
<child link="FR_thigh_rotor" />
</joint>
<link name="FR_calf">
<inertial>
<origin xyz="0.00548 0.000975 -0.115" rpy="0 0 0" />
<mass value="0.154" />
<inertia ixx="0.00108" ixy="-3.4E-07" ixz="1.72E-05" iyy="0.0011" iyz="-8.28E-06" izz="3.29E-05" />
</inertial>
<visual>
<origin xyz="0 0 0" rpy="0 0 0" />
<geometry>
<mesh filename="package://go2_description/meshes/calf_mirror.dae" />
</geometry>
<material name="黑色足端_006-effect"><color rgba="0 0 0 1" /></material><material name="深色橡胶_007-effect"><color rgba="0.6717055 0.6924257 0.7742702 1" /></material></visual>
<collision>
<origin rpy="0 -0.2 0" xyz="0.01 0 -0.06" />
<geometry>
<cylinder length="0.12" radius="0.013" />
</geometry>
</collision>
</link>
<joint name="FR_calf_joint" type="revolute">
<origin xyz="0 0 -0.213" rpy="0 0 0" />
<parent link="FR_thigh" />
<child link="FR_calf" />
<axis xyz="0 1 0" />
<limit lower="-2.7227" upper="-0.83776" effort="45.43" velocity="15.70" />
</joint>
<link name="FR_calflower">
<collision>
<origin rpy="0 0 0" xyz="0 0 0" />
<geometry>
<cylinder length="0.065" radius="0.011" />
</geometry>
</collision>
</link>
<joint name="FR_calflower_joint" type="fixed">
<origin xyz="0.020 0 -0.148" rpy="0 0.05 0" />
<parent link="FR_calf" />
<child link="FR_calflower" />
<axis xyz="0 0 0" />
</joint>
<link name="FR_calflower1">
<collision>
<origin rpy="0 0 0" xyz="0 0 0" />
<geometry>
<cylinder length="0.03" radius="0.0155" />
</geometry>
</collision>
</link>
<joint name="FR_calflower1_joint" type="fixed">
<origin xyz="-0.01 0 -0.04" rpy="0 0.48 0" />
<parent link="FR_calflower" />
<child link="FR_calflower1" />
<axis xyz="0 0 0" />
</joint>
<link name="FR_calf_rotor">
<inertial>
<origin rpy="0 0 0" xyz="0.0 0.0 0.0" />
<mass value="0.089" />
<inertia ixx="0.000059647" ixy="0.0" ixz="0.0" iyy="0.000111842" iyz="0.0" izz="0.000059647" />
</inertial>
</link>
<joint name="FR_calf_rotor_joint" type="fixed">
<origin rpy="0 0 0" xyz="0 0.03235 0" />
<parent link="FR_thigh" />
<child link="FR_calf_rotor" />
</joint>
<link name="FR_foot">
<inertial>
<origin xyz="0 0 0" rpy="0 0 0" />
<mass value="0.04" />
<inertia ixx="9.6e-06" ixy="0" ixz="0" iyy="9.6e-06" iyz="0" izz="9.6e-06" />
</inertial>
<visual>
<origin xyz="0 0 0" rpy="0 0 0" />
<geometry>
<mesh filename="package://go2_description/meshes/foot.dae" />
</geometry>
<material name="黑色足端_007-effect"><color rgba="0 0 0 1" /></material></visual>
<collision>
<origin rpy="0 0 0" xyz="-0.002 0 0" />
<geometry>
<sphere radius="0.022" />
</geometry>
</collision>
</link>
<joint name="FR_foot_joint" type="fixed" dont_collapse="true">
<origin xyz="0 0 -0.213" rpy="0 0 0" />
<parent link="FR_calf" />
<child link="FR_foot" />
<axis xyz="0 0 0" />
</joint>
<link name="RL_hip">
<inertial>
<origin xyz="0.0054 0.00194 -0.000105" rpy="0 0 0" />
<mass value="0.678" />
<inertia ixx="0.00048" ixy="3.01E-06" ixz="-1.11E-06" iyy="0.000884" iyz="-1.42E-06" izz="0.000596" />
</inertial>
<visual>
<origin xyz="0 0 0" rpy="0 3.1415 0" />
<geometry>
<mesh filename="package://go2_description/meshes/hip.dae" />
</geometry>
<material name="深色橡胶_001-effect"><color rgba="0.6717055 0.6924257 0.7742702 1" /></material><material name="Material-effect"><color rgba="0.9 0.95 0.95 1" /></material></visual>
<collision>
<origin rpy="1.5707963267948966 0 0" xyz="0 0.08 0" />
<geometry>
<cylinder length="0.04" radius="0.046" />
</geometry>
</collision>
</link>
<joint name="RL_hip_joint" type="revolute">
<origin xyz="-0.1934 0.0465 0" rpy="0 0 0" />
<parent link="base" />
<child link="RL_hip" />
<axis xyz="1 0 0" />
<limit lower="-1.0472" upper="1.0472" effort="23.7" velocity="30.1" />
</joint>
<link name="RL_hip_rotor">
<inertial>
<origin rpy="0 0 0" xyz="0.0 0.0 0.0" />
<mass value="0.089" />
<inertia ixx="0.000111842" ixy="0.0" ixz="0.0" iyy="0.000059647" iyz="0.0" izz="0.000059647" />
</inertial>
</link>
<joint name="RL_hip_rotor_joint" type="fixed">
<origin rpy="0 0 0" xyz="-0.11215 0.04675 0" />
<parent link="base" />
<child link="RL_hip_rotor" />
</joint>
<link name="RL_thigh">
<inertial>
<origin xyz="-0.00374 -0.0223 -0.0327" rpy="0 0 0" />
<mass value="1.152" />
<inertia ixx="0.00584" ixy="8.72E-05" ixz="-0.000289" iyy="0.0058" iyz="0.000808" izz="0.00103" />
</inertial>
<visual>
<origin xyz="0 0 0" rpy="0 0 0" />
<geometry>
<mesh filename="package://go2_description/meshes/thigh.dae" />
</geometry>
<material name="深色橡胶_003-effect"><color rgba="0.6717055 0.6924257 0.7742702 1" /></material><material name="Material_009-effect"><color rgba="0.9 0.95 0.95 1" /></material></visual>
<collision>
<origin rpy="0 1.5707963267948966 0" xyz="0 0 -0.1065" />
<geometry>
<box size="0.213 0.0245 0.034" />
</geometry>
</collision>
</link>
<joint name="RL_thigh_joint" type="revolute">
<origin xyz="0 0.0955 0" rpy="0 0 0" />
<parent link="RL_hip" />
<child link="RL_thigh" />
<axis xyz="0 1 0" />
<limit lower="-0.5236" upper="4.5379" effort="23.7" velocity="30.1" />
</joint>
<link name="RL_thigh_rotor">
<inertial>
<origin rpy="0 0 0" xyz="0.0 0.0 0.0" />
<mass value="0.089" />
<inertia ixx="0.000059647" ixy="0.0" ixz="0.0" iyy="0.000111842" iyz="0.0" izz="0.000059647" />
</inertial>
</link>
<joint name="RL_thigh_rotor_joint" type="fixed">
<origin rpy="0 0 0" xyz="0 -0.00015 0" />
<parent link="RL_hip" />
<child link="RL_thigh_rotor" />
</joint>
<link name="RL_calf">
<inertial>
<origin xyz="0.00548 -0.000975 -0.115" rpy="0 0 0" />
<mass value="0.154" />
<inertia ixx="0.00108" ixy="3.4E-07" ixz="1.72E-05" iyy="0.0011" iyz="8.28E-06" izz="3.29E-05" />
</inertial>
<visual>
<origin xyz="0 0 0" rpy="0 0 0" />
<geometry>
<mesh filename="package://go2_description/meshes/calf.dae" />
</geometry>
<material name="黑色足端_005-effect"><color rgba="0 0 0 1" /></material><material name="深色橡胶_005-effect"><color rgba="0.6717055 0.6924257 0.7742702 1" /></material></visual>
<collision>
<origin rpy="0 -0.2 0" xyz="0.01 0 -0.06" />
<geometry>
<cylinder length="0.12" radius="0.013" />
</geometry>
</collision>
</link>
<joint name="RL_calf_joint" type="revolute">
<origin xyz="0 0 -0.213" rpy="0 0 0" />
<parent link="RL_thigh" />
<child link="RL_calf" />
<axis xyz="0 1 0" />
<limit lower="-2.7227" upper="-0.83776" effort="45.43" velocity="15.70" />
</joint>
<link name="RL_calflower">
<collision>
<origin rpy="0 0 0" xyz="0 0 0" />
<geometry>
<cylinder length="0.065" radius="0.011" />
</geometry>
</collision>
</link>
<joint name="RL_calflower_joint" type="fixed">
<origin xyz="0.020 0 -0.148" rpy="0 0.05 0" />
<parent link="RL_calf" />
<child link="RL_calflower" />
<axis xyz="0 0 0" />
</joint>
<link name="RL_calflower1">
<collision>
<origin rpy="0 0 0" xyz="0 0 0" />
<geometry>
<cylinder length="0.03" radius="0.0155" />
</geometry>
</collision>
</link>
<joint name="RL_calflower1_joint" type="fixed">
<origin xyz="-0.01 0 -0.04" rpy="0 0.48 0" />
<parent link="RL_calflower" />
<child link="RL_calflower1" />
<axis xyz="0 0 0" />
</joint>
<link name="RL_calf_rotor">
<inertial>
<origin rpy="0 0 0" xyz="0.0 0.0 0.0" />
<mass value="0.089" />
<inertia ixx="0.000059647" ixy="0.0" ixz="0.0" iyy="0.000111842" iyz="0.0" izz="0.000059647" />
</inertial>
</link>
<joint name="RL_calf_rotor_joint" type="fixed">
<origin rpy="0 0 0" xyz="0 -0.03235 0" />
<parent link="RL_thigh" />
<child link="RL_calf_rotor" />
</joint>
<link name="RL_foot">
<inertial>
<origin xyz="0 0 0" rpy="0 0 0" />
<mass value="0.04" />
<inertia ixx="9.6e-06" ixy="0" ixz="0" iyy="9.6e-06" iyz="0" izz="9.6e-06" />
</inertial>
<visual>
<origin xyz="0 0 0" rpy="0 0 0" />
<geometry>
<mesh filename="package://go2_description/meshes/foot.dae" />
</geometry>
<material name="黑色足端_007-effect"><color rgba="0 0 0 1" /></material></visual>
<collision>
<origin rpy="0 0 0" xyz="-0.002 0 0" />
<geometry>
<sphere radius="0.022" />
</geometry>
</collision>
</link>
<joint name="RL_foot_joint" type="fixed" dont_collapse="true">
<origin xyz="0 0 -0.213" rpy="0 0 0" />
<parent link="RL_calf" />
<child link="RL_foot" />
<axis xyz="0 0 0" />
</joint>
<link name="RR_hip">
<inertial>
<origin xyz="0.0054 -0.00194 -0.000105" rpy="0 0 0" />
<mass value="0.678" />
<inertia ixx="0.00048" ixy="-3.01E-06" ixz="-1.11E-06" iyy="0.000884" iyz="1.42E-06" izz="0.000596" />
</inertial>
<visual>
<origin xyz="0 0 0" rpy="3.1415 3.1415 0" />
<geometry>
<mesh filename="package://go2_description/meshes/hip.dae" />
</geometry>
<material name="深色橡胶_001-effect"><color rgba="0.6717055 0.6924257 0.7742702 1" /></material><material name="Material-effect"><color rgba="0.9 0.95 0.95 1" /></material></visual>
<collision>
<origin rpy="1.5707963267948966 0 0" xyz="0 -0.08 0" />
<geometry>
<cylinder length="0.04" radius="0.046" />
</geometry>
</collision>
</link>
<joint name="RR_hip_joint" type="revolute">
<origin xyz="-0.1934 -0.0465 0" rpy="0 0 0" />
<parent link="base" />
<child link="RR_hip" />
<axis xyz="1 0 0" />
<limit lower="-1.0472" upper="1.0472" effort="23.7" velocity="30.1" />
</joint>
<link name="RR_hip_rotor">
<inertial>
<origin rpy="0 0 0" xyz="0.0 0.0 0.0" />
<mass value="0.089" />
<inertia ixx="0.000111842" ixy="0.0" ixz="0.0" iyy="0.000059647" iyz="0.0" izz="0.000059647" />
</inertial>
</link>
<joint name="RR_hip_rotor_joint" type="fixed">
<origin rpy="0 0 0" xyz="-0.11215 -0.04675 0" />
<parent link="base" />
<child link="RR_hip_rotor" />
</joint>
<link name="RR_thigh">
<inertial>
<origin xyz="-0.00374 0.0223 -0.0327" rpy="0 0 0" />
<mass value="1.152" />
<inertia ixx="0.00584" ixy="-8.72E-05" ixz="-0.000289" iyy="0.0058" iyz="-0.000808" izz="0.00103" />
</inertial>
<visual>
<origin xyz="0 0 0" rpy="0 0 0" />
<geometry>
<mesh filename="package://go2_description/meshes/thigh_mirror.dae" />
</geometry>
<material name="深色橡胶_012-effect"><color rgba="0.6717055 0.6924257 0.7742702 1" /></material><material name="Material_011-effect"><color rgba="0.9 0.95 0.95 1" /></material></visual>
<collision>
<origin rpy="0 1.5707963267948966 0" xyz="0 0 -0.1065" />
<geometry>
<box size="0.213 0.0245 0.034" />
</geometry>
</collision>
</link>
<joint name="RR_thigh_joint" type="revolute">
<origin xyz="0 -0.0955 0" rpy="0 0 0" />
<parent link="RR_hip" />
<child link="RR_thigh" />
<axis xyz="0 1 0" />
<limit lower="-0.5236" upper="4.5379" effort="23.7" velocity="30.1" />
</joint>
<link name="RR_thigh_rotor">
<inertial>
<origin rpy="0 0 0" xyz="0.0 0.0 0.0" />
<mass value="0.089" />
<inertia ixx="0.000059647" ixy="0.0" ixz="0.0" iyy="0.000111842" iyz="0.0" izz="0.000059647" />
</inertial>
</link>
<joint name="RR_thigh_rotor_joint" type="fixed">
<origin rpy="0 0 0" xyz="0 0.00015 0" />
<parent link="RR_hip" />
<child link="RR_thigh_rotor" />
</joint>
<link name="RR_calf">
<inertial>
<origin xyz="0.00548 0.000975 -0.115" rpy="0 0 0" />
<mass value="0.154" />
<inertia ixx="0.00108" ixy="-3.4E-07" ixz="1.72E-05" iyy="0.0011" iyz="-8.28E-06" izz="3.29E-05" />
</inertial>
<visual>
<origin xyz="0 0 0" rpy="0 0 0" />
<geometry>
<mesh filename="package://go2_description/meshes/calf_mirror.dae" />
</geometry>
<material name="黑色足端_006-effect"><color rgba="0 0 0 1" /></material><material name="深色橡胶_007-effect"><color rgba="0.6717055 0.6924257 0.7742702 1" /></material></visual>
<collision>
<origin rpy="0 -0.2 0" xyz="0.01 0 -0.06" />
<geometry>
<cylinder length="0.12" radius="0.013" />
</geometry>
</collision>
</link>
<joint name="RR_calf_joint" type="revolute">
<origin xyz="0 0 -0.213" rpy="0 0 0" />
<parent link="RR_thigh" />
<child link="RR_calf" />
<axis xyz="0 1 0" />
<limit lower="-2.7227" upper="-0.83776" effort="45.43" velocity="15.70" />
</joint>
<link name="RR_calflower">
<collision>
<origin rpy="0 0 0" xyz="0 0 0" />
<geometry>
<cylinder length="0.065" radius="0.011" />
</geometry>
</collision>
</link>
<joint name="RR_calflower_joint" type="fixed">
<origin xyz="0.020 0 -0.148" rpy="0 0.05 0" />
<parent link="RR_calf" />
<child link="RR_calflower" />
<axis xyz="0 0 0" />
</joint>
<link name="RR_calflower1">
<collision>
<origin rpy="0 0 0" xyz="0 0 0" />
<geometry>
<cylinder length="0.03" radius="0.0155" />
</geometry>
</collision>
</link>
<joint name="RR_calflower1_joint" type="fixed">
<origin xyz="-0.01 0 -0.04" rpy="0 0.48 0" />
<parent link="RR_calflower" />
<child link="RR_calflower1" />
<axis xyz="0 0 0" />
</joint>
<link name="RR_calf_rotor">
<inertial>
<origin rpy="0 0 0" xyz="0.0 0.0 0.0" />
<mass value="0.089" />
<inertia ixx="0.000059647" ixy="0.0" ixz="0.0" iyy="0.000111842" iyz="0.0" izz="0.000059647" />
</inertial>
</link>
<joint name="RR_calf_rotor_joint" type="fixed">
<origin rpy="0 0 0" xyz="0 0.03235 0" />
<parent link="RR_thigh" />
<child link="RR_calf_rotor" />
</joint>
<link name="RR_foot">
<inertial>
<origin xyz="0 0 0" rpy="0 0 0" />
<mass value="0.04" />
<inertia ixx="9.6e-06" ixy="0" ixz="0" iyy="9.6e-06" iyz="0" izz="9.6e-06" />
</inertial>
<visual>
<origin xyz="0 0 0" rpy="0 0 0" />
<geometry>
<mesh filename="package://go2_description/meshes/foot.dae" />
</geometry>
<material name="黑色足端_007-effect"><color rgba="0 0 0 1" /></material></visual>
<collision>
<origin rpy="0 0 0" xyz="-0.002 0 0" />
<geometry>
<sphere radius="0.022" />
</geometry>
</collision>
</link>
<joint name="RR_foot_joint" type="fixed" dont_collapse="true">
<origin xyz="0 0 -0.213" rpy="0 0 0" />
<parent link="RR_calf" />
<child link="RR_foot" />
<axis xyz="0 0 0" />
</joint>
<link name="imu">
<inertial>
<origin xyz="0 0 0" rpy="0 0 0" />
<mass value="0" />
<inertia ixx="0" ixy="0" ixz="0" iyy="0" iyz="0" izz="0" />
</inertial>
</link>
<joint name="imu_joint" type="fixed">
<origin xyz="-0.02557 0 0.04232" rpy="0 0 0" />
<parent link="base" />
<child link="imu" />
<axis xyz="0 0 0" />
</joint>
<link name="radar">
<inertial>
<origin xyz="0 0 0" rpy="0 0 0" />
<mass value="0" />
<inertia ixx="0" ixy="0" ixz="0" iyy="0" iyz="0" izz="0" />
</inertial>
</link>
<joint name="radar_joint" type="fixed">
<origin xyz="0.28945 0 -0.046825" rpy="0 2.8782 0" />
<parent link="base" />
<child link="radar" />
<axis xyz="0 0 0" />
</joint>
<joint name="front_camera_joint" type="fixed">
<origin rpy="0 0 0" xyz="0.32715 -0.00003 0.04297" />
<parent link="base" />
<child link="front_camera" />
</joint>
<link name="front_camera">
</link>
</robot>

3
source/robot_lab/pyproject.toml Normal file
View File

@@ -0,0 +1,3 @@
[build-system]
requires = ["setuptools", "wheel", "toml"]
build-backend = "setuptools.build_meta"

12
source/robot_lab/robot_lab/__init__.py Normal file
View File

@@ -0,0 +1,12 @@
# Copyright (c) 2024-2025 Ziqi Fan
# SPDX-License-Identifier: Apache-2.0
"""
Python module serving as a project/extension template.
"""
# Register Gym environments.
# from .tasks import *
# Register UI extensions.
from .ui_extension_example import *

29
source/robot_lab/robot_lab/assets/__init__.py Normal file
View File

@@ -0,0 +1,29 @@
# Copyright (c) 2024-2025 Ziqi Fan
# SPDX-License-Identifier: Apache-2.0
# Copyright (c) 2022-2025, The Isaac Lab Project Developers.
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause
"""Package containing asset and sensor configurations."""
import os
import toml
##
# Configuration for different assets.
##
# Conveniences to other module directories via relative paths
ISAACLAB_ASSETS_EXT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../"))
"""Path to the extension source directory."""
ISAACLAB_ASSETS_DATA_DIR = os.path.join(ISAACLAB_ASSETS_EXT_DIR, "data")
"""Path to the extension data directory."""
ISAACLAB_ASSETS_METADATA = toml.load(os.path.join(ISAACLAB_ASSETS_EXT_DIR, "config", "extension.toml"))
"""Extension metadata dictionary parsed from the extension.toml file."""
# Configure the module-level variables
__version__ = ISAACLAB_ASSETS_METADATA["package"]["version"]

70
source/robot_lab/robot_lab/assets/unitree.py Normal file
View File

@@ -0,0 +1,70 @@
# Copyright (c) 2024-2025 Ziqi Fan
# SPDX-License-Identifier: Apache-2.0
"""Configuration for Unitree robots.
Reference: https://github.com/unitreerobotics/unitree_ros
"""
import isaaclab.sim as sim_utils
from isaaclab.actuators import DCMotorCfg
from isaaclab.assets.articulation import ArticulationCfg
from robot_lab.assets import ISAACLAB_ASSETS_DATA_DIR
##
# Configuration
##
UNITREE_GO2_CFG = ArticulationCfg(
spawn=sim_utils.UrdfFileCfg(
fix_base=False,
merge_fixed_joints=True,
replace_cylinders_with_capsules=False,
asset_path=f"{ISAACLAB_ASSETS_DATA_DIR}/Robots/unitree/go2_description/urdf/go2_description.urdf",
activate_contact_sensors=True,
rigid_props=sim_utils.RigidBodyPropertiesCfg(
disable_gravity=False,
retain_accelerations=False,
linear_damping=0.0,
angular_damping=0.0,
max_linear_velocity=1000.0,
max_angular_velocity=1000.0,
max_depenetration_velocity=1.0,
),
articulation_props=sim_utils.ArticulationRootPropertiesCfg(
enabled_self_collisions=False,
solver_position_iteration_count=4,
solver_velocity_iteration_count=0,
),
joint_drive=sim_utils.UrdfConverterCfg.JointDriveCfg(
gains=sim_utils.UrdfConverterCfg.JointDriveCfg.PDGainsCfg(
stiffness=0, damping=0
)
),
),
init_state=ArticulationCfg.InitialStateCfg(
pos=(0.0, 0.0, 0.38),
joint_pos={
".*L_hip_joint": 0.0,
".*R_hip_joint": -0.0,
"F.*_thigh_joint": 0.8,
"R.*_thigh_joint": 0.8,
".*_calf_joint": -1.5,
},
joint_vel={".*": 0.0},
),
soft_joint_pos_limit_factor=0.9,
actuators={
"legs": DCMotorCfg(
joint_names_expr=[".*"],
effort_limit=23.5,
saturation_effort=23.5,
velocity_limit=30.0,
stiffness=25.0,
damping=0.5,
friction=0.0,
),
},
)
"""Configuration of Unitree Go2 using DC motor.
"""

24
source/robot_lab/robot_lab/tasks/__init__.py Normal file
View File

@@ -0,0 +1,24 @@
# Copyright (c) 2024-2025 Ziqi Fan
# SPDX-License-Identifier: Apache-2.0
# Copyright (c) 2022-2025, The Isaac Lab Project Developers.
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause
"""Package containing task implementations for various robotic environments."""
import os
import toml
from isaaclab_tasks.utils import import_packages
##
# Register Gym environments.
##
# The blacklist is used to prevent importing configs from sub-packages
_BLACKLIST_PKGS = ["utils"]
# Import all configs in this package
import_packages(__name__, _BLACKLIST_PKGS)

34
source/robot_lab/robot_lab/tasks/go2/__init__.py Normal file
View File

@@ -0,0 +1,34 @@
# Copyright (c) 2024-2025 Ziqi Fan
# SPDX-License-Identifier: Apache-2.0
# Copyright (c) 2022-2025, The Isaac Lab Project Developers.
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause
"""Package containing task implementations for various robotic environments."""
import os
import toml
import gymnasium as gym
from isaaclab_tasks.utils import import_packages
##
# Register Gym environments.
##
gym.register(
id="Robotlab-Go2-v0",
entry_point="robot_lab.tasks.go2.env.go2_env:ActionDelayGo2Env",
disable_env_checker=True,
kwargs={
"env_cfg_entry_point": f"{__name__}.env_cfg:Go2EnvCfg",
# "env_cfg_entry_point": f"{__name__}.origin_env_cfg:UnitreeGo2RoughEnvCfg",
"rsl_rl_cfg_entry_point": f"{__name__}.rsl_rl_cfg:MoECTSRunnerCfg",
},
)
# The blacklist is used to prevent importing configs from sub-packages
_BLACKLIST_PKGS = ["utils"]
# Import all configs in this package
import_packages(__name__, _BLACKLIST_PKGS)

128
source/robot_lab/robot_lab/tasks/go2/env/go2_env.py vendored Normal file
View File

@@ -0,0 +1,128 @@
from isaaclab.envs import ManagerBasedRLEnv, ManagerBasedRLEnvCfg, VecEnvStepReturn
from robot_lab.tasks.go2.manager.action_manager import ActionManagerWithDelay
import torch
from isaaclab.ui.widgets import ManagerLiveVisualizer
class ActionDelayGo2Env(ManagerBasedRLEnv):
cfg: ManagerBasedRLEnvCfg
def __init__(self, cfg: ManagerBasedRLEnvCfg, render_mode: str | None = None, **kwargs):
"""
Initialize the ActionDelayGo2Env with the given configuration.
Args:
cfg: The configuration for the environment.
render_mode: Rendering mode for the environment, e.g., "human" or "rgb_array". Default is None.
**kwargs: Additional keyword arguments for customization.
"""
# Call the parent class initializer
super().__init__(cfg=cfg, render_mode=render_mode, **kwargs)
print(
"[WARNING] You are using ActionDelayGo2Env; "
"make sure all ActionTerms support multiple calls to process_actions() "
"within a single step()."
)
def load_managers(self):
super().load_managers()
# override action manager
self.action_manager = ActionManagerWithDelay(self.cfg.actions, self)
print("[INFO] Overriding action manager with ActionManagerWithDelay: ", self.action_manager)
def step(self, action: torch.Tensor) -> VecEnvStepReturn:
"""Execute one time-step of the environment's dynamics and reset terminated environments.
Args:
action: The actions to apply on the environment. Shape is (num_envs, action_dim).
Returns:
A tuple containing the observations, rewards, resets (terminated and truncated) and extras.
IMPORTANT NOTE:
We intentionally call action_manager.process_action_with_delay() multiple times inside step
(contrary to its original intent) to implement action-delay functionality.
So, we assume that all ActionTerm.process_actions() are allowed to be
called multiple times within a single step call.
"""
# call update_action once per step to set action and prev action
self.action_manager.update_action(action.to(self.device))
# randomly determine when to start applying actions within the decimation steps for each environment
actions_start_decimation = torch.randint(0, self.cfg.decimation+1, (self.num_envs, 1), device=self.device)
self.recorder_manager.record_pre_step()
# check if we need to do rendering within the physics loop
# note: checked here once to avoid multiple checks within the loop
is_rendering = self.sim.has_gui() or self.sim.has_rtx_sensors()
# perform physics stepping
for i in range(self.cfg.decimation):
self._sim_step_counter += 1
# determine which envs should apply delayed action at this decimation step
action_delay_masks = (i <= actions_start_decimation)
self.action_manager.process_action_with_delay(action_delay_masks)
# set actions into buffers
self.action_manager.apply_action()
# set actions into simulator
self.scene.write_data_to_sim()
# simulate
self.sim.step(render=False)
# render between steps only if the GUI or an RTX sensor needs it
# note: we assume the render interval to be the shortest accepted rendering interval.
# If a camera needs rendering at a faster frequency, this will lead to unexpected behavior.
if self._sim_step_counter % self.cfg.sim.render_interval == 0 and is_rendering:
self.sim.render()
# update buffers at sim dt
self.scene.update(dt=self.physics_dt)
# post-step:
# -- update env counters (used for curriculum generation)
self.episode_length_buf += 1 # step in current episode (per env)
self.common_step_counter += 1 # total step (common for all envs)
# -- check terminations
self.reset_buf = self.termination_manager.compute()
self.reset_terminated = self.termination_manager.terminated
self.reset_time_outs = self.termination_manager.time_outs
# -- reward computation
self.reward_buf = self.reward_manager.compute(dt=self.step_dt)
if len(self.recorder_manager.active_terms) > 0:
# update observations for recording if needed
self.obs_buf = self.observation_manager.compute()
self.recorder_manager.record_post_step()
# -- reset envs that terminated/timed-out and log the episode information
reset_env_ids = self.reset_buf.nonzero(as_tuple=False).squeeze(-1)
if len(reset_env_ids) > 0:
# trigger recorder terms for pre-reset calls
self.recorder_manager.record_pre_reset(reset_env_ids)
self._reset_idx(reset_env_ids)
# update articulation kinematics
self.scene.write_data_to_sim()
self.sim.forward()
# if sensors are added to the scene, make sure we render to reflect changes in reset
if self.sim.has_rtx_sensors() and self.cfg.rerender_on_reset:
self.sim.render()
# trigger recorder terms for post-reset calls
self.recorder_manager.record_post_reset(reset_env_ids)
# -- update command
self.command_manager.compute(dt=self.step_dt)
# -- step interval events
if "interval" in self.event_manager.available_modes:
self.event_manager.apply(mode="interval", dt=self.step_dt)
# -- compute observations
# note: done after reset to get the correct observations for reset envs
self.obs_buf = self.observation_manager.compute(update_history=True)
# return observations, rewards, resets and extras
return self.obs_buf, self.reward_buf, self.reset_terminated, self.reset_time_outs, self.extras

615
source/robot_lab/robot_lab/tasks/go2/env_cfg.py Normal file
View File

@@ -0,0 +1,615 @@
import math
import isaaclab.sim as sim_utils
from isaaclab.assets import ArticulationCfg, AssetBaseCfg
from isaaclab.envs import ManagerBasedRLEnvCfg
from isaaclab.managers import CurriculumTermCfg as CurrTerm
from isaaclab.managers import EventTermCfg as EventTerm
from isaaclab.managers import ObservationGroupCfg as ObsGroup
from isaaclab.managers import ObservationTermCfg as ObsTerm
from isaaclab.managers import RewardTermCfg as RewTerm
from isaaclab.managers import SceneEntityCfg
from isaaclab.managers import TerminationTermCfg as DoneTerm
from isaaclab.scene import InteractiveSceneCfg
from isaaclab.sensors import ContactSensorCfg, RayCasterCfg, patterns
from isaaclab.terrains import TerrainImporterCfg
from isaaclab.utils import configclass
from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR, ISAACLAB_NUCLEUS_DIR
from isaaclab.utils.noise import AdditiveUniformNoiseCfg as Unoise
import isaaclab.terrains as terrain_gen
import robot_lab.tasks.go2.mdp as mdp
from robot_lab.assets.unitree import UNITREE_GO2_CFG
JOINT_NAMES = [
"FL_hip_joint", "FL_thigh_joint", "FL_calf_joint",
"FR_hip_joint", "FR_thigh_joint", "FR_calf_joint",
"RL_hip_joint", "RL_thigh_joint", "RL_calf_joint",
"RR_hip_joint", "RR_thigh_joint", "RR_calf_joint",
]
BASE_LINK_NAME = "base"
FOOT_LINK_NAME = ".*_foot"
##
# Terrain definition
##
TERRAIN_CFG = terrain_gen.TerrainGeneratorCfg(
size=(8.0, 8.0),
border_width=20.0,
num_rows=10,
num_cols=20,
horizontal_scale=0.1,
vertical_scale=0.005,
slope_threshold=0.75,
use_cache=False,
sub_terrains={
"pyramid_stairs": terrain_gen.MeshPyramidStairsTerrainCfg(
proportion=0.15,
step_height_range=(0.05, 0.23),
step_width=0.3,
platform_width=3.0,
border_width=1.0,
holes=False,
),
"pyramid_stairs_inv": terrain_gen.MeshInvertedPyramidStairsTerrainCfg(
proportion=0.20,
step_height_range=(0.05, 0.23),
step_width=0.3,
platform_width=3.0,
border_width=1.0,
holes=False,
),
"boxes": terrain_gen.MeshRandomGridTerrainCfg(
proportion=0.15, grid_width=0.45, grid_height_range=(0.05, 0.2), platform_width=2.0
),
"random_rough": terrain_gen.HfRandomUniformTerrainCfg(
proportion=0.1, noise_range=(0.02, 0.10), noise_step=0.02, border_width=0.25
),
"flat": terrain_gen.MeshPlaneTerrainCfg(proportion=0.2),
"hf_pyramid_slope": terrain_gen.HfPyramidSlopedTerrainCfg(
proportion=0.1, slope_range=(0.0, 0.4), platform_width=2.0, border_width=0.25
),
"hf_pyramid_slope_inv": terrain_gen.HfInvertedPyramidSlopedTerrainCfg(
proportion=0.1, slope_range=(0.0, 0.4), platform_width=2.0, border_width=0.25
),
},
)
##
# Scene definition
##
@configclass
class Go2SceneCfg(InteractiveSceneCfg):
"""Configuration for the terrain scene with the Go2 robot."""
terrain = TerrainImporterCfg(
prim_path="/World/ground",
terrain_type="generator",
terrain_generator=TERRAIN_CFG,
max_init_terrain_level=5,
collision_group=-1,
physics_material=sim_utils.RigidBodyMaterialCfg(
friction_combine_mode="multiply",
static_friction=1.0,
dynamic_friction=1.0,
),
visual_material=sim_utils.MdlFileCfg(
mdl_path=f"{ISAACLAB_NUCLEUS_DIR}/Materials/TilesMarbleSpiderWhiteBrickBondHoned/TilesMarbleSpiderWhiteBrickBondHoned.mdl",
project_uvw=True,
texture_scale=(0.25, 0.25),
),
debug_vis=False
)
robot: ArticulationCfg = UNITREE_GO2_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot")
height_scanner = RayCasterCfg(
prim_path="{ENV_REGEX_NS}/Robot/base",
offset=RayCasterCfg.OffsetCfg(pos=(0.0, 0.0, 20.0)),
ray_alignment="yaw",
pattern_cfg=patterns.GridPatternCfg(resolution=0.1, size=[1.6, 1.0]),
debug_vis=False,
mesh_prim_paths=["/World/ground"],
)
height_scanner_base = RayCasterCfg(
prim_path="{ENV_REGEX_NS}/Robot/base",
offset=RayCasterCfg.OffsetCfg(pos=(0.0, 0.0, 20.0)),
ray_alignment="yaw",
pattern_cfg=patterns.GridPatternCfg(resolution=0.05, size=(0.1, 0.1)),
debug_vis=False,
mesh_prim_paths=["/World/ground"],
)
contact_forces = ContactSensorCfg(
prim_path="{ENV_REGEX_NS}/Robot/.*",
history_length=3,
track_air_time=True,
)
# 灯光
sky_light = AssetBaseCfg(
prim_path="/World/skyLight",
spawn=sim_utils.DomeLightCfg(
intensity=750.0,
texture_file=f"{ISAAC_NUCLEUS_DIR}/Materials/Textures/Skies/PolyHaven/kloofendal_43d_clear_puresky_4k.hdr",
),
)
##
# MDP settings
##
@configclass
class CommandsCfg:
"""Command specifications for the MDP."""
base_velocity = mdp.UniformVelTerrainCmdCfg(
asset_name="robot",
resampling_time_range=(5.0, 5.0),
rel_standing_envs=0.1,
rel_heading_envs=1.0,
heading_command=False, # note here, no heading command!
heading_control_stiffness=0.5,
debug_vis=True,
cycle_time=0.5,
use_phase_cmd=False,
ranges = mdp.UniformVelocityCommandCfg.Ranges(
lin_vel_x=(-0.5, 0.5), lin_vel_y=(-0.5, 0.5), ang_vel_z=(-1.0, 1.0), heading=(-1.57, 1.57)
),
terrain_max_ranges={
"pyramid_stairs": mdp.UniformVelocityCommandCfg.Ranges(
lin_vel_x=(-1.0, 1.0), lin_vel_y=(-1.0, 1.0), ang_vel_z=(-1.0, 1.0), heading=(-1.57, 1.57)
),
"pyramid_stairs_inv": mdp.UniformVelocityCommandCfg.Ranges(
lin_vel_x=(-1.0, 1.0), lin_vel_y=(-1.0, 1.0), ang_vel_z=(-1.0, 1.0), heading=(-1.57, 1.57)
),
"box": mdp.UniformVelocityCommandCfg.Ranges(
lin_vel_x=(-1.0, 1.0), lin_vel_y=(-1.0, 1.0), ang_vel_z=(-1.5, 1.5), heading=(-1.57, 1.57)
),
"random_rough": mdp.UniformVelocityCommandCfg.Ranges(
lin_vel_x=(-1.5, 1.5), lin_vel_y=(-1.0, 1.0), ang_vel_z=(-1.5, 1.5), heading=(-1.57, 1.57)
),
"flat": mdp.UniformVelocityCommandCfg.Ranges(
lin_vel_x=(-2.0, 2.0), lin_vel_y=(-1.0, 1.0), ang_vel_z=(-2.0, 2.0), heading=(-1.57, 1.57)
),
"hf_pyramid_slope": mdp.UniformVelocityCommandCfg.Ranges(
lin_vel_x=(-1.5, 1.5), lin_vel_y=(-1.0, 1.0), ang_vel_z=(-1.5, 1.5), heading=(-1.57, 1.57)
),
"hf_pyramid_slope_inv": mdp.UniformVelocityCommandCfg.Ranges(
lin_vel_x=(-1.5, 1.5), lin_vel_y=(-1.0, 1.0), ang_vel_z=(-1.5, 1.5), heading=(-1.57, 1.57)
),
},
curriculum_schedule=[
{'iter': 50000, 'lin_vel_x': [-2.0, 2.0], 'lin_vel_y': [-1.0, 1.0], 'ang_vel_yaw': [-2.0, 2.0], 'heading': [-1.57, 1.57]},
{'iter': 20000, 'lin_vel_x': [-1.0, 1.0], 'lin_vel_y': [-1.0, 1.0], 'ang_vel_yaw': [-1.5, 1.5], 'heading': [-1.57, 1.57]},
], # 根据iter倒序排布
)
@configclass
class ActionsCfg:
"""Action specifications for the MDP."""
# 腿部关节：位置控制
joint_pos = mdp.JointPositionActionCfg(
asset_name="robot",
joint_names=JOINT_NAMES,
scale={".*_hip_joint": 0.25, "^(?!.*_hip_joint).*": 0.25},
use_default_offset=True,
clip=None,
preserve_order=True
)
@configclass
class ObservationsCfg:
@configclass
class PolicyCfg(ObsGroup):
"""Observations for policy group."""
base_ang_vel = ObsTerm(
func=mdp.base_ang_vel,
noise=Unoise(n_min=-0.2, n_max=0.2),
clip=(-100.0, 100.0),
scale=0.25,
)
projected_gravity = ObsTerm(
func=mdp.projected_gravity,
noise=Unoise(n_min=-0.05, n_max=0.05),
clip=(-100.0, 100.0),
scale=1.0,
)
velocity_commands = ObsTerm(
func=mdp.generated_commands,
params={"command_name": "base_velocity"},
clip=(-100.0, 100.0),
scale=1.0,
)
joint_pos = ObsTerm(
func=mdp.joint_pos_rel,
params={"asset_cfg": SceneEntityCfg("robot", joint_names=JOINT_NAMES, preserve_order=True)},
noise=Unoise(n_min=-0.03, n_max=0.03),
clip=(-100.0, 100.0),
scale=1.0,
)
joint_vel = ObsTerm(
func=mdp.joint_vel_rel,
params={"asset_cfg": SceneEntityCfg("robot", joint_names=JOINT_NAMES, preserve_order=True)},
noise=Unoise(n_min=-2.0, n_max=2.0),
clip=(-100.0, 100.0),
scale=0.05,
)
actions = ObsTerm(
func=mdp.last_action,
clip=(-100.0, 100.0),
scale=1.0,
)
def __post_init__(self):
self.history_length = 5
self.enable_corruption = True
self.concatenate_terms = True
self.flatten_history_dim = True
@configclass
class CriticCfg(ObsGroup):
base_lin_vel = ObsTerm(
func=mdp.base_lin_vel,
clip=(-100.0, 100.0),
scale=2.0,
)
base_ang_vel = ObsTerm(
func=mdp.base_ang_vel,
clip=(-100.0, 100.0),
scale=0.25,
)
projected_gravity = ObsTerm(
func=mdp.projected_gravity,
clip=(-100.0, 100.0),
scale=1.0,
)
velocity_commands = ObsTerm(
func=mdp.generated_commands,
params={"command_name": "base_velocity"},
clip=(-100.0, 100.0),
scale=1.0,
)
joint_pos = ObsTerm(
func=mdp.joint_pos_rel,
params={"asset_cfg": SceneEntityCfg("robot", joint_names=JOINT_NAMES, preserve_order=True)},
clip=(-100.0, 100.0),
scale=1.0,
)
joint_vel = ObsTerm(
func=mdp.joint_vel_rel,
params={"asset_cfg": SceneEntityCfg("robot", joint_names=JOINT_NAMES, preserve_order=True)},
clip=(-100.0, 100.0),
scale=0.05,
)
actions = ObsTerm(
func=mdp.last_action,
clip=(-100.0, 100.0),
scale=1.0,
)
joint_acc = ObsTerm(
func=mdp.joint_acc,
params={"asset_cfg": SceneEntityCfg("robot", joint_names=JOINT_NAMES, preserve_order=True)},
clip=(-100.0, 100.0),
scale=1e-4,
)
joint_torque = ObsTerm(
func=mdp.joint_effort,
params={"asset_cfg": SceneEntityCfg("robot", joint_names=JOINT_NAMES, preserve_order=True)},
clip=(-100.0, 100.0),
scale=0.01,
)
contact_force = ObsTerm(
func=mdp.foot_contact_force_norm,
params={"sensor_cfg": SceneEntityCfg("contact_forces", body_names=FOOT_LINK_NAME)},
clip=(-100.0, 100.0),
scale=1e-3,
)
height_scan = ObsTerm(
func=mdp.height_scan,
params={"sensor_cfg": SceneEntityCfg("height_scanner")},
clip=(-1.0, 1.0),
scale=2.5,
)
def __post_init__(self):
self.enable_corruption = False
self.concatenate_terms = True
@configclass
class SingleObsCfg(PolicyCfg):
def __post_init__(self):
super().__post_init__()
self.history_length = 1
policy: PolicyCfg = PolicyCfg()
critic: CriticCfg = CriticCfg()
single_obs: SingleObsCfg = SingleObsCfg() # Used to obtain the current-timestep observation for the MoE CTS model
@configclass
class EventCfg:
"""Configuration for events."""
randomize_rigid_body_mass_base = EventTerm(
func=mdp.randomize_rigid_body_mass,
mode="startup",
params={
"asset_cfg": SceneEntityCfg("robot", body_names=BASE_LINK_NAME), # Override
"mass_distribution_params": (-1.0, 1.0),
"operation": "add",
"recompute_inertia": True,
},
)
randomize_rigid_body_mass_others = EventTerm(
func=mdp.randomize_rigid_body_mass,
mode="startup",
params={
"asset_cfg": SceneEntityCfg("robot", body_names="^(?!.*base).*"), # Override
"mass_distribution_params": (0.8, 1.2),
"operation": "scale",
"recompute_inertia": True,
},
)
randomize_com_positions = EventTerm(
func=mdp.randomize_rigid_body_com,
mode="startup",
params={
"asset_cfg": SceneEntityCfg("robot", body_names=BASE_LINK_NAME), # Override
"com_range": {"x": (-0.05, 0.05), "y": (-0.05, 0.05), "z": (-0.05, 0.05)},
},
)
randomize_com_positions_other = EventTerm(
func=mdp.randomize_rigid_body_com,
mode="startup",
params={
"asset_cfg": SceneEntityCfg("robot", body_names="^(?!.*base).*"), # Override
"com_range": {"x": (-0.05, 0.05), "y": (-0.05, 0.05), "z": (-0.05, 0.05)},
},
)
reset_robot_joints = EventTerm(
func=mdp.reset_joints_by_scale,
mode="reset",
params={
"position_range": (0.8, 1.2),
"velocity_range": (0.0, 0.0),
},
)
# when use DelayedPDActuator, the range should be narrowed, range that is too wide like (0.5, 2.0) will cause training to fail.
randomize_actuator_gains = EventTerm(
func=mdp.randomize_actuator_gains,
mode="reset",
params={
"asset_cfg": SceneEntityCfg("robot", joint_names=".*"),
"stiffness_distribution_params": (0.9, 1.1),
"damping_distribution_params": (0.9, 1.1),
"operation": "scale",
"distribution": "uniform",
},
)
randomize_push_robot = EventTerm(
func=mdp.push_by_setting_velocity,
mode="interval",
interval_range_s=(4.0, 4.0),
params={
"velocity_range": {
"x": (-0.5, 0.5),
"y": (-0.5, 0.5),
"roll": (-0.7, 0.7),
"pitch": (-0.7, 0.7),
"yaw": (-0.7, 0.7)
}
}
)
randomize_rigid_body_mass_base = EventTerm(
func=mdp.randomize_rigid_body_mass,
mode="startup",
params={
"asset_cfg": SceneEntityCfg("robot", body_names=BASE_LINK_NAME),
"mass_distribution_params": (-5.0, 10.0),
"operation": "add",
"recompute_inertia": True,
},
)
randomize_rigid_body_material = EventTerm(
func=mdp.randomize_rigid_body_material,
mode="startup",
params={
"asset_cfg": SceneEntityCfg("robot", body_names=".*"),
"static_friction_range": (0.5, 1.5),
"dynamic_friction_range": (0.5, 1.5),
"restitution_range": (0.0, 0.5),
"num_buckets": 512,
"make_consistent": True
},
)
randomize_apply_external_force_torque = EventTerm(
func=mdp.apply_external_force_torque,
mode="reset",
params={
"asset_cfg": SceneEntityCfg("robot", body_names=BASE_LINK_NAME),
"force_range": (-10.0, 10.0),
"torque_range": (-10.0, 10.0),
},
)
reset_base = EventTerm(
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.2), "yaw": (-3.14, 3.14)},
"velocity_range": {
"x": (-0.5, 0.5),
"y": (-0.5, 0.5),
"z": (-0.5, 0.5),
"roll": (-0.5, 0.5),
"pitch": (-0.5, 0.5),
"yaw": (-0.5, 0.5),
},
},
)
@configclass
class RewardsCfg:
"""Reward terms for the MDP."""
track_lin_vel_xy_exp = RewTerm(
func=mdp.track_lin_vel_xy_exp,
weight=1.0,
params={"command_name": "base_velocity", "std": 0.5}
)
track_ang_vel_z_exp = RewTerm(
func=mdp.track_ang_vel_z_exp,
weight=0.5,
params={"command_name": "base_velocity", "std": 0.5}
)
lin_vel_z_l2 = RewTerm(func=mdp.lin_vel_z_l2, weight=-2.0)
ang_vel_xy_l2 = RewTerm(func=mdp.ang_vel_xy_l2, weight=-0.05)
flat_orientation_l2 = RewTerm(
func=mdp.flat_orientation_l2,
weight=-3.0
)
base_height_l2 = RewTerm(
func=mdp.base_height_l2,
weight=-50.0,
params={
"asset_cfg": SceneEntityCfg("robot", body_names=BASE_LINK_NAME),
"target_height": 0.38
}
)
dof_acc_l2 = RewTerm(
func=mdp.joint_acc_l2,
weight=-2.0e-7,
params={"asset_cfg": SceneEntityCfg("robot", joint_names=JOINT_NAMES)}
)
action_rate_l2 = RewTerm(func=mdp.action_rate_l2, weight=-0.01)
action_smoothness_l2 = RewTerm(func=mdp.action_smoothness_l2, weight=-0.01)
undesired_contacts = RewTerm(
func=mdp.undesired_contacts,
weight=-1.0,
params={"sensor_cfg": SceneEntityCfg("contact_forces", body_names=".*_thigh|.*_calf"), "threshold": 5.0},
)
hip_pos_penalty = RewTerm(
func=mdp.hip_pos_penalty,
weight=-0.85,
params={
"command_name": "base_velocity",
"asset_cfg": SceneEntityCfg("robot", joint_names=".*_hip_joint"),
"stand_still_scale": 1.0,
"command_threshold": 0.1,
},
)
joint_torques_l2 = RewTerm(
func=mdp.joint_torques_l2,
weight=-1e-4,
params={"asset_cfg": SceneEntityCfg("robot", joint_names=JOINT_NAMES)}
)
joint_power = RewTerm(
func=mdp.joint_power,
weight=-2e-5,
params={"asset_cfg": SceneEntityCfg("robot", joint_names=JOINT_NAMES)}
)
@configclass
class TerminationsCfg:
"""Termination terms for the MDP."""
time_out = DoneTerm(func=mdp.time_out, time_out=True)
bad_orientation = DoneTerm(
func=mdp.bad_orientation,
params={
"asset_cfg": SceneEntityCfg("robot", body_names=BASE_LINK_NAME),
"limit_angle": 2 * math.pi / 3
},
)
@configclass
class CurriculumCfg:
"""Curriculum terms for the MDP."""
terrain_levels = CurrTerm(func=mdp.terrain_levels_vel_gym)
base_linear_velocity = CurrTerm(mdp.gradual_reward_weight_modification, params={
"term_name": "lin_vel_z_l2", "initial_weight": -2.0, "final_weight": -0.0, "start_it": 0, "end_it": 2000
})
ref_stand_envs = CurrTerm(mdp.gradual_ref_stand_modification, params={
"term_name": "base_velocity", "initial": 0.0, "final": 0.1, "start_it": 0, "end_it": 1500
})
command_curr = CurrTerm(mdp.command_curriculum, params={
"command_term_name": "base_velocity", "num_steps_per_iter": 24,
})
##
# Environment configuration
##
@configclass
class Go2EnvCfg(ManagerBasedRLEnvCfg):
"""Merged configuration for the Go2 robot on rough terrain."""
# Scene settings
scene: Go2SceneCfg = Go2SceneCfg(num_envs=8192, env_spacing=2.5)
# Basic settings
observations: ObservationsCfg = ObservationsCfg()
actions: ActionsCfg = ActionsCfg()
commands: CommandsCfg = CommandsCfg()
# MDP settings
rewards: RewardsCfg = RewardsCfg()
terminations: TerminationsCfg = TerminationsCfg()
events: EventCfg = EventCfg()
curriculum: CurriculumCfg = CurriculumCfg()
def __post_init__(self):
"""Post initialization."""
# General settings
self.decimation = 4
self.episode_length_s = 25.0
# Simulation settings
self.sim.dt = 0.005
self.sim.render_interval = self.decimation
# Physics material settings from subclass
self.sim.physics_material = self.scene.terrain.physics_material
self.sim.physx.gpu_max_rigid_patch_count = 10 * 2**15
self.sim.physx.gpu_collision_stack_size = int(128 * 1024 * 1024) # 128 MB
self.sim.physx.enable_external_forces_every_iteration = True
# Update sensor periods
if self.scene.height_scanner is not None:
self.scene.height_scanner.update_period = self.decimation * self.sim.dt
if self.scene.contact_forces is not None:
self.scene.contact_forces.update_period = self.sim.dt
# Scale down terrain grid for small robot
if self.scene.terrain.terrain_generator is not None:
if "boxes" in self.scene.terrain.terrain_generator.sub_terrains:
self.scene.terrain.terrain_generator.sub_terrains["boxes"].grid_height_range = (0.025, 0.1)
if "random_rough" in self.scene.terrain.terrain_generator.sub_terrains:
self.scene.terrain.terrain_generator.sub_terrains["random_rough"].noise_range = (0.01, 0.06)
self.scene.terrain.terrain_generator.sub_terrains["random_rough"].noise_step = 0.01
# Handle curriculum for terrain generator
if getattr(self.curriculum, "terrain_levels", None) is not None:
if self.scene.terrain.terrain_generator is not None:
self.scene.terrain.terrain_generator.curriculum = True
else:
if self.scene.terrain.terrain_generator is not None:
self.scene.terrain.terrain_generator.curriculum = False

47
source/robot_lab/robot_lab/tasks/go2/manager/action_manager.py Normal file
View File

@@ -0,0 +1,47 @@
from isaaclab.managers import ActionManager
import torch
from collections.abc import Sequence
class ActionManagerWithDelay(ActionManager):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self._prev_prev_action = torch.zeros_like(self._action)
def reset(self, env_ids: Sequence[int] | None = None) -> dict[str, torch.Tensor]:
super().reset(env_ids)
self._prev_prev_action.zero_()
return {}
def process_action(self, action: torch.Tensor):
raise NotImplementedError("Use update_action() and process_action_with_delay() instead for ActionManagerWithDelay.")
def process_action_with_delay(self, action_delay_masks: torch.Tensor):
"""Processes the actions sent to the environment.
Important Note:
This function can be called multiple times within a single step() call to implement action delay.
Args:
action_delay_masks: A tensor of shape (num_envs, 1) indicating which actions to apply at this time step.
"""
# action_delay_masks == True means delay (use prev action)
action = torch.where(action_delay_masks, self._prev_action, self._action)
# split the actions and apply to each tensor
# NOTE: we also assume that all term.process_actions can be called multiple times within a single step() call
idx = 0
for term in self._terms.values():
term_actions = action[:, idx : idx + term.action_dim]
term.process_actions(term_actions)
idx += term.action_dim
def update_action(self, action: torch.Tensor):
if self.total_action_dim != action.shape[1]:
raise ValueError(f"Invalid action shape, expected: {self.total_action_dim}, received: {action.shape[1]}.")
self._prev_prev_action[:] = self._prev_action
self._prev_action[:] = self._action
self._action[:] = action.to(self.device)
@property
def prev_prev_action(self):
return self._prev_prev_action

19
source/robot_lab/robot_lab/tasks/go2/mdp/__init__.py Normal file
View File

@@ -0,0 +1,19 @@
# Copyright (c) 2024-2025 Ziqi Fan
# SPDX-License-Identifier: Apache-2.0
# Copyright (c) 2022-2025, The Isaac Lab Project Developers.
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause
"""This sub-module contains the functions that are specific to the locomotion environments."""
from isaaclab.envs.mdp import * # noqa: F401, F403
from isaaclab_tasks.manager_based.locomotion.velocity.mdp import * # noqa: F401, F403
from .commands import * # noqa: F401, F403
from .curriculums import * # noqa: F401, F403
from .events import * # noqa: F401, F403
from .observations import * # noqa: F401, F403
from .rewards import * # noqa: F401, F403
from .utils import * # noqa: F401, F403

397
source/robot_lab/robot_lab/tasks/go2/mdp/commands.py Normal file
View File

@@ -0,0 +1,397 @@
# Copyright (c) 2024-2025 Ziqi Fan
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
import time
import torch
import copy
from collections.abc import Sequence
from typing import TYPE_CHECKING
from dataclasses import MISSING
from itertools import product
from isaaclab.managers import CommandTerm
from isaaclab.utils import configclass
from isaaclab.assets import Articulation
from isaaclab.envs.mdp import UniformVelocityCommandCfg
import isaaclab.utils.math as math_utils
from isaaclab.markers import VisualizationMarkers
# SPDX-License-Identifier: BSD-3-Clause
# Copyright (c) 2022-2025, The Isaac Lab Project Developers
from .utils import is_robot_on_terrain
if TYPE_CHECKING:
from isaaclab.envs import ManagerBasedEnv
class UniformVelTerrainCmd(CommandTerm):
cfg: UniformVelTerrainCmdCfg
def __init__(self, cfg: UniformVelTerrainCmdCfg, env: ManagerBasedRLEnv):
super().__init__(cfg, env)
# obtain the robot asset
# -- robot
self.robot: Articulation = env.scene[cfg.asset_name]
# crete buffers to store the command
# -- command: x vel, y vel, yaw vel, heading
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.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)
# -- metrics
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)
if self.cfg.use_phase_cmd:
self.phase = torch.zeros(self.num_envs, 1, device=self.device)
self.cycle_time = cfg.cycle_time
self.max_angular_envs = cfg.max_angular_envs
self.limit_vel_envs = cfg.limit_vel_envs
self.stop_heading = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device)
# 获取当前环境的地形类型索引
self.terrain_types = list(env.scene.terrain.cfg.terrain_generator.sub_terrains.keys())
if not isinstance(self.cfg.ranges, dict):
single_range = cfg.ranges
expanded_ranges = {}
for t_name in self.terrain_types:
expanded_ranges[t_name] = copy.deepcopy(single_range)
self.cfg.ranges = expanded_ranges
assert set(self.terrain_types) == set(list(self.cfg.ranges.keys())), \
"Terrain types in cfg.ranges do not match those in terrain generator config."
self.terrain_type_to_id = {name: i for i, name in enumerate(self.terrain_types)}
self.global_env_terrain_idx = torch.full((self.num_envs,), -1, dtype=torch.long, device=self.device)
for t_type in self.terrain_types:
ids = is_robot_on_terrain(self._env, t_type).nonzero(as_tuple=True)[0]
if len(ids) > 0:
self.global_env_terrain_idx[ids] = self.terrain_type_to_id[t_type]
# 初始化累计指令和最大位移
self.commands_xy_accumulation = torch.zeros(self.num_envs, 2, device=self.device)
self.dt = self._env.step_dt
self.max_episode_length = self._env.max_episode_length
self.target_dist = env.scene.terrain.cfg.terrain_generator.size[0] * 0.625
self.max_move_distance = torch.zeros(self.num_envs, device=self.device)
self.env_origins = env.scene.env_origins
# 极限指令组合
limit_options_x = [-1, 1]
limit_options_y = [-1, 1]
limit_options_z = [-1, 0, 1] # 允许 Yaw 为 0，即直线冲刺
self.limit_vel_combinations = torch.tensor(
list(product(limit_options_x, limit_options_y, limit_options_z)),
dtype=torch.long,
device=self.device
)
self.last_is_limit_vel = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device)
def __str__(self) -> str:
"""Return a string representation of the command generator."""
msg = "UniformVelocityCommand:\n"
msg += f"\tCommand dimension: {tuple(self.command.shape[1:])}\n"
msg += f"\tResampling time range: {self.cfg.resampling_time_range}\n"
msg += f"\tHeading command: {self.cfg.heading_command}\n"
if self.cfg.heading_command:
msg += f"\tHeading probability: {self.cfg.rel_heading_envs}\n"
msg += f"\tStanding probability: {self.cfg.rel_standing_envs}\n"
msg += f"\tMax angular vel envs: {self.max_angular_envs}\n"
msg += f"\tLimit linear vel envs: {self.limit_vel_envs}\n"
for t_type in self.terrain_types:
msg += f"\tTerrain type '{t_type}' command ranges: {self.cfg.ranges[t_type]}"
return msg
@property
def command(self) -> torch.Tensor:
if self.cfg.use_phase_cmd:
phase = self.phase * 2 * torch.pi
return torch.cat([self.vel_command_b, torch.sin(phase), torch.cos(phase)], dim=-1)
return self.vel_command_b
def _update_metrics(self):
# time for which the command was executed
max_command_time = self.cfg.resampling_time_range[1]
max_command_step = max_command_time / self._env.step_dt
# logs data
self.metrics["error_vel_xy"] += (
torch.norm(self.vel_command_b[:, :2] - self.robot.data.root_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_ang_vel_b[:, 2]) / max_command_step
)
def reset(self, env_ids: Sequence[int] | None = None):
self.commands_xy_accumulation[env_ids] = 0.0
self.max_move_distance[env_ids] = 0.0
extra = super().reset(env_ids)
return extra
def _terrain_cmd_sample(self, dim: str, env_ids: torch.Tensor, min_abs_vel: torch.Tensor | None = None):
"""根据不同地形sample不同指令"""
cmd = torch.zeros(len(env_ids), device=self.device)
batch_terrain_idxs = self.global_env_terrain_idx[env_ids]
for t_type in self.terrain_types:
t_id = self.terrain_type_to_id[t_type]
mask = (batch_terrain_idxs == t_id)
if mask.any():
count = mask.sum().item()
if dim == "x":
r = self.cfg.ranges[t_type].lin_vel_x
elif dim == "y":
r = self.cfg.ranges[t_type].lin_vel_y
elif dim == "z":
r = self.cfg.ranges[t_type].ang_vel_z
elif dim == "heading":
r = self.cfg.ranges[t_type].heading
sampled_vals = torch.empty(count, device=self.device).uniform_(*r) # type: ignore
if min_abs_vel is not None and (dim == "x" or dim == "y"):
lower_bound = min_abs_vel[mask]
max_abs = max(abs(r[0]), abs(r[1]))
lower_bound = torch.clamp(lower_bound, max=max_abs)
# 重新采样 Magnitude: [lower_bound, max_abs]
mag = torch.empty(count, device=self.device).uniform_(0, 1) * (max_abs - lower_bound) + lower_bound
if r[0] >= 0:
sampled_vals = mag
elif r[1] <= 0:
sampled_vals = -mag
else:
sign = torch.sign(torch.empty(count, device=self.device).uniform_(-1, 1))
sign = torch.where(sign == 0, torch.ones_like(sign), sign)
sampled_vals = mag * sign
cmd[mask] = sampled_vals
return cmd
def _get_cmd_board(self, dim: str, env_ids: torch.Tensor):
"""根据预计算的地形类型获取速度指令边界"""
cmd_min = torch.zeros(len(env_ids), device=self.device)
cmd_max = torch.zeros(len(env_ids), device=self.device)
batch_terrain_idxs = self.global_env_terrain_idx[env_ids]
for t_type in self.terrain_types:
t_id = self.terrain_type_to_id[t_type]
mask = (batch_terrain_idxs == t_id)
if mask.any():
if dim == "x":
r = self.cfg.ranges[t_type].lin_vel_x
elif dim == "y":
r = self.cfg.ranges[t_type].lin_vel_y
elif dim == "z":
r = self.cfg.ranges[t_type].ang_vel_z
else:
continue
cmd_min[mask] = r[0]
cmd_max[mask] = r[1]
return cmd_min, cmd_max
def _resample_command(self, env_ids: Sequence[int]):
_env_ids = torch.tensor(env_ids, device=self.device)
# 计算剩余距离: 目标距离 - 已经指令累积走过的距离 * 上一次重采样时间
dist_covered = torch.norm(self.commands_xy_accumulation[env_ids], dim=1) * self.cfg.resampling_time_range[0]
remaining_dist = torch.clamp(self.target_dist - dist_covered, min=0.0)
# 计算剩余时间: (最大步数 - 当前步数) * dt
time_left = (self.max_episode_length - self._env.episode_length_buf[env_ids]) * self.dt
# 计算下限速度: 距离 / 时间
vel_low_bound = torch.zeros(len(env_ids), device=self.device)
# 避免除以零或负数时间
valid_time_mask = time_left > 1e-4
if valid_time_mask.any():
vel_low_bound[valid_time_mask] = remaining_dist[valid_time_mask] / time_left[valid_time_mask]
# 根据不同地形类型采样速度指令
self.vel_command_b[env_ids, 0] = self._terrain_cmd_sample("x", _env_ids, min_abs_vel=vel_low_bound)
self.vel_command_b[env_ids, 0] = torch.where(
self.vel_command_b[env_ids, 0].abs() < 0.1,
torch.zeros_like(self.vel_command_b[env_ids, 0]),
self.vel_command_b[env_ids, 0]
)
self.vel_command_b[env_ids, 1] = self._terrain_cmd_sample("y", _env_ids, min_abs_vel=vel_low_bound)
self.vel_command_b[env_ids, 1] = torch.where(
self.vel_command_b[env_ids, 1].abs() < 0.1,
torch.zeros_like(self.vel_command_b[env_ids, 1]),
self.vel_command_b[env_ids, 1]
)
self.vel_command_b[env_ids, 2] = self._terrain_cmd_sample("z", _env_ids)
self.vel_command_b[env_ids, 2] = torch.where(
self.vel_command_b[env_ids, 2].abs() < 0.1,
torch.zeros_like(self.vel_command_b[env_ids, 2]),
self.vel_command_b[env_ids, 2]
)
# 计算heading和stand
r = torch.empty(len(env_ids), device=self.device)
if self.cfg.heading_command:
self.heading_target[env_ids] = self._terrain_cmd_sample("heading", _env_ids)
# update heading envs
self.is_heading_env[env_ids] = r.uniform_(0.0, 1.0) <= self.cfg.rel_heading_envs
# update standing envs
self.is_standing_env[env_ids] = r.uniform_(0.0, 1.0) <= self.cfg.rel_standing_envs
# 静止环境sample最大转向
standing_env_ids = self.is_standing_env.nonzero(as_tuple=False).flatten()
if len(standing_env_ids) > 0:
self.vel_command_b[standing_env_ids, :] = 0.0
ang_vel_rand = torch.rand(len(standing_env_ids), device=self.device)
add_ang_mask = ang_vel_rand < self.max_angular_envs
add_ang_env_ids = standing_env_ids[add_ang_mask]
if self.max_angular_envs > 0 and len(add_ang_env_ids) > 0:
direction_rand = torch.rand(len(add_ang_env_ids), device=self.device)
min_z, max_z = self._get_cmd_board("z", add_ang_env_ids)
self.vel_command_b[add_ang_env_ids, 2] = torch.where(direction_rand < 0.5, min_z, max_z)
self.stop_heading[add_ang_env_ids] = True
limit_vel_env_ids = (self.is_standing_env == 0).nonzero(as_tuple=False).flatten()
# 非静止环境sample极限速度
current_ids = env_ids
mask = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device)
mask[current_ids] = True
target_limit_ids = limit_vel_env_ids[mask[limit_vel_env_ids]]
if self.limit_vel_envs > 0 and len(target_limit_ids) > 0:
limit_prob = torch.rand(len(target_limit_ids), device=self.device)
execute_limit_mask = limit_prob < self.limit_vel_envs
execute_ids = target_limit_ids[execute_limit_mask]
if len(execute_ids) > 0:
num_combs = self.limit_vel_combinations.shape[0]
comb_indices = torch.randint(0, num_combs, (len(execute_ids),), device=self.device)
selected_combs = self.limit_vel_combinations[comb_indices] # (N, 3)
min_x, max_x = self._get_cmd_board("x", execute_ids)
min_y, max_y = self._get_cmd_board("y", execute_ids)
min_z, max_z = self._get_cmd_board("z", execute_ids)
vals_x = torch.zeros_like(min_x)
vals_x = torch.where(selected_combs[:, 0] == -1, min_x, vals_x)
vals_x = torch.where(selected_combs[:, 0] == 1, max_x, vals_x)
vals_y = torch.zeros_like(min_y)
vals_y = torch.where(selected_combs[:, 1] == -1, min_y, vals_y)
vals_y = torch.where(selected_combs[:, 1] == 1, max_y, vals_y)
vals_z = torch.zeros_like(min_z)
vals_z = torch.where(selected_combs[:, 2] == -1, min_z, vals_z)
vals_z = torch.where(selected_combs[:, 2] == 1, max_z, vals_z)
self.vel_command_b[execute_ids, 0] = vals_x
self.vel_command_b[execute_ids, 1] = vals_y
self.vel_command_b[execute_ids, 2] = vals_z
self.stop_heading[execute_ids] = True
if self.cfg.use_phase_cmd:
self.phase[env_ids] = torch.rand((len(env_ids), 1), device=self.device)
self.commands_xy_accumulation[env_ids] += self.vel_command_b[env_ids, :2]
def _update_command(self):
# 停止heading更新
stop_heading_env_ids = self.stop_heading.nonzero(as_tuple=False).flatten()
if len(stop_heading_env_ids) > 0:
self.is_heading_env[stop_heading_env_ids] = False
# Compute angular velocity from heading direction
if self.cfg.heading_command:
# resolve indices of heading envs
env_ids = self.is_heading_env.nonzero(as_tuple=False).flatten()
# compute angular velocity
if len(env_ids) > 0:
heading_error = math_utils.wrap_to_pi(self.heading_target[env_ids] - self.robot.data.heading_w[env_ids])
min_z, max_z = self._get_cmd_board("z", env_ids)
self.vel_command_b[env_ids, 2] = torch.clip(
self.cfg.heading_control_stiffness * heading_error,
min=min_z,
max=max_z,
)
if self.cfg.use_phase_cmd:
self.phase = self._env.episode_length_buf[:, None] * self._env.step_dt / self.cycle_time
current_dist = torch.norm(self.robot.data.root_pos_w[:, :2] - self.env_origins[:, :2], dim=1)
self.max_move_distance = torch.max(self.max_move_distance, current_dist)
def _set_debug_vis_impl(self, debug_vis: bool):
# set visibility of markers
# note: parent only deals with callbacks. not their visibility
if debug_vis:
# create markers if necessary for the first time
if not hasattr(self, "goal_vel_visualizer"):
# -- goal
self.goal_vel_visualizer = VisualizationMarkers(self.cfg.goal_vel_visualizer_cfg)
# -- current
self.current_vel_visualizer = VisualizationMarkers(self.cfg.current_vel_visualizer_cfg)
# set their visibility to true
self.goal_vel_visualizer.set_visibility(True)
self.current_vel_visualizer.set_visibility(True)
else:
if hasattr(self, "goal_vel_visualizer"):
self.goal_vel_visualizer.set_visibility(False)
self.current_vel_visualizer.set_visibility(False)
def _debug_vis_callback(self, event):
# check if robot is initialized
# note: this is needed in-case the robot is de-initialized. we can't access the data
if not self.robot.is_initialized:
return
# get marker location
# -- base state
base_pos_w = self.robot.data.root_pos_w.clone()
base_pos_w[:, 2] += 0.5
# -- resolve the scales and quaternions
vel_des_arrow_scale, vel_des_arrow_quat = self._resolve_xy_velocity_to_arrow(self.command[:, :2])
vel_arrow_scale, vel_arrow_quat = self._resolve_xy_velocity_to_arrow(self.robot.data.root_lin_vel_b[:, :2])
# display markers
self.goal_vel_visualizer.visualize(base_pos_w, vel_des_arrow_quat, vel_des_arrow_scale)
self.current_vel_visualizer.visualize(base_pos_w, vel_arrow_quat, vel_arrow_scale)
"""
Internal helpers.
"""
def _resolve_xy_velocity_to_arrow(self, xy_velocity: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
"""Converts the XY base velocity command to arrow direction rotation."""
# obtain default scale of the marker
default_scale = self.goal_vel_visualizer.cfg.markers["arrow"].scale
# arrow-scale
arrow_scale = torch.tensor(default_scale, device=self.device).repeat(xy_velocity.shape[0], 1)
arrow_scale[:, 0] *= torch.linalg.norm(xy_velocity, dim=1) * 3.0
# arrow-direction
heading_angle = torch.atan2(xy_velocity[:, 1], xy_velocity[:, 0])
zeros = torch.zeros_like(heading_angle)
arrow_quat = math_utils.quat_from_euler_xyz(zeros, zeros, heading_angle)
# convert everything back from base to world frame
base_quat_w = self.robot.data.root_quat_w
arrow_quat = math_utils.quat_mul(base_quat_w, arrow_quat)
return arrow_scale, arrow_quat
@configclass
class UniformVelTerrainCmdCfg(UniformVelocityCommandCfg):
class_type: type = UniformVelTerrainCmd
cycle_time:float = 0.5
max_angular_envs: float = 0.2
limit_vel_envs: float = 0.2
ranges: dict[str, UniformVelocityCommandCfg.Ranges] | UniformVelocityCommandCfg.Ranges = MISSING # 当前指令范围
terrain_max_ranges: dict[str, UniformVelocityCommandCfg.Ranges] = MISSING # type:ignore 地形最大指令范围
curriculum_schedule: list[dict] | None = None
use_phase_cmd: bool = True

216
source/robot_lab/robot_lab/tasks/go2/mdp/curriculums.py Normal file
View File

@@ -0,0 +1,216 @@
# Copyright (c) 2024-2025 Ziqi Fan
# SPDX-License-Identifier: Apache-2.0
"""Common functions that can be used to create curriculum for the learning environment.
The functions can be passed to the :class:`isaaclab.managers.CurriculumTermCfg` object to enable
the curriculum introduced by the function.
"""
from __future__ import annotations
import torch
from collections.abc import Sequence
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from isaaclab.envs import ManagerBasedRLEnv
def command_levels_lin_vel(
env: ManagerBasedRLEnv,
env_ids: Sequence[int],
reward_term_name: str,
range_multiplier: Sequence[float] = (0.1, 1.0),
) -> None:
"""command_levels_lin_vel"""
base_velocity_ranges = env.command_manager.get_term("base_velocity").cfg.ranges
# Get original velocity ranges (ONLY ON FIRST EPISODE)
if env.common_step_counter == 0:
env._original_vel_x = torch.tensor(base_velocity_ranges.lin_vel_x, device=env.device)
env._original_vel_y = torch.tensor(base_velocity_ranges.lin_vel_y, device=env.device)
env._initial_vel_x = env._original_vel_x * range_multiplier[0]
env._final_vel_x = env._original_vel_x * range_multiplier[1]
env._initial_vel_y = env._original_vel_y * range_multiplier[0]
env._final_vel_y = env._original_vel_y * range_multiplier[1]
# Initialize command ranges to initial values
base_velocity_ranges.lin_vel_x = env._initial_vel_x.tolist()
base_velocity_ranges.lin_vel_y = env._initial_vel_y.tolist()
# avoid updating command curriculum at each step since the maximum command is common to all envs
if env.common_step_counter % env.max_episode_length == 0:
episode_sums = env.reward_manager._episode_sums[reward_term_name]
reward_term_cfg = env.reward_manager.get_term_cfg(reward_term_name)
delta_command = torch.tensor([-0.1, 0.1], device=env.device)
# If the tracking reward is above 80% of the maximum, increase the range of commands
if torch.mean(episode_sums[env_ids]) / env.max_episode_length_s > 0.8 * reward_term_cfg.weight:
new_vel_x = torch.tensor(base_velocity_ranges.lin_vel_x, device=env.device) + delta_command
new_vel_y = torch.tensor(base_velocity_ranges.lin_vel_y, device=env.device) + delta_command
# Clamp to ensure we don't exceed final ranges
new_vel_x = torch.clamp(new_vel_x, min=env._final_vel_x[0], max=env._final_vel_x[1])
new_vel_y = torch.clamp(new_vel_y, min=env._final_vel_y[0], max=env._final_vel_y[1])
# Update ranges
base_velocity_ranges.lin_vel_x = new_vel_x.tolist()
base_velocity_ranges.lin_vel_y = new_vel_y.tolist()
return torch.tensor(base_velocity_ranges.lin_vel_x[1], device=env.device)
def command_levels_ang_vel(
env: ManagerBasedRLEnv,
env_ids: Sequence[int],
reward_term_name: str,
range_multiplier: Sequence[float] = (0.1, 1.0),
) -> None:
"""command_levels_ang_vel"""
base_velocity_ranges = env.command_manager.get_term("base_velocity").cfg.ranges
# Get original angular velocity ranges (ONLY ON FIRST EPISODE)
if env.common_step_counter == 0:
env._original_ang_vel_z = torch.tensor(base_velocity_ranges.ang_vel_z, device=env.device)
env._initial_ang_vel_z = env._original_ang_vel_z * range_multiplier[0]
env._final_ang_vel_z = env._original_ang_vel_z * range_multiplier[1]
# Initialize command ranges to initial values
base_velocity_ranges.ang_vel_z = env._initial_ang_vel_z.tolist()
# avoid updating command curriculum at each step since the maximum command is common to all envs
if env.common_step_counter % env.max_episode_length == 0:
episode_sums = env.reward_manager._episode_sums[reward_term_name]
reward_term_cfg = env.reward_manager.get_term_cfg(reward_term_name)
delta_command = torch.tensor([-0.1, 0.1], device=env.device)
# If the tracking reward is above 80% of the maximum, increase the range of commands
if torch.mean(episode_sums[env_ids]) / env.max_episode_length_s > 0.8 * reward_term_cfg.weight:
new_ang_vel_z = torch.tensor(base_velocity_ranges.ang_vel_z, device=env.device) + delta_command
# Clamp to ensure we don't exceed final ranges
new_ang_vel_z = torch.clamp(new_ang_vel_z, min=env._final_ang_vel_z[0], max=env._final_ang_vel_z[1])
# Update ranges
base_velocity_ranges.ang_vel_z = new_ang_vel_z.tolist()
return torch.tensor(base_velocity_ranges.ang_vel_z[1], device=env.device)
def command_curriculum(
env: ManagerBasedRLEnv,
env_ids: Sequence[int],
command_term_name: str,
num_steps_per_iter: int = 24,
) -> float:
"""
阶跃式指令课程 (数据存储在 CommandCfg 中)。
"""
try:
cmd_term = env.command_manager.get_term(command_term_name)
cmd_cfg = cmd_term.cfg
except LookupError:
return 0.0
current_iter = env.common_step_counter // num_steps_per_iter
schedule = cmd_cfg.curriculum_schedule
for i in range(len(schedule) - 1, -1, -1):
stage = schedule[i]
if current_iter >= stage['iter']:
for t_name, active_range in cmd_cfg.ranges.items():
hard_limit = cmd_cfg.terrain_max_ranges.get(t_name)
if hard_limit is None:
continue
def get_intersection(target_val, limit_val):
new_min = max(target_val[0], limit_val[0])
new_max = min(target_val[1], limit_val[1])
return (new_min, new_max)
if 'lin_vel_x' in stage:
new_range = get_intersection(stage['lin_vel_x'], hard_limit.lin_vel_x)
active_range.lin_vel_x = new_range
if 'lin_vel_y' in stage:
new_range = get_intersection(stage['lin_vel_y'], hard_limit.lin_vel_y)
active_range.lin_vel_y = new_range
if 'ang_vel_yaw' in stage:
new_range = get_intersection(stage['ang_vel_yaw'], hard_limit.ang_vel_z)
active_range.ang_vel_z = new_range
if 'heading' in stage and hasattr(active_range, 'heading'):
new_range = get_intersection(stage['heading'], hard_limit.heading)
active_range.heading = new_range
schedule.pop(i)
break
last_key = list(cmd_cfg.ranges.keys())[-1]
return cmd_cfg.ranges[last_key].lin_vel_x[1]
def gradual_ref_stand_modification(
env: ManagerBasedRLEnv,
env_ids: Sequence[int],
term_name: str,
initial: float,
final: float,
start_it: int,
end_it: int,
):
current_it = env.common_step_counter // 24
if current_it < start_it:
return
if current_it >= end_it:
new = final
else:
new = (current_it - start_it) / (end_it - start_it) * (final - initial) + initial
term = env.command_manager.get_term(term_name)
term.cfg.rel_standing_envs = new
def gradual_reward_weight_modification(
env: ManagerBasedRLEnv,
env_ids: Sequence[int],
term_name: str,
initial_weight: float,
final_weight: float,
start_it: int,
end_it: int,
):
"""Curriculum that gradually modifies a reward weight between an initial and final value over a range of steps."""
current_it = env.common_step_counter // 24
if current_it < start_it:
return
if current_it >= end_it:
new_weight = final_weight
else:
new_weight = (current_it - start_it) / (end_it - start_it) * (final_weight - initial_weight) + initial_weight
term_cfg = env.reward_manager.get_term_cfg(term_name)
term_cfg.weight = new_weight
env.reward_manager.set_term_cfg(term_name, term_cfg)
def terrain_levels_vel_gym(env: ManagerBasedRLEnv, env_ids: Sequence[int]) -> float:
"""
使用 max_move_distance 而非 reset 时的瞬间位移, 比较标准基于 commands_xy_accumulation
"""
terrain = env.scene.terrain
command = env.command_manager.get_term("base_velocity")
max_move_dist = command.max_move_distance[env_ids]
cmd_accum = command.commands_xy_accumulation[env_ids]
resampling_time = command.cfg.resampling_time_range[0]
zero_prob = command.cfg.rel_standing_envs
move_up = max_move_dist > terrain.cfg.terrain_generator.size[0] / 2
target_dist = torch.norm(cmd_accum, dim=1) * (resampling_time * (1 - zero_prob))
move_down = (max_move_dist < target_dist * 0.5) * ~move_up
terrain.update_env_origins(env_ids, move_up, move_down)
return torch.mean(terrain.terrain_levels.float()) # type: ignore

269
source/robot_lab/robot_lab/tasks/go2/mdp/events.py Normal file
View File

@@ -0,0 +1,269 @@
# Copyright (c) 2024-2025 Ziqi Fan
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
import torch
from typing import TYPE_CHECKING, Literal
import isaaclab.utils.math as math_utils
from isaaclab.assets import Articulation, RigidObject
from isaaclab.managers import SceneEntityCfg
from .utils import is_env_assigned_to_terrain
if TYPE_CHECKING:
from isaaclab.envs import ManagerBasedEnv
def randomize_rigid_body_inertia(
env: ManagerBasedEnv,
env_ids: torch.Tensor | None,
asset_cfg: SceneEntityCfg,
inertia_distribution_params: tuple[float, float],
operation: Literal["add", "scale", "abs"],
distribution: Literal["uniform", "log_uniform", "gaussian"] = "uniform",
):
"""Randomize the inertia tensors of the bodies by adding, scaling, or setting random values.
This function allows randomizing only the diagonal inertia tensor components (xx, yy, zz) of the bodies.
The function samples random values from the given distribution parameters and adds, scales, or sets the values
into the physics simulation based on the operation.
.. tip::
This function uses CPU tensors to assign the body inertias. It is recommended to use this function
only during the initialization of the environment.
"""
# extract the used quantities (to enable type-hinting)
asset: RigidObject | Articulation = env.scene[asset_cfg.name]
# resolve environment ids
if env_ids is None:
env_ids = torch.arange(env.scene.num_envs, device="cpu")
else:
env_ids = env_ids.cpu()
# resolve body indices
if asset_cfg.body_ids == slice(None):
body_ids = torch.arange(asset.num_bodies, dtype=torch.int, device="cpu")
else:
body_ids = torch.tensor(asset_cfg.body_ids, dtype=torch.int, device="cpu")
# get the current inertia tensors of the bodies (num_assets, num_bodies, 9 for articulations or 9 for rigid objects)
inertias = asset.root_physx_view.get_inertias()
# apply randomization on default values
inertias[env_ids[:, None], body_ids, :] = asset.data.default_inertia[env_ids[:, None], body_ids, :].clone()
# randomize each diagonal element (xx, yy, zz -> indices 0, 4, 8)
for idx in [0, 4, 8]:
# Extract and randomize the specific diagonal element
randomized_inertias = _randomize_prop_by_op(
inertias[:, :, idx],
inertia_distribution_params,
env_ids,
body_ids,
operation,
distribution,
)
# Assign the randomized values back to the inertia tensor
inertias[env_ids[:, None], body_ids, idx] = randomized_inertias
# set the inertia tensors into the physics simulation
asset.root_physx_view.set_inertias(inertias, env_ids)
def randomize_com_positions(
env: ManagerBasedEnv,
env_ids: torch.Tensor | None,
asset_cfg: SceneEntityCfg,
com_distribution_params: tuple[float, float],
operation: Literal["add", "scale", "abs"],
distribution: Literal["uniform", "log_uniform", "gaussian"] = "uniform",
):
"""Randomize the center of mass (COM) positions for the rigid bodies.
This function allows randomizing the COM positions of the bodies in the physics simulation. The positions can be
randomized by adding, scaling, or setting random values sampled from the specified distribution.
.. tip::
This function is intended for initialization or offline adjustments, as it modifies physics properties directly.
Args:
env (ManagerBasedEnv): The simulation environment.
env_ids (torch.Tensor | None): Specific environment indices to apply randomization, or None for all environments.
asset_cfg (SceneEntityCfg): The configuration for the target asset whose COM will be randomized.
com_distribution_params (tuple[float, float]): Parameters of the distribution (e.g., min and max for uniform).
operation (Literal["add", "scale", "abs"]): The operation to apply for randomization.
distribution (Literal["uniform", "log_uniform", "gaussian"]): The distribution to sample random values from.
"""
# Extract the asset (Articulation or RigidObject)
asset: RigidObject | Articulation = env.scene[asset_cfg.name]
# Resolve environment indices
if env_ids is None:
env_ids = torch.arange(env.scene.num_envs, device="cpu")
else:
env_ids = env_ids.cpu()
# Resolve body indices
if asset_cfg.body_ids == slice(None):
body_ids = torch.arange(asset.num_bodies, dtype=torch.int, device="cpu")
else:
body_ids = torch.tensor(asset_cfg.body_ids, dtype=torch.int, device="cpu")
# Get the current COM offsets (num_assets, num_bodies, 3)
com_offsets = asset.root_physx_view.get_coms()
for dim_idx in range(3): # Randomize x, y, z independently
randomized_offset = _randomize_prop_by_op(
com_offsets[:, :, dim_idx],
com_distribution_params,
env_ids,
body_ids,
operation,
distribution,
)
com_offsets[env_ids[:, None], body_ids, dim_idx] = randomized_offset[env_ids[:, None], body_ids]
# Set the randomized COM offsets into the simulation
asset.root_physx_view.set_coms(com_offsets, env_ids)
"""
Internal helper functions.
"""
def _randomize_prop_by_op(
data: torch.Tensor,
distribution_parameters: tuple[float | torch.Tensor, float | torch.Tensor],
dim_0_ids: torch.Tensor | None,
dim_1_ids: torch.Tensor | slice,
operation: Literal["add", "scale", "abs"],
distribution: Literal["uniform", "log_uniform", "gaussian"],
) -> torch.Tensor:
"""Perform data randomization based on the given operation and distribution.
Args:
data: The data tensor to be randomized. Shape is (dim_0, dim_1).
distribution_parameters: The parameters for the distribution to sample values from.
dim_0_ids: The indices of the first dimension to randomize.
dim_1_ids: The indices of the second dimension to randomize.
operation: The operation to perform on the data. Options: 'add', 'scale', 'abs'.
distribution: The distribution to sample the random values from. Options: 'uniform', 'log_uniform'.
Returns:
The data tensor after randomization. Shape is (dim_0, dim_1).
Raises:
NotImplementedError: If the operation or distribution is not supported.
"""
# resolve shape
# -- dim 0
if dim_0_ids is None:
n_dim_0 = data.shape[0]
dim_0_ids = slice(None)
else:
n_dim_0 = len(dim_0_ids)
if not isinstance(dim_1_ids, slice):
dim_0_ids = dim_0_ids[:, None]
# -- dim 1
if isinstance(dim_1_ids, slice):
n_dim_1 = data.shape[1]
else:
n_dim_1 = len(dim_1_ids)
# resolve the distribution
if distribution == "uniform":
dist_fn = math_utils.sample_uniform
elif distribution == "log_uniform":
dist_fn = math_utils.sample_log_uniform
elif distribution == "gaussian":
dist_fn = math_utils.sample_gaussian
else:
raise NotImplementedError(
f"Unknown distribution: '{distribution}' for joint properties randomization."
" Please use 'uniform', 'log_uniform', 'gaussian'."
)
# perform the operation
if operation == "add":
data[dim_0_ids, dim_1_ids] += dist_fn(*distribution_parameters, (n_dim_0, n_dim_1), device=data.device)
elif operation == "scale":
data[dim_0_ids, dim_1_ids] *= dist_fn(*distribution_parameters, (n_dim_0, n_dim_1), device=data.device)
elif operation == "abs":
data[dim_0_ids, dim_1_ids] = dist_fn(*distribution_parameters, (n_dim_0, n_dim_1), device=data.device)
else:
raise NotImplementedError(
f"Unknown operation: '{operation}' for property randomization. Please use 'add', 'scale', or 'abs'."
)
return data
def reset_root_state_uniform(
env: ManagerBasedEnv,
env_ids: torch.Tensor,
pose_range: dict[str, tuple[float, float]],
velocity_range: dict[str, tuple[float, float]],
asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
):
"""Reset the asset root state to a random position and velocity uniformly within the given ranges.
This function randomizes the root position and velocity of the asset.
* It samples the root position from the given ranges and adds them to the default root position, before setting
them into the physics simulation.
* It samples the root orientation from the given ranges and sets them into the physics simulation.
* It samples the root velocity from the given ranges and sets them into the physics simulation.
The function takes a dictionary of pose and velocity ranges for each axis and rotation. The keys of the
dictionary are ``x``, ``y``, ``z``, ``roll``, ``pitch``, and ``yaw``. The values are tuples of the form
``(min, max)``. If the dictionary does not contain a key, the position or velocity is set to zero for that axis.
Note: If "pits" terrain exists, environments on pit terrain will be reset to default state without random
perturbations to avoid the robot falling into the pit.
"""
# extract the used quantities (to enable type-hinting)
asset: RigidObject | Articulation = env.scene[asset_cfg.name]
# Separate pit and non-pit environments
# Check which environments are assigned to pit terrain (not random reset)
assigned_to_pits = is_env_assigned_to_terrain(env, "pits")
pit_env_ids = env_ids[assigned_to_pits[env_ids]]
non_pit_env_ids = env_ids[~assigned_to_pits[env_ids]]
# Reset pit environments to default state (no random perturbations)
if len(pit_env_ids) > 0:
root_states = asset.data.default_root_state[pit_env_ids].clone()
positions = root_states[:, 0:3] + env.scene.env_origins[pit_env_ids]
orientations = root_states[:, 3:7]
velocities = torch.zeros_like(root_states[:, 7:13])
asset.write_root_pose_to_sim(torch.cat([positions, orientations], dim=-1), env_ids=pit_env_ids)
asset.write_root_velocity_to_sim(velocities, env_ids=pit_env_ids)
# Reset non-pit environments with random perturbations
if len(non_pit_env_ids) > 0:
root_states = asset.data.default_root_state[non_pit_env_ids].clone()
# poses
range_list = [pose_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z", "roll", "pitch", "yaw"]]
ranges = torch.tensor(range_list, device=asset.device)
rand_samples = math_utils.sample_uniform(
ranges[:, 0], ranges[:, 1], (len(non_pit_env_ids), 6), device=asset.device
)
positions = root_states[:, 0:3] + env.scene.env_origins[non_pit_env_ids] + rand_samples[:, 0:3]
orientations_delta = math_utils.quat_from_euler_xyz(rand_samples[:, 3], rand_samples[:, 4], rand_samples[:, 5])
orientations = math_utils.quat_mul(root_states[:, 3:7], orientations_delta)
# velocities
range_list = [velocity_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z", "roll", "pitch", "yaw"]]
ranges = torch.tensor(range_list, device=asset.device)
rand_samples = math_utils.sample_uniform(
ranges[:, 0], ranges[:, 1], (len(non_pit_env_ids), 6), device=asset.device
)
velocities = root_states[:, 7:13] + rand_samples
# set into the physics simulation
asset.write_root_pose_to_sim(torch.cat([positions, orientations], dim=-1), env_ids=non_pit_env_ids)
asset.write_root_velocity_to_sim(velocities, env_ids=non_pit_env_ids)

49
source/robot_lab/robot_lab/tasks/go2/mdp/observations.py Normal file
View File

@@ -0,0 +1,49 @@
# Copyright (c) 2024-2025 Ziqi Fan
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
import torch
from typing import TYPE_CHECKING
from isaaclab.assets import Articulation
from isaaclab.managers import SceneEntityCfg
from isaaclab.sensors import ContactSensor
if TYPE_CHECKING:
from isaaclab.envs import ManagerBasedEnv, ManagerBasedRLEnv
def joint_pos_rel_without_wheel(
env: ManagerBasedEnv,
asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
wheel_asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
) -> torch.Tensor:
"""The joint positions of the asset w.r.t. the default joint positions.(Without the wheel joints)"""
# extract the used quantities (to enable type-hinting)
asset: Articulation = env.scene[asset_cfg.name]
joint_pos_rel = asset.data.joint_pos[:, asset_cfg.joint_ids] - asset.data.default_joint_pos[:, asset_cfg.joint_ids]
joint_pos_rel[:, wheel_asset_cfg.joint_ids] = 0
return joint_pos_rel
def phase(env: ManagerBasedRLEnv, cycle_time: float) -> torch.Tensor:
if not hasattr(env, "episode_length_buf") or env.episode_length_buf is None:
env.episode_length_buf = torch.zeros(env.num_envs, device=env.device, dtype=torch.long)
phase = env.episode_length_buf[:, None] * env.step_dt / cycle_time
phase_tensor = torch.cat([torch.sin(2 * torch.pi * phase), torch.cos(2 * torch.pi * phase)], dim=-1)
return phase_tensor
def joint_acc(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
asset: Articulation = env.scene[asset_cfg.name]
return asset.data.joint_acc[:, asset_cfg.joint_ids]
def foot_contact_force_norm(env: ManagerBasedRLEnv, sensor_cfg: SceneEntityCfg) -> torch.Tensor:
contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name]
net_contact_forces = contact_sensor.data.net_forces_w_history # [B, T_hist, num_bodies, 3]
contact_force_norm = torch.norm(net_contact_forces[:, :, sensor_cfg.body_ids], dim=-1) # [B, T_hist, num_legs]
max_contact_force_norm, _ = torch.max(contact_force_norm, dim=1) # [B, num_legs]
contact_force_norm = torch.concat([max_contact_force_norm.unsqueeze(1), contact_force_norm], dim=1) # [B, T_hist+1, num_legs]
return contact_force_norm.flatten(start_dim=-2) # [B, (T_hist+1)*num_legs]

557
source/robot_lab/robot_lab/tasks/go2/mdp/rewards.py Normal file
View File

@@ -0,0 +1,557 @@
# Copyright (c) 2024-2025 Ziqi Fan
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
import torch
from typing import TYPE_CHECKING
import isaaclab.utils.math as math_utils
from isaaclab.assets import Articulation, RigidObject
from isaaclab.envs import mdp
from isaaclab.managers import ManagerTermBase
from isaaclab.managers import RewardTermCfg as RewTerm
from isaaclab.managers import SceneEntityCfg
from isaaclab.sensors import ContactSensor, RayCaster
from isaaclab.utils.math import quat_apply_inverse, yaw_quat
if TYPE_CHECKING:
from isaaclab.envs import ManagerBasedRLEnv
def track_lin_vel_xy_exp(
env: ManagerBasedRLEnv, std: float, command_name: str, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")
) -> torch.Tensor:
"""Reward tracking of linear velocity commands (xy axes) using exponential kernel."""
# extract the used quantities (to enable type-hinting)
asset: RigidObject = env.scene[asset_cfg.name]
# compute the error
lin_vel_error = torch.sum(
torch.square(env.command_manager.get_command(command_name)[:, :2] - asset.data.root_lin_vel_b[:, :2]),
dim=1,
)
reward = torch.exp(-lin_vel_error / std**2)
reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
return reward
def track_ang_vel_z_exp(
env: ManagerBasedRLEnv, std: float, command_name: str, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")
) -> torch.Tensor:
"""Reward tracking of angular velocity commands (yaw) using exponential kernel."""
# extract the used quantities (to enable type-hinting)
asset: RigidObject = env.scene[asset_cfg.name]
# compute the error
ang_vel_error = torch.square(env.command_manager.get_command(command_name)[:, 2] - asset.data.root_ang_vel_b[:, 2])
reward = torch.exp(-ang_vel_error / std**2)
reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
return reward
def joint_power(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
"""Reward joint_power"""
# extract the used quantities (to enable type-hinting)
asset: Articulation = env.scene[asset_cfg.name]
# compute the reward
reward = torch.sum(
torch.abs(asset.data.joint_vel[:, asset_cfg.joint_ids] * asset.data.applied_torque[:, asset_cfg.joint_ids]),
dim=1,
)
return reward
def stand_still(
env: ManagerBasedRLEnv,
command_name: str,
command_threshold: float = 0.06,
asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
) -> torch.Tensor:
"""Penalize offsets from the default joint positions when the command is very small."""
# Penalize motion when command is nearly zero.
reward = mdp.joint_deviation_l1(env, asset_cfg)
reward *= torch.norm(env.command_manager.get_command(command_name), dim=1) < command_threshold
reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
return reward
def joint_pos_penalty(
env: ManagerBasedRLEnv,
command_name: str,
asset_cfg: SceneEntityCfg,
stand_still_scale: float,
velocity_threshold: float,
command_threshold: float,
) -> torch.Tensor:
"""Penalize joint position error from default on the articulation."""
# extract the used quantities (to enable type-hinting)
asset: Articulation = env.scene[asset_cfg.name]
cmd = torch.linalg.norm(env.command_manager.get_command(command_name), dim=1)
body_vel = torch.linalg.norm(asset.data.root_lin_vel_b[:, :2], dim=1)
running_reward = torch.linalg.norm(
(asset.data.joint_pos[:, asset_cfg.joint_ids] - asset.data.default_joint_pos[:, asset_cfg.joint_ids]), dim=1
)
reward = torch.where(
torch.logical_or(cmd > command_threshold, body_vel > velocity_threshold),
running_reward,
stand_still_scale * running_reward,
)
reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
return reward
def joint_mirror(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg, mirror_joints: list[list[str]]) -> torch.Tensor:
# extract the used quantities (to enable type-hinting)
asset: Articulation = env.scene[asset_cfg.name]
if not hasattr(env, "joint_mirror_joints_cache") or env.joint_mirror_joints_cache is None:
# Cache joint positions for all pairs
env.joint_mirror_joints_cache = [
[asset.find_joints(joint_name) for joint_name in joint_pair] for joint_pair in mirror_joints
]
reward = torch.zeros(env.num_envs, device=env.device)
# Iterate over all joint pairs
for joint_pair in env.joint_mirror_joints_cache:
# Calculate the difference for each pair and add to the total reward
diff = torch.sum(
torch.square(asset.data.joint_pos[:, joint_pair[0][0]] - asset.data.joint_pos[:, joint_pair[1][0]]),
dim=-1,
)
reward += diff
reward *= 1 / len(mirror_joints) if len(mirror_joints) > 0 else 0
reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
return reward
def action_mirror(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg, mirror_joints: list[list[str]]) -> torch.Tensor:
# extract the used quantities (to enable type-hinting)
asset: Articulation = env.scene[asset_cfg.name]
if not hasattr(env, "action_mirror_joints_cache") or env.action_mirror_joints_cache is None:
# Cache joint positions for all pairs
env.action_mirror_joints_cache = [
[asset.find_joints(joint_name) for joint_name in joint_pair] for joint_pair in mirror_joints
]
reward = torch.zeros(env.num_envs, device=env.device)
# Iterate over all joint pairs
for joint_pair in env.action_mirror_joints_cache:
# Calculate the difference for each pair and add to the total reward
diff = torch.sum(
torch.square(
torch.abs(env.action_manager.action[:, joint_pair[0][0]])
- torch.abs(env.action_manager.action[:, joint_pair[1][0]])
),
dim=-1,
)
reward += diff
reward *= 1 / len(mirror_joints) if len(mirror_joints) > 0 else 0
reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
return reward
def action_sync(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg, joint_groups: list[list[str]]) -> torch.Tensor:
# extract the used quantities (to enable type-hinting)
asset: Articulation = env.scene[asset_cfg.name]
# Cache joint indices if not already done
if not hasattr(env, "action_sync_joint_cache") or env.action_sync_joint_cache is None:
env.action_sync_joint_cache = [
[asset.find_joints(joint_name) for joint_name in joint_group] for joint_group in joint_groups
]
reward = torch.zeros(env.num_envs, device=env.device)
# Iterate over each joint group
for joint_group in env.action_sync_joint_cache:
if len(joint_group) < 2:
continue # need at least 2 joints to compare
# Get absolute actions for all joints in this group
actions = torch.stack(
[torch.abs(env.action_manager.action[:, joint[0]]) for joint in joint_group], dim=1
) # shape: (num_envs, num_joints_in_group)
# Calculate mean action for each environment
mean_actions = torch.mean(actions, dim=1, keepdim=True)
# Calculate variance from mean for each joint
variance = torch.mean(torch.square(actions - mean_actions), dim=1)
# Add to reward (we want to minimize this variance)
reward += variance.squeeze()
reward *= 1 / len(joint_groups) if len(joint_groups) > 0 else 0
reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
return reward
def feet_air_time(
env: ManagerBasedRLEnv, command_name: str, sensor_cfg: SceneEntityCfg, threshold: float
) -> torch.Tensor:
"""Reward long steps taken by the feet using L2-kernel.
This function rewards the agent for taking steps that are longer than a threshold. This helps ensure
that the robot lifts its feet off the ground and takes steps. The reward is computed as the sum of
the time for which the feet are in the air.
If the commands are small (i.e. the agent is not supposed to take a step), then the reward is zero.
"""
# extract the used quantities (to enable type-hinting)
contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name]
# compute the reward
first_contact = contact_sensor.compute_first_contact(env.step_dt)[:, sensor_cfg.body_ids]
last_air_time = contact_sensor.data.last_air_time[:, sensor_cfg.body_ids]
reward = torch.sum((last_air_time - threshold) * first_contact, dim=1)
# no reward for zero command
reward *= torch.norm(env.command_manager.get_command(command_name), dim=1) > 0.1
reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
return reward
def feet_air_time_positive_biped(env, command_name: str, threshold: float, sensor_cfg: SceneEntityCfg) -> torch.Tensor:
"""Reward long steps taken by the feet for bipeds.
This function rewards the agent for taking steps up to a specified threshold and also keep one foot at
a time in the air.
If the commands are small (i.e. the agent is not supposed to take a step), then the reward is zero.
"""
contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name]
# compute the reward
air_time = contact_sensor.data.current_air_time[:, sensor_cfg.body_ids]
contact_time = contact_sensor.data.current_contact_time[:, sensor_cfg.body_ids]
in_contact = contact_time > 0.0
in_mode_time = torch.where(in_contact, contact_time, air_time)
single_stance = torch.sum(in_contact.int(), dim=1) == 1
reward = torch.min(torch.where(single_stance.unsqueeze(-1), in_mode_time, 0.0), dim=1)[0]
reward = torch.clamp(reward, max=threshold)
# no reward for zero command
reward *= torch.norm(env.command_manager.get_command(command_name), dim=1) > 0.1
reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
return reward
def feet_air_time_variance_penalty(env: ManagerBasedRLEnv, sensor_cfg: SceneEntityCfg) -> torch.Tensor:
"""Penalize variance in the amount of time each foot spends in the air/on the ground relative to each other"""
# extract the used quantities (to enable type-hinting)
contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name]
# compute the reward
last_air_time = contact_sensor.data.last_air_time[:, sensor_cfg.body_ids]
last_contact_time = contact_sensor.data.last_contact_time[:, sensor_cfg.body_ids]
reward = torch.var(torch.clip(last_air_time, max=0.5), dim=1) + torch.var(
torch.clip(last_contact_time, max=0.5), dim=1
)
reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
return reward
def feet_contact(
env: ManagerBasedRLEnv, command_name: str, expect_contact_num: int, sensor_cfg: SceneEntityCfg
) -> torch.Tensor:
"""Reward feet contact"""
# extract the used quantities (to enable type-hinting)
contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name]
# compute the reward
contact = contact_sensor.compute_first_contact(env.step_dt)[:, sensor_cfg.body_ids]
contact_num = torch.sum(contact, dim=1)
reward = (contact_num != expect_contact_num).float()
# no reward for zero command
reward *= torch.linalg.norm(env.command_manager.get_command(command_name), dim=1) > 0.1
reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
return reward
def feet_contact_without_cmd(env: ManagerBasedRLEnv, command_name: str, sensor_cfg: SceneEntityCfg) -> torch.Tensor:
"""Reward feet contact"""
# extract the used quantities (to enable type-hinting)
contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name]
# compute the reward
contact = contact_sensor.compute_first_contact(env.step_dt)[:, sensor_cfg.body_ids]
reward = torch.sum(contact, dim=-1).float()
reward *= torch.linalg.norm(env.command_manager.get_command(command_name), dim=1) < 0.1
reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
return reward
def feet_stumble(env: ManagerBasedRLEnv, sensor_cfg: SceneEntityCfg) -> torch.Tensor:
# extract the used quantities (to enable type-hinting)
contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name]
forces_z = torch.abs(contact_sensor.data.net_forces_w[:, sensor_cfg.body_ids, 2])
forces_xy = torch.linalg.norm(contact_sensor.data.net_forces_w[:, sensor_cfg.body_ids, :2], dim=2)
# Penalize feet hitting vertical surfaces
reward = torch.any(forces_xy > 4 * forces_z, dim=1).float()
reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
return reward
def feet_distance_y_exp(
env: ManagerBasedRLEnv, stance_width: float, std: float, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")
) -> torch.Tensor:
asset: RigidObject = env.scene[asset_cfg.name]
cur_footsteps_translated = asset.data.body_link_pos_w[:, asset_cfg.body_ids, :] - asset.data.root_link_pos_w[
:, :
].unsqueeze(1)
n_feet = len(asset_cfg.body_ids)
footsteps_in_body_frame = torch.zeros(env.num_envs, n_feet, 3, device=env.device)
for i in range(n_feet):
footsteps_in_body_frame[:, i, :] = math_utils.quat_apply(
math_utils.quat_conjugate(asset.data.root_link_quat_w), cur_footsteps_translated[:, i, :]
)
side_sign = torch.tensor(
[1.0 if i % 2 == 0 else -1.0 for i in range(n_feet)],
device=env.device,
)
stance_width_tensor = stance_width * torch.ones([env.num_envs, 1], device=env.device)
desired_ys = stance_width_tensor / 2 * side_sign.unsqueeze(0)
stance_diff = torch.square(desired_ys - footsteps_in_body_frame[:, :, 1])
reward = torch.exp(-torch.sum(stance_diff, dim=1) / (std**2))
reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
return reward
def feet_distance_xy_exp(
env: ManagerBasedRLEnv,
stance_width: float,
stance_length: float,
std: float,
asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
) -> torch.Tensor:
asset: RigidObject = env.scene[asset_cfg.name]
# Compute the current footstep positions relative to the root
cur_footsteps_translated = asset.data.body_link_pos_w[:, asset_cfg.body_ids, :] - asset.data.root_link_pos_w[
:, :
].unsqueeze(1)
footsteps_in_body_frame = torch.zeros(env.num_envs, 4, 3, device=env.device)
for i in range(4):
footsteps_in_body_frame[:, i, :] = math_utils.quat_apply(
math_utils.quat_conjugate(asset.data.root_link_quat_w), cur_footsteps_translated[:, i, :]
)
# Desired x and y positions for each foot
stance_width_tensor = stance_width * torch.ones([env.num_envs, 1], device=env.device)
stance_length_tensor = stance_length * torch.ones([env.num_envs, 1], device=env.device)
desired_xs = torch.cat(
[stance_length_tensor / 2, stance_length_tensor / 2, -stance_length_tensor / 2, -stance_length_tensor / 2],
dim=1,
)
desired_ys = torch.cat(
[stance_width_tensor / 2, -stance_width_tensor / 2, stance_width_tensor / 2, -stance_width_tensor / 2], dim=1
)
# Compute differences in x and y
stance_diff_x = torch.square(desired_xs - footsteps_in_body_frame[:, :, 0])
stance_diff_y = torch.square(desired_ys - footsteps_in_body_frame[:, :, 1])
# Combine x and y differences and compute the exponential penalty
stance_diff = stance_diff_x + stance_diff_y
reward = torch.exp(-torch.sum(stance_diff, dim=1) / std**2)
reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
return reward
def feet_height(
env: ManagerBasedRLEnv,
command_name: str,
asset_cfg: SceneEntityCfg,
target_height: float,
tanh_mult: float,
) -> torch.Tensor:
"""Reward the swinging feet for clearing a specified height off the ground"""
asset: RigidObject = env.scene[asset_cfg.name]
foot_z_target_error = torch.square(asset.data.body_pos_w[:, asset_cfg.body_ids, 2] - target_height)
foot_velocity_tanh = torch.tanh(
tanh_mult * torch.linalg.norm(asset.data.body_lin_vel_w[:, asset_cfg.body_ids, :2], dim=2)
)
reward = torch.sum(foot_z_target_error * foot_velocity_tanh, dim=1)
# no reward for zero command
reward *= torch.linalg.norm(env.command_manager.get_command(command_name), dim=1) > 0.1
reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
return reward
def feet_height_body(
env: ManagerBasedRLEnv,
command_name: str,
asset_cfg: SceneEntityCfg,
target_height: float,
tanh_mult: float,
) -> torch.Tensor:
"""Reward the swinging feet for clearing a specified height off the ground"""
asset: RigidObject = env.scene[asset_cfg.name]
cur_footpos_translated = asset.data.body_pos_w[:, asset_cfg.body_ids, :] - asset.data.root_pos_w[:, :].unsqueeze(1)
footpos_in_body_frame = torch.zeros(env.num_envs, len(asset_cfg.body_ids), 3, device=env.device)
cur_footvel_translated = asset.data.body_lin_vel_w[:, asset_cfg.body_ids, :] - asset.data.root_lin_vel_w[
:, :
].unsqueeze(1)
footvel_in_body_frame = torch.zeros(env.num_envs, len(asset_cfg.body_ids), 3, device=env.device)
for i in range(len(asset_cfg.body_ids)):
footpos_in_body_frame[:, i, :] = math_utils.quat_apply_inverse(
asset.data.root_quat_w, cur_footpos_translated[:, i, :]
)
footvel_in_body_frame[:, i, :] = math_utils.quat_apply_inverse(
asset.data.root_quat_w, cur_footvel_translated[:, i, :]
)
foot_z_target_error = torch.square(footpos_in_body_frame[:, :, 2] - target_height).view(env.num_envs, -1)
foot_velocity_tanh = torch.tanh(tanh_mult * torch.norm(footvel_in_body_frame[:, :, :2], dim=2))
reward = torch.sum(foot_z_target_error * foot_velocity_tanh, dim=1)
reward *= torch.linalg.norm(env.command_manager.get_command(command_name), dim=1) > 0.1
reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
return reward
def feet_slide(
env: ManagerBasedRLEnv, sensor_cfg: SceneEntityCfg, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")
) -> torch.Tensor:
"""Penalize feet sliding.
This function penalizes the agent for sliding its feet on the ground. The reward is computed as the
norm of the linear velocity of the feet multiplied by a binary contact sensor. This ensures that the
agent is penalized only when the feet are in contact with the ground.
"""
# Penalize feet sliding
contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name]
contacts = contact_sensor.data.net_forces_w_history[:, :, sensor_cfg.body_ids, :].norm(dim=-1).max(dim=1)[0] > 1.0
asset: RigidObject = env.scene[asset_cfg.name]
# feet_vel = asset.data.body_lin_vel_w[:, asset_cfg.body_ids, :2]
# reward = torch.sum(feet_vel.norm(dim=-1) * contacts, dim=1)
cur_footvel_translated = asset.data.body_lin_vel_w[:, asset_cfg.body_ids, :] - asset.data.root_lin_vel_w[
:, :
].unsqueeze(1)
footvel_in_body_frame = torch.zeros(env.num_envs, len(asset_cfg.body_ids), 3, device=env.device)
for i in range(len(asset_cfg.body_ids)):
footvel_in_body_frame[:, i, :] = math_utils.quat_apply_inverse(
asset.data.root_quat_w, cur_footvel_translated[:, i, :]
)
foot_leteral_vel = torch.sqrt(torch.sum(torch.square(footvel_in_body_frame[:, :, :2]), dim=2)).view(
env.num_envs, -1
)
reward = torch.sum(foot_leteral_vel * contacts, dim=1)
reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
return reward
# def smoothness_1(env: ManagerBasedRLEnv) -> torch.Tensor:
# # Penalize changes in actions
# diff = torch.square(env.action_manager.action - env.action_manager.prev_action)
# diff = diff * (env.action_manager.prev_action[:, :] != 0) # ignore first step
# return torch.sum(diff, dim=1)
# def smoothness_2(env: ManagerBasedRLEnv) -> torch.Tensor:
# # Penalize changes in actions
# diff = torch.square(env.action_manager.action - 2 * env.action_manager.prev_action + env.action_manager.prev_prev_action)
# diff = diff * (env.action_manager.prev_action[:, :] != 0) # ignore first step
# diff = diff * (env.action_manager.prev_prev_action[:, :] != 0) # ignore second step
# return torch.sum(diff, dim=1)
def action_smoothness_l2(env: ManagerBasedRLEnv) -> torch.Tensor:
# Penalize changes in actions
diff = torch.square(env.action_manager.action - 2 * env.action_manager.prev_action + env.action_manager.prev_prev_action)
diff = diff * (env.action_manager.prev_action[:, :] != 0) # ignore first step
diff = diff * (env.action_manager.prev_prev_action[:, :] != 0) # ignore second step
return torch.sum(diff, dim=1)
def upward(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
"""Penalize z-axis base linear velocity using L2 squared kernel."""
# extract the used quantities (to enable type-hinting)
asset: RigidObject = env.scene[asset_cfg.name]
reward = torch.square(1 - asset.data.projected_gravity_b[:, 2])
return reward
def base_height_l2(
env: ManagerBasedRLEnv,
target_height: float,
asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
sensor_cfg: SceneEntityCfg | None = None,
) -> torch.Tensor:
"""Penalize asset height from its target using L2 squared kernel.
Note:
For flat terrain, target height is in the world frame. For rough terrain,
sensor readings can adjust the target height to account for the terrain.
"""
# extract the used quantities (to enable type-hinting)
asset: RigidObject = env.scene[asset_cfg.name]
if sensor_cfg is not None:
sensor: RayCaster = env.scene[sensor_cfg.name]
# Adjust the target height using the sensor data
ray_hits = sensor.data.ray_hits_w[..., 2]
if torch.isnan(ray_hits).any() or torch.isinf(ray_hits).any() or torch.max(torch.abs(ray_hits)) > 1e6:
adjusted_target_height = asset.data.root_link_pos_w[:, 2]
else:
adjusted_target_height = target_height + torch.mean(ray_hits, dim=1)
else:
# Use the provided target height directly for flat terrain
adjusted_target_height = target_height
# Compute the L2 squared penalty
reward = torch.square(asset.data.root_pos_w[:, 2] - adjusted_target_height)
reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
return reward
def lin_vel_z_l2(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
"""Penalize z-axis base linear velocity using L2 squared kernel."""
# extract the used quantities (to enable type-hinting)
asset: RigidObject = env.scene[asset_cfg.name]
reward = torch.square(asset.data.root_lin_vel_b[:, 2])
reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
return reward
def ang_vel_xy_l2(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
"""Penalize xy-axis base angular velocity using L2 squared kernel."""
# extract the used quantities (to enable type-hinting)
asset: RigidObject = env.scene[asset_cfg.name]
reward = torch.sum(torch.square(asset.data.root_ang_vel_b[:, :2]), dim=1)
reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
return reward
def undesired_contacts(env: ManagerBasedRLEnv, threshold: float, sensor_cfg: SceneEntityCfg) -> torch.Tensor:
"""Penalize undesired contacts as the number of violations that are above a threshold."""
# extract the used quantities (to enable type-hinting)
contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name]
# check if contact force is above threshold
net_contact_forces = contact_sensor.data.net_forces_w_history
is_contact = torch.max(torch.norm(net_contact_forces[:, :, sensor_cfg.body_ids], dim=-1), dim=1)[0] > threshold
# sum over contacts for each environment
reward = torch.sum(is_contact, dim=1).float()
reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
return reward
def flat_orientation_l2(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
"""Penalize non-flat base orientation using L2 squared kernel.
This is computed by penalizing the xy-components of the projected gravity vector.
"""
# extract the used quantities (to enable type-hinting)
asset: RigidObject = env.scene[asset_cfg.name]
reward = torch.sum(torch.square(asset.data.projected_gravity_b[:, :2]), dim=1)
reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
return reward
def hip_pos_penalty(
env: ManagerBasedRLEnv,
command_name: str,
asset_cfg: SceneEntityCfg,
stand_still_scale: float,
command_threshold: float,
) -> torch.Tensor:
"""Penalize joint position error from default on the articulation."""
# extract the used quantities (to enable type-hinting)
asset: Articulation = env.scene[asset_cfg.name]
command = env.command_manager.get_command(command_name)[:, [1, 2]]
cmd_large = torch.any(torch.abs(command) > command_threshold, dim=1)
running_reward = torch.linalg.norm(
(asset.data.joint_pos[:, asset_cfg.joint_ids] - asset.data.default_joint_pos[:, asset_cfg.joint_ids]), dim=1
)
reward = torch.where(
cmd_large,
running_reward,
stand_still_scale * running_reward
)
return reward

126
source/robot_lab/robot_lab/tasks/go2/mdp/utils.py Normal file
View File

@@ -0,0 +1,126 @@
# Copyright (c) 2024-2025 Ziqi Fan
# SPDX-License-Identifier: Apache-2.0
"""Utility functions for terrain-aware operations."""
from __future__ import annotations
import torch
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from isaaclab.envs import ManagerBasedEnv
def _get_terrain_column_range(terrain_cfg, terrain_name: str, device) -> tuple[int, int] | None:
"""Helper function to calculate column range for a terrain type.
Args:
terrain_cfg: The terrain generator configuration.
terrain_name: Name of the terrain.
device: Torch device.
Returns:
Tuple of (col_start, col_end) or None if terrain not found.
"""
if terrain_cfg.sub_terrains is None or terrain_name not in terrain_cfg.sub_terrains:
return None
sub_terrain_names = list(terrain_cfg.sub_terrains.keys())
proportions = torch.tensor([sub_cfg.proportion for sub_cfg in terrain_cfg.sub_terrains.values()], device=device)
proportions = proportions / proportions.sum()
cumsum_props = torch.cumsum(proportions, dim=0)
terrain_idx = sub_terrain_names.index(terrain_name)
# Use round() instead of int() to properly allocate columns
col_start = round((0.0 if terrain_idx == 0 else cumsum_props[terrain_idx - 1].item()) * terrain_cfg.num_cols)
col_end = round(cumsum_props[terrain_idx].item() * terrain_cfg.num_cols)
return (col_start, col_end)
def is_env_assigned_to_terrain(env: ManagerBasedEnv, terrain_name: str) -> torch.Tensor:
"""Check which environments are initially assigned to the specified terrain type.
Each environment is assigned to a specific terrain cell at initialization.
This function returns a mask indicating which environments were assigned to the given terrain type.
Args:
env: The environment instance.
terrain_name: Name of the terrain to check (e.g., "pits", "stairs").
Returns:
Boolean tensor of shape (num_envs,) where True means the environment is assigned to this terrain.
"""
# Check if terrain and terrain generator are available
terrain = getattr(env.scene, "terrain", None)
if terrain is None or not hasattr(terrain, "terrain_types"):
return torch.zeros(env.num_envs, dtype=torch.bool, device=env.device)
if terrain.cfg.terrain_type != "generator" or terrain.cfg.terrain_generator is None:
return torch.zeros(env.num_envs, dtype=torch.bool, device=env.device)
terrain_cfg = terrain.cfg.terrain_generator
col_range = _get_terrain_column_range(terrain_cfg, terrain_name, env.device)
if col_range is None:
return torch.zeros(env.num_envs, dtype=torch.bool, device=env.device)
col_start, col_end = col_range
# terrain_types directly stores column indices, so just check if they're in range
return (terrain.terrain_types >= col_start) & (terrain.terrain_types < col_end)
def is_robot_on_terrain(env: ManagerBasedEnv, terrain_name: str, asset_name: str = "robot") -> torch.Tensor:
"""Check which robots are currently standing on the specified terrain type.
This function calculates which terrain grid cell each robot is on based on its world position,
then checks if that cell's terrain type matches the specified terrain.
Args:
env: The environment instance.
terrain_name: Name of the terrain to check (e.g., "pits", "stairs").
asset_name: Name of the robot asset. Defaults to "robot".
Returns:
Boolean tensor of shape (num_envs,) where True means the robot is currently on this terrain.
"""
# Check if terrain and terrain generator are available
terrain = getattr(env.scene, "terrain", None)
if terrain is None or not hasattr(terrain, "terrain_types"):
return torch.zeros(env.num_envs, dtype=torch.bool, device=env.device)
if terrain.cfg.terrain_type != "generator" or terrain.cfg.terrain_generator is None:
return torch.zeros(env.num_envs, dtype=torch.bool, device=env.device)
terrain_cfg = terrain.cfg.terrain_generator
col_range = _get_terrain_column_range(terrain_cfg, terrain_name, env.device)
if col_range is None:
return torch.zeros(env.num_envs, dtype=torch.bool, device=env.device)
col_start, col_end = col_range
# Get robot positions in world frame
asset = env.scene[asset_name]
robot_pos_w = asset.data.root_pos_w[:, :2] # [num_envs, 2] (x, y)
# Get terrain grid information
terrain_origins = terrain.terrain_origins # [num_rows, num_cols, 3]
num_rows, num_cols, _ = terrain_origins.shape
# Use terrain_origins to directly compute which cell each robot is in
# terrain_origins[r, c, :2] is the center of cell (r, c)
# We need to find the closest terrain origin for each robot
# Reshape terrain_origins for distance calculation
terrain_origins_2d = terrain_origins[:, :, :2].reshape(num_rows * num_cols, 2) # [num_rows*num_cols, 2]
# Calculate distances from each robot to all terrain origins
distances = torch.cdist(robot_pos_w, terrain_origins_2d) # [num_envs, num_rows*num_cols]
# Find the closest terrain origin for each robot
closest_flat_idx = torch.argmin(distances, dim=1) # [num_envs]
# Convert flat index to column index
# flat_idx = row * num_cols + col
col_idx = closest_flat_idx % num_cols # [num_envs]
# Check if the robot's current terrain column is in the specified terrain's range
return (col_idx >= col_start) & (col_idx < col_end)

86
source/robot_lab/robot_lab/tasks/go2/rsl_rl_cfg.py Normal file
View File

@@ -0,0 +1,86 @@
from isaaclab.utils import configclass
from isaaclab_rl.rsl_rl import RslRlOnPolicyRunnerCfg, RslRlPpoActorCriticCfg, RslRlPpoAlgorithmCfg
@configclass
class PPORunnerCfg(RslRlOnPolicyRunnerCfg):
num_steps_per_env = 24
max_iterations = 150000
save_interval = 500
experiment_name = "go2_rough"
policy = RslRlPpoActorCriticCfg(
init_noise_std=1.0,
actor_obs_normalization=False,
critic_obs_normalization=False,
actor_hidden_dims=[512, 256, 128],
critic_hidden_dims=[512, 256, 128],
activation="elu",
)
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,
)
@configclass
class RslRlMoeCtsActorCriticCfg(RslRlPpoActorCriticCfg):
class_name = "ActorCriticMoECTS"
init_noise_std = 1.0
expert_num = 8 # number of experts in the student model
latent_dim = 32
norm_type = 'l2norm' # normalization type for encoders: l2norm, simnorm
teacher_encoder_hidden_dims = [512, 256, 128]
student_encoder_hidden_dims = [512, 256, 128]
actor_hidden_dims=[512, 256, 128]
critic_hidden_dims=[512, 256, 128]
activation="elu"
actor_obs_normalization = False
critic_obs_normalization = False
@configclass
class RslRlMoeCtsAlgorithmCfg(RslRlPpoAlgorithmCfg):
class_name = "MoECTS"
value_loss_coef = 1.0
load_balance_coef = 0.01 # coefficient for load balance loss
use_clipped_value_loss = True
clip_param = 0.2
entropy_coef = 0.01
num_learning_epochs = 5
num_mini_batches = 4
learning_rate = 1e-3
student_encoder_learning_rate = 1e-3
schedule = "adaptive"
gamma = 0.99
lam = 0.95
betas = (0.9, 0.999)
weight_decay = 0.0
desired_kl = 0.01
max_grad_norm = 1.0
teacher_env_ratio = 0.75 # percentage of envs assigned to teacher
@configclass
class MoECTSRunnerCfg(RslRlOnPolicyRunnerCfg):
experiment_name = "go2_moe_cts"
class_name = "OnPolicyRunnerCTS"
num_steps_per_env = 24
max_iterations = 150000
save_interval = 500
policy = RslRlMoeCtsActorCriticCfg()
algorithm = RslRlMoeCtsAlgorithmCfg()
# concat elu inspired by concat relu from https://arxiv.org/pdf/2303.07507
@configclass
class MoECTSCatELURunnerCfg(MoECTSRunnerCfg):
def __post_init__(self):
super().__post_init__()
self.policy.activation = 'cat_elu'

49
source/robot_lab/robot_lab/ui_extension_example.py Normal file
View File

@@ -0,0 +1,49 @@
# Copyright (c) 2024-2025 Ziqi Fan
# SPDX-License-Identifier: Apache-2.0
# Copyright (c) 2022-2025, The Isaac Lab Project Developers.
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause
import omni.ext
# Functions and vars are available to other extension as usual in python: `example.python_ext.some_public_function(x)`
def some_public_function(x: int):
print("[robot_lab] some_public_function was called with x: ", x)
return x**x
# Any class derived from `omni.ext.IExt` in top level module (defined in `python.modules` of `extension.toml`) will be
# instantiated when extension gets enabled and `on_startup(ext_id)` will be called. Later when extension gets disabled
# on_shutdown() is called.
class ExampleExtension(omni.ext.IExt):
# ext_id is current extension id. It can be used with extension manager to query additional information, like where
# this extension is located on filesystem.
def on_startup(self, ext_id):
print("[robot_lab] startup")
self._count = 0
self._window = omni.ui.Window("My Window", width=300, height=300)
with self._window.frame:
with omni.ui.VStack():
label = omni.ui.Label("")
def on_click():
self._count += 1
label.text = f"count: {self._count}"
def on_reset():
self._count = 0
label.text = "empty"
on_reset()
with omni.ui.HStack():
omni.ui.Button("Add", clicked_fn=on_click)
omni.ui.Button("Reset", clicked_fn=on_reset)
def on_shutdown(self):
print("[robot_lab] shutdown")

53
source/robot_lab/setup.py Normal file
View File

@@ -0,0 +1,53 @@
# Copyright (c) 2024-2025 Ziqi Fan
# SPDX-License-Identifier: Apache-2.0
"""Installation script for the 'robot_lab' python package."""
import os
import toml
from setuptools import setup
# Obtain the extension data from the extension.toml file
EXTENSION_PATH = os.path.dirname(os.path.realpath(__file__))
# Read the extension.toml file
EXTENSION_TOML_DATA = toml.load(os.path.join(EXTENSION_PATH, "config", "extension.toml"))
# Minimum dependencies required prior to installation
INSTALL_REQUIRES = [
# base
"psutil",
"colorama",
"xacrodoc",
# amp
"numpy",
"pandas",
"pinocchio",
# rl
"cusrl[all]",
]
# Installation operation
setup(
name="robot_lab",
packages=["robot_lab"],
author=EXTENSION_TOML_DATA["package"]["author"],
maintainer=EXTENSION_TOML_DATA["package"]["maintainer"],
url=EXTENSION_TOML_DATA["package"]["repository"],
version=EXTENSION_TOML_DATA["package"]["version"],
description=EXTENSION_TOML_DATA["package"]["description"],
keywords=EXTENSION_TOML_DATA["package"]["keywords"],
install_requires=INSTALL_REQUIRES,
license="Apache License 2.0",
include_package_data=True,
python_requires=">=3.10",
classifiers=[
"Natural Language :: English",
"Programming Language :: Python :: 3.10",
"Programming Language :: Python :: 3.11",
"Isaac Sim :: 4.5.0",
"Isaac Sim :: 5.0.0",
"Isaac Sim :: 5.1.0",
],
zip_safe=False,
)