Training a locomotion policy with RL#

This tutorial trains a walking policy for the Unitree Go2 quadruped and rolls it out in the viewer. It is a deliberately minimal pipeline: a gym-style environment, PPO from an external RL library, and a short evaluation loop. The reward terms are simplified to get a policy walking quickly, so treat it as a starting point rather than a production locomotion stack.

The three scripts that make up the example are the source of truth for the complete code:

go2_env.py: the environment: scene, control, observations, rewards, and resets.
go2_train.py: hyperparameters and the training entry point.
go2_eval.py: loading a checkpoint and running the policy.

The design borrows its structure from Legged Gym.

Prerequisites#

Training uses PPO from rsl-rl, an external library, and logs to TensorBoard. Install both alongside Genesis World:

pip install tensorboard "rsl-rl-lib>=5.0.0"

The scripts require rsl-rl-lib>=5.0.0 and raise an ImportError at startup if an older version is present, because the example relies on the observation-group and TensorDict conventions introduced in that release.

The environment#

Go2Env is a standard RL environment: it owns a Genesis World scene built across many environments, and it exposes reset() and step(actions). Training runs thousands of copies of the robot in parallel on one GPU, so every buffer carries a leading n_envs dimension and the step logic stays vectorized on the device.

The constructor builds the scene once for all environments:

self.scene = gs.Scene(
    sim_options=gs.options.SimOptions(
        dt=self.dt,  # 0.02 s -> 50 Hz control, matching the real robot
        substeps=2,
    ),
    rigid_options=gs.options.RigidOptions(
        enable_self_collision=False,
        max_collision_pairs=20,  # a walking Go2 rarely exceeds this; a low cap saves memory
    ),
    ...
)

self.scene.build(n_envs=num_envs)

n_envs is set at build time, not at construction, and everything downstream is sized to match it.

Control#

The robot is driven by a PD controller in joint-position space. The 12 leg joints are located by name, and their gains are set once after the build:

self.robot.set_dofs_kp([self.env_cfg["kp"]] * self.num_actions, self.motors_dof_idx)  # kp = 20.0
self.robot.set_dofs_kv([self.env_cfg["kd"]] * self.num_actions, self.motors_dof_idx)  # kd = 0.5

The policy does not output torques or absolute joint angles. It outputs a residual around the robot’s default standing pose, which step scales and adds to that pose to form a position target:

target_dof_pos = exec_actions * self.env_cfg["action_scale"] + self.default_dof_pos  # action_scale = 0.25, rad
self.robot.control_dofs_position(target_dof_pos[:, self.actions_dof_idx], slice(6, 18))

Actuating around a default pose keeps early exploration close to a stable stance, which is what lets the policy find a gait in a few hundred iterations. The environment also feeds the previous step’s action to the controller (simulate_action_latency) to reproduce the one-step (~20 ms) actuation delay a real Go2 shows, narrowing the sim-to-real gap.

Observations#

step reads the robot state back after scene.step(), converting base linear and angular velocity into the robot’s own frame, then assembles the policy input:

self.obs_buf = torch.concatenate(
    (
        self.base_ang_vel * self.obs_scales["ang_vel"],           # 3, base-frame angular velocity (rad/s)
        self.projected_gravity,                                   # 3, gravity direction in base frame
        self.commands * self.commands_scale,                      # 3, target lin_vel_x, lin_vel_y, ang_vel_yaw
        (self.dof_pos - self.default_dof_pos) * self.obs_scales["dof_pos"],  # 12, joint-position residual (rad)
        self.dof_vel * self.obs_scales["dof_vel"],                # 12, joint velocity (rad/s)
        self.actions,                                             # 12, previous action
    ),
    dim=-1,
)

The observation is a 45-dimensional vector, shape (n_envs, 45). Each term is multiplied by a fixed scale so the components sit on comparable magnitudes before they reach the network. get_observations wraps the buffer in a TensorDict under the key "policy", which is the format rsl-rl expects:

def get_observations(self):
    return TensorDict({"policy": self.obs_buf}, batch_size=[self.num_envs])

The command is the task: a per-environment target of forward velocity, lateral velocity, and yaw rate. Commands are resampled at random within configured ranges every resampling_time_s and on reset, so a single policy learns to follow a distribution of velocity commands rather than one fixed gait.

Rewards#

The total reward is a weighted sum of six terms, each a method on the environment. Two terms reward the task and four regularize the motion:

Reward term	Purpose	Scale
`tracking_lin_vel`	follow the commanded xy velocity	`1.0`
`tracking_ang_vel`	follow the commanded yaw rate	`0.2`
`lin_vel_z`	penalize vertical bouncing	`-1.0`
`base_height`	hold the base near its target height	`-50.0`
`action_rate`	penalize jerky action changes	`-0.005`
`similar_to_default`	stay near the default joint pose	`-0.1`

The two tracking terms use an exponential of the squared error, which rewards getting close without overweighting near-perfect tracking:

def _reward_tracking_lin_vel(self):
    # Tracking of linear velocity commands (xy axes)
    lin_vel_error = torch.sum(torch.square(self.commands[:, :2] - self.base_lin_vel[:, :2]), dim=1)
    return torch.exp(-lin_vel_error / self.reward_cfg["tracking_sigma"])

Each raw term is multiplied by its scale and by dt, so the weights read as per-second rates and stay independent of the control frequency. All reward configuration lives in get_cfgs in go2_train.py; tuning behavior means editing those scales, not the environment.

Resets#

step marks an environment for reset when its episode reaches episode_length_s, when the base roll or pitch exceeds 10 degrees (the robot has fallen), or when the physics solver flags a numerical error. Flagged environments are reset in place to the initial pose without interrupting the others, and a fresh command is drawn for each.

Training#

With the environment defined, go2_train.py hands it to rsl-rl’s OnPolicyRunner, which trains a PPO policy. The actor and critic are both three-layer MLPs ([512, 256, 128], ELU); the full hyperparameter set is in get_train_cfg. Genesis World is initialized on the GPU with performance_mode=True, which bakes static tensor shapes into the compiled kernels for faster stepping at the cost of recompiling when the scene changes, a worthwhile trade for a long training run:

gs.init(backend=gs.gpu, precision="32", logging_level="warning", seed=args.seed, performance_mode=True)

Start training from the repository root:

python examples/locomotion/go2_train.py

By default this runs 4096 environments in parallel for 101 iterations under the experiment name go2-walking, saving checkpoints and a cfgs.pkl (the exact configs used) into logs/go2-walking/. Override any of these with -B/--num_envs, --max_iterations, or -e/--exp_name.

Monitor progress with TensorBoard:

tensorboard --logdir logs

TensorBoard reward curves rising over training iterations for the Go2 walking policy

Evaluation#

go2_eval.py reloads the saved configs, rebuilds the environment with a single visualized instance, restores a checkpoint, and steps the policy in a loop:

runner.load(os.path.join(log_dir, f"model_{args.ckpt}.pt"))
policy = runner.get_inference_policy(device=gs.device)

obs_dict = env.reset()
with torch.no_grad():
    while True:
        actions = policy(obs_dict)
        obs_dict, rews, dones, infos = env.step(actions)

Evaluation initializes Genesis World on the CPU (gs.init(backend=gs.cpu)), which is sufficient for a single environment with the viewer open. Run it with the experiment name and checkpoint to load:

python examples/locomotion/go2_eval.py -e go2-walking --ckpt 100

The trained policy is a standard joint-position controller and can be deployed to a physical Go2.

Variations#

go2_backflip.py reuses the same Go2Env machinery for a dynamic backflip. It subclasses the environment to add a phase-based observation and loads a pre-trained TorchScript policy, showing how far the base environment stretches beyond flat walking:

python examples/locomotion/go2_backflip.py -e single