The goal of this project is to provide a simple, minimal working example of sim2real transfer for the Unitree Go2 robot, including training, visualization, and deployment to real hardware.
showcase.mp4
Currently, only the joystick environment from MuJoCo Playground is implemented. Note, however, that the official MuJoCo Playground does not support the Unitree Go2. Therefore, a fork of the Playground is used in this project.
The Policy Controller is written in Python and can be deployed either on the robot's internal Nvidia Jetson module or run directly on a workstation. For latency reasons, deploying on the Jetson module is recommended.
- Switch easily between simulation and real hardware
- Simple API to add new policies
- Capture API, nice way to store and visualize data
- Security features to prevent hardware damage
- Joystick control after policy phase (i. e. without restarting the robot)
- Joystick controllable policy
- Supports Vicon and Go2 Odometry for base estimation
If you plan to use this code for your own project, please be aware of the following: While this implementation may work for the provided policy, Python is generally not ideal for real-time control tasks. To achieve the best possible results, ensure that your policy adapter is as efficient as possible. Whenever feasible, use JAX-compiled functions to minimize latency and avoid jitter in control loops.
General Note on high frequency applications: If you experience jitter or if even the provided example policy immediately crashes in simulation, try disabling all energy-saving modes, increase the process priority, and ensure that your machine is not under heavy load. Running the simulation or control code on an overloaded system can cause significant lag and instability.
When operating in low state mode, the Go2 sport mode state is unavailable. As a result, your policy cannot access the robot’s linear velocity or position via the sport mode interface. If your policy requires position or linear velocity inputs, use either the simulation environment (which provides perfect data) or rely on Vicon / Go2 Odometry when running on real hardware.
# Clone the repository including submodules
git clone --recurse-submodules https://github.com/DerSimi/unitree_go2_sim2real
cd unitree_go2_sim2realBefore continuing, ensure that CycloneDDS is installed and properly sourced:
cd ~
git clone https://github.com/eclipse-cyclonedds/cyclonedds -b releases/0.10.x
cd cyclonedds && mkdir build install && cd build
cmake .. -DCMAKE_INSTALL_PREFIX=../install
cmake --build . --target install
# Add this to your .bashrc or .zshrc
export CYCLONEDDS_HOME="$HOME/cyclonedds/install"If you want to render videos using the demo code, you need to install ffmpeg.
sudo apt install ffmpegInstall the Unitree SDK2:
- Clone the repository: unitree_sdk2
- To ensure compatibility with
unitree_mujoco, check out the following commit:3a4680ae9b00df59e60f7e63cfb0fcc432a9d08d - Follow the instructions in the repository.
Install MuJoCo, version 3.2.7. Note: According to unitree_mujoco, you must build MuJoCo from source!
For unitree_mujoco, use the unitree_mujoco submodule and build it as described in its repository.
uv venv --python 3.12
source .venv/bin/activate
uv pip install -U "jax[cuda12]"
cd mujoco_playground
uv pip install -e .
cd ..
uv pip install -r requirements.txt
cd unitree_sdk2_python
uv pip install -e .
cd ..Now install the sim2real package itself:
uv pip install -e .Before using sim2real, you need to train a policy. This project contains multiple policies to choose from, the simplest is the walker_policy in the policies directory.
./policies/walker_policy/train.shAfterwards, open policies/walker_policy/walker_policy.py and specify the correct checkpoint for the trained policy.
You can also experiment with test/interference.py to use the policy outside of the Unitree SDK bridge.
Note: In
train.sh, you can choose whether to use Weights & Biases (wandb). If you do, make sure to enter the correct project and organization name inpolicies/train_policy.py.
After training, you can deploy the policy to real hardware.
There are two options for deployment:
This is not the real robot, but a simulated environment that closely mimics the robot. It allows you to test your policies before deploying it on real hardware.
First, start the simulation environment:
./mujoco.shThen run for example the walker policy:
python policies/walker_policy/walker_policy.pyTo deploy on the real robot, you need to determine the network interface that connects to the robot. If you are working from a workstation, see the instructions here.
To launch the policy, run the following command and replace <network interface> with the correct interface name:
python sim2real/run_policy.py <network interface>This project provides a simple yet powerful way to add new policies to sim2real. Each policy must implement the PolicyAdapter interface:
class PolicyAdapter(ABC):
def __init__(
self,
dt_sim: float,
dt_ctrl: float,
default_pose: jax.Array,
kp: float,
kd: float,
high_state_type: HighStateType = HighStateType.NONE,
**kwargs,
):
"""
dt_sim: Simulation time step, in seconds
dt_ctrl: Control time step, in seconds. Note, that dt_sim <= dt_ctrl
default_pose: The default starting pose the policy is expecting. Note: Here the order used in simulaton is expected!
kp: Proportional constant for PD controller
kd: Derivative constant for PD controller
high_state_type: Specify the base state estimation when running on real hardware (options: VICON or GO2_ODOMETRY). In simulation, if set, the base state from MuJoCo will be used. NONE means no base estimation is available.
kwargs: When using vicon, the system expects vicon_ip and vicon_obj_name as input.
"""
super().__init__()
assert (
dt_sim <= dt_ctrl
), "Simulation time step must be smaller or equal to control time step"
self.dt_sim = dt_sim
self.dt_ctrl = dt_ctrl
self.ctrl_per_sim = int(dt_ctrl / dt_sim)
self.default_pose = default_pose
self.kp = kp
self.kd = kd
self.high_state_type = high_state_type
self.kwargs = kwargs
# will be overriden in the policy controller
# True when policy runns in simulation, false on real hardware
self.simulation = False
# Time since the policy started (in seconds and float)
self.start_time = 0
@abstractmethod
def next_command(low_state: LowState_, high_state: HighState) -> jax.Array:
"""Returns the current command as a jax.Array."""
pass
@abstractmethod
def get_observation(
self,
low_state: LowState_,
high_state: HighState,
last_raw_action: jax.Array,
command: jax.Array,
) -> jax.Array:
"""Builds the observation for the policy given the low-level state,
last raw action, and command.
Args:
low_state: The low-level state.
high_state: Either data from sport mode state, vicon or go2odometry.
last_raw_action: The last raw action taken.
command: Next command
Returns:
The policy observation
"""
pass
@abstractmethod
def format_action(self, action: jax.Array, env_cfg): # Ignore type
"""Depending on the used model, the raw action can not be used directly and needs preprocessing"""
passConsider implementations in policies/walker_policy as examples. You also see how to use the vicon system.
As a workflow, I recommend: Test your policies locally first using Unitree MuJoCo!
If you are working with policies that may result in high pitch angles, please note that the controller includes a safety mechanism to immediately stop the policy if the robot is at risk of crashing:
class DefaultEmergencyStop(EmergencyStop):
def kill_system(self, roll, pitch, low_state, high_state):
# 0.4 rad = 22,9183°
return abs(pitch) >= 0.4 or abs(roll) >= 0.4If you need different emergency stop behavior, implement your own EmergencyStop interface and pass it to the policy controller.
class PolicyController:
def __init__(
self,
policy: Policy,
simulation: bool,
policy_adapter: PolicyAdapter,
emergency_stop: EmergencyStop = DefaultEmergencyStop(),
capture: Capture = None,
):
"""
policy: The expected policy, see example in walker_policy.py
simulation: True when the policy is running on the simulation, false when on real hardware
policy_adapter: The policy adapter implements all methods to communicate with the policy bridge
emergency_stop: Control the precise moment when your policy should terminate immediately.
capture: Optional field. You can define a own capture method or use the default ones provided in capture/
"""
...
)As some workflows require capturing data, this project includes a simple interface:
class Capture(ABC):
"""
Abstract base class for capturing policy and control data during policy operation.
Be careful: Inefficient code in `policy_sample` and `send_ctrl` can introduce jitter
in the control loop. Keep implementations lightweight and non-blocking.
"""
def __init__(self):
# Overwritten by policy controller, ignore
self.adapter = None
super().__init__()
@abstractmethod
def record_observation_input(
self,
low_state: LowState_,
high_state: HighState,
last_raw_action: jax.Array,
command: jax.Array,
):
"""
Triggered each time a new action is sampled from the policy.
Expected frequency is dt_ctrl defined in the policy adapter.
low_state: Unitree low state containing imu, motor information, etc.
high_state: Sim2real high state, either built from unitree sport state mode or other means
last_raw_action: The last raw action taken.
command: Next command
"""
pass
@abstractmethod
def record_observation_output(
self,
observation: jax.Array,
raw_act: jax.Array,
fmt_act: jax.Array,
policy_ctrl: jax.Array,
):
"""
Triggered each time a new action is sampled from the policy.
Expected frequency is dt_ctrl defined in the policy adapter.
observation: Observation used by the policy
raw_act: The raw action the policy predicts
fmt_act: Formatted action (i. e. with right action scale, ...)
ctrl: Policy ctrl input. This is never send to the robot directly!
Note, ctrl is always in the order as the robot expects it!
"""
pass
@abstractmethod
def record_ctrl(self, ctrl: jax.Array):
"""
Triggered when an actual control (joint goal positions) is sent to the robot.
Expected frequency is dt_sim defined in the policy adapter.
ctrl: Control in order the robot expects!
"""
pass
@abstractmethod
def store(self):
"""
Triggered when the policy is killed.
"""
pass
@abstractmethod
def display(self):
"""
Method for displaying graphs, called shortly before termination.
Exceptions are supressed by default.
"""
passThree useful default implementations are available:
ctrl_capture: Captures all control commands sent to the robot and saves them to a file. You can find an example script for processing this data in thetestfolder.delay_capture: Plots the delay between observation function calls, making it easier to identify jitter issues in the control loop.velocity_capture: Plots the robot linear velocity, only works when the high state mode is enabled.
You can also stack multiple capture implementations before passing them to the policy controller:
capture = StackedCapture([CtrlCapture(), VelocityCapture(), DelayCapture()])There is a default implementation in policy_bridge/storage/ctrl_capture, only capturing controls sent to the robot.
You can terminate the policy at any time by pressing Ctrl + C in the terminal. The robot will immediately disengage from the policy and transition to a safe sitting position.
If the robot has not crashed, you can resume manual control using the joystick.