This repository is the starter code for the NYU Reinforcement Learning and Optimal Control project in which students train a Unitree Go2 walking policy in Isaac Lab starting from a minimal baseline and improve it via reward shaping and robustness strategies. Please read this README fully before starting and follow the exact workflow and naming rules below to ensure your runs integrate correctly with the cluster scripts and grading pipeline.
- Fork this repository and do not change the repository name in your fork.
- Your fork must be named rob6323_go2_project so cluster scripts and paths work without modification.
- GitHub Account: You must have a GitHub account to fork this repository and manage your code. If you do not have one, sign up here.
- Project Webpage: https://machines-in-motion.github.io/RL_class_go2_project/
- Project Tutorial: https://github.com/machines-in-motion/rob6323_go2_project/blob/master/tutorial/tutorial.md
- Connect to the NYU Greene HPC via SSH; if you are off-campus or not on NYU Wi‑Fi, you must connect through the NYU VPN before SSHing to Greene.
- The official instructions include example SSH config snippets and commands for greene.hpc.nyu.edu and dtn.hpc.nyu.edu as well as VPN and gateway options: https://sites.google.com/nyu.edu/nyu-hpc/accessing-hpc?authuser=0#h.7t97br4zzvip.
After logging into Greene, cd into your home directory (cd $HOME). You must clone your fork into $HOME only (not scratch or archive). This ensures subsequent scripts and paths resolve correctly on the cluster. Since this is a private repository, you need to authenticate with GitHub. You have two options:
The easiest way to avoid managing keys manually is to configure VS Code Remote SSH. If set up correctly, VS Code forwards your local credentials to the cluster.
- Follow the NYU HPC VS Code guide to set up the connection.
Tip: Once connected to Greene in VS Code, you can clone directly without using the terminal:
- Sign in to GitHub: Click the "Accounts" icon (user profile picture) in the bottom-left sidebar. If you aren't signed in, click "Sign in with GitHub" and follow the browser prompts to authorize VS Code.
- Clone the Repo: Open the Command Palette (
Ctrl+Shift+PorCmd+Shift+P), type Git: Clone, and select it.- Select Destination: When prompted, select your home directory (
/home/<netid>/) as the clone location.For more details, see the VS Code Version Control Documentation.
If you prefer using a standard terminal, you must generate a unique SSH key on the Greene cluster and add it to your GitHub account:
- Generate a key: Run the
ssh-keygencommand on Greene (follow the official GitHub documentation on generating a new SSH key). - Add the key to GitHub: Copy the output of your public key (e.g.,
cat ~/.ssh/id_ed25519.pub) and add it to your account settings (follow the GitHub documentation on adding a new SSH key).
Once authenticated, run the following commands. Replace <your-git-ssh-url> with the SSH URL of your fork (e.g., git@github.com:YOUR_USERNAME/rob6323_go2_project.git).
cd $HOME
git clone <your-git-ssh-url> rob6323_go2_project
Note: You must ensure the target directory is named exactly rob6323_go2_project. This ensures subsequent scripts and paths resolve correctly on the cluster.
- Enter the project directory and run the installer to set up required dependencies and cluster-side tooling.
cd $HOME/rob6323_go2_project
./install.sh
Do not skip this step, as it configures the environment expected by the training and evaluation scripts. It will launch a job in burst to set up things and clone the IsaacLab repo inside your greene storage. You must wait until the job in burst is complete before launching your first training. To check the progress of the job, you can run ssh burst "squeue -u $USER", and the job should disappear from there once it's completed. It takes around 30 minutes to complete.
You should see something similar to the screenshot below (captured from Greene):
In this output, the ST (state) column indicates the job status:
PD= pending in the queue (waiting for resources).CF= instance is being configured.R= job is running.
On burst, it is common for an instance to fail to configure; in that case, the provided scripts automatically relaunch the job when this happens, so you usually only need to wait until the job finishes successfully and no longer appears in squeue.
- In this project you'll only have to modify the two files below, which define the Isaac Lab task and its configuration (including PPO hyperparameters).
- source/rob6323_go2/rob6323_go2/tasks/direct/rob6323_go2/rob6323_go2_env.py
- source/rob6323_go2/rob6323_go2/tasks/direct/rob6323_go2/rob6323_go2_env_cfg.py PPO hyperparameters are defined in source/rob6323_go2/rob6323_go2/tasks/direct/rob6323_go2/agents/rsl_rl_ppo_cfg.py, but you shouldn't need to modify them.
- Option A (recommended): Use VS Code Remote SSH from your laptop to edit files on Greene; follow the NYU HPC VS Code guide and connect to a compute node as instructed (VPN required off‑campus) (https://sites.google.com/nyu.edu/nyu-hpc/training-support/general-hpc-topics/vs-code). If you set it correctly, it makes the login process easier, among other things, e.g., cloning a private repo.
- Option B: Edit directly on Greene using a terminal editor such as nano.
nano source/rob6323_go2/rob6323_go2/tasks/direct/rob6323_go2/rob6323_go2_env.py
- Option C: Develop locally on your machine, push to your fork, then pull changes on Greene within your $HOME/rob6323_go2_project clone.
Tip: Don't forget to regularly push your work to github
- From $HOME/rob6323_go2_project on Greene, submit a training job via the provided script.
cd "$HOME/rob6323_go2_project"
./train.sh
- Check job status with SLURM using squeue on the burst head node as shown below.
ssh burst "squeue -u $USER"
Be aware that jobs can be canceled and requeued by the scheduler or underlying provider policies when higher-priority work preempts your resources, which is normal behavior on shared clusters using preemptible partitions.
- When a job completes, logs are written under logs in your project clone on Greene in the structure logs/[job_id]/rsl_rl/go2_flat_direct/[date_time]/.
- Inside each run directory you will find a TensorBoard events file (events.out.tfevents...), neural network checkpoints (model_[epoch].pt), YAML files with the exact PPO and environment parameters, and a rollout video under videos/play/ that showcases the trained policy.
Use rsync to copy results from the cluster to your local machine. It is faster and can resume interrupted transfers. Run this on your machine (NOT on Greene):
rsync -avzP -e 'ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null' <netid>@dtn.hpc.nyu.edu:/home/<netid>/rob6323_go2_project/logs ./
Explanation of flags:
-a: Archive mode (preserves permissions, times, and recursive).-v: Verbose output.-z: Compresses data during transfer (faster over network).-P: Shows progress bar and allows resuming partial transfers.
You can inspect training metrics (reward curves, loss values, episode lengths) using TensorBoard. This requires installing it on your local machine.
-
Install TensorBoard: On your local computer (do NOT run this on Greene), install the package:
pip install tensorboard -
Launch the Server: Navigate to the folder where you downloaded your logs and start the server:
# Assuming you are in the directory containing the 'logs' folder tensorboard --logdir ./logs -
View Metrics: Open your browser to the URL shown (usually
http://localhost:6006/).
Burst storage is accessible only from a job running on burst, not from the burst login node. The provided scripts do not automatically synchronize error logs back to your home directory on Greene. However, you will need access to these logs to debug failed jobs. These error logs differ from the logs in the previous section.
The suggested way to inspect these logs is via the Open OnDemand web interface:
- Navigate to https://ood-burst-001.hpc.nyu.edu.
- Select Files > Home Directory from the top menu.
- You will see a list of files, including your
.errlog files. - Click on any
.errfile to view its content directly in the browser.
Important: Do not modify anything inside the
rob6323_go2_projectfolder on burst storage. This directory is managed by the job scripts, and manual changes may cause synchronization issues or job failures.
- The assignment expects you to go beyond velocity tracking by adding principled reward terms (posture stabilization, foot clearance, slip minimization, smooth actions, contact and collision penalties), robustness via domain randomization, and clear benchmarking metrics for evaluation as described in the course guidelines.
- Keep your repository organized, document your changes in the README, and ensure your scripts are reproducible, as these factors are part of grading alongside policy quality and the short demo video deliverable.
- Isaac Lab documentation — Everything you need to know about IsaacLab, and more!
- Isaac Lab ANYmal C environment — This targets ANYmal C (not Unitree Go2), so use it as a reference and adapt robot config, assets, and reward to Go2.
- DMO (IsaacGym) Go2 walking project page • Go2 walking environment used by the authors • Config file used by the authors — Look at the function
compute_reward_CaT(beware that some reward terms have a weight of 0 and thus are deactivated, check weights in the config file); this implementation includes strong reward shaping, domain randomization, and training disturbances for robust sim‑to‑real, but it is written for legacy IsaacGym and the challenge is to re-implement it in Isaac Lab. - API References:
- ArticulationData (
robot.data) — Containsroot_pos_w,joint_pos,projected_gravity_b, etc. - ContactSensorData (
_contact_sensor.data) — Containsnet_forces_w(contact forces).
- ArticulationData (
Students should only edit README.md below this ligne.
We apply the suggested changes from the tutorial as-is for parts 1 to 4. We document our implementations of parts 5 to 6 below, as well as additional featuers such as the actuator friction model and rough terrain environment training.
To achieve stable and natural-looking locomotion, we added penalties for unwanted behaviors that the basic velocity tracking rewards don't address.
Add the following reward scales to rob6323_go2_env_cfg.py:
# Additional reward scales
orient_reward_scale = -5.0
lin_vel_z_reward_scale = -0.02
dof_vel_reward_scale = -0.0001
ang_vel_xy_reward_scale = -0.001Add these reward calculations in _get_rewards() method in rob6323_go2_env.py:
# Penalize non-vertical orientation
rew_orient = torch.sum(torch.square(self.robot.data.projected_gravity_b[:, :2]), dim=1)
# Penalize vertical velocity
rew_lin_vel_z = torch.square(self.robot.data.root_lin_vel_b[:, 2])
# Penalize high joint velocities
rew_dof_vel = torch.sum(torch.square(self.robot.data.joint_vel), dim=1)
# Penalize angular velocity in XY plane
rew_ang_vel_xy = torch.sum(torch.square(self.robot.data.root_ang_vel_b[:, :2]), dim=1)
# Add to rewards dictionary
rewards = {
...
"orient": rew_orient * self.cfg.orient_reward_scale,
"lin_vel_z": rew_lin_vel_z * self.cfg.lin_vel_z_reward_scale,
"dof_vel": rew_dof_vel * self.cfg.dof_vel_reward_scale,
"ang_vel_xy": rew_ang_vel_xy * self.cfg.ang_vel_xy_reward_scale,
}Update logging dictionary in __init__:
self._episode_sums = {
key: torch.zeros(self.num_envs, dtype=torch.float, device=self.device)
for key in [
"track_lin_vel_xy_exp",
"track_ang_vel_z_exp",
"rew_action_rate",
"raibert_heuristic",
"orient",
"lin_vel_z",
"dof_vel",
"ang_vel_xy",
]
}In __init__, add separate indices for robot bodies and contact sensor:
self._feet_ids = []
self._feet_ids_sensor = []
foot_names = ["FL_foot", "FR_foot", "RL_foot", "RR_foot"]
for name in foot_names:
id_list, _ = self.robot.find_bodies(name)
self._feet_ids.append(id_list[0])
sensor_id_list, _ = self._contact_sensor.find_bodies(name)
self._feet_ids_sensor.append(sensor_id_list[0])Add this code in _get_rewards():
# Foot clearance reward
phases = 1 - torch.abs(1.0 - torch.clip((self.foot_indices * 2.0) - 1.0, 0.0, 1.0) * 2.0)
foot_height = self.foot_positions_w[:, :, 2]
target_height = 0.08 * phases + 0.02
rew_foot_clearance = torch.square(target_height - foot_height) * (1 - self.desired_contact_states)
rew_feet_clearance = torch.sum(rew_foot_clearance, dim=1)Add this code in _get_rewards():
# Contact force tracking reward
foot_forces = torch.norm(self._contact_sensor.data.net_forces_w[:, self._feet_ids_sensor, :], dim=-1)
desired_contact = self.desired_contact_states
rew_tracking_contacts_shaped_force = torch.zeros(self.num_envs, device=self.device)
for i in range(4):
rew_tracking_contacts_shaped_force += -(1 - desired_contact[:, i]) * (1 - torch.exp(-1 * foot_forces[:, i] ** 2 / 100.))
rew_tracking_contacts_shaped_force /= 4Update _get_rewards():
rewards = {
...
"feet_clearance": rew_feet_clearance * self.cfg.feet_clearance_reward_scale,
"tracking_contacts_shaped_force": rew_tracking_contacts_shaped_force * self.cfg.tracking_contacts_shaped_force_reward_scale,
}Add friction parameters to rob6323_go2_env_cfg.py:
# Friction ranges for actuator
actuator_mu_range_min = 0.001
actuator_mu_range_max = 0.3
actuator_st_range_min = 0.001
actuator_st_range_max = 2.5In __init__ of rob6323_go2_env.py:
# Actuator friction coefficients
self.stiction_coeff = torch.zeros(self.num_envs, 12, device=self.device)
self.viscous_coeff = torch.zeros(self.num_envs, 12, device=self.device)In _reset_idx():
# Sample new friction coefficients
self.stiction_coeff[env_ids] = torch.zeros_like(self.stiction_coeff[env_ids]).uniform_(
self.cfg.actuator_st_range_min,
self.cfg.actuator_st_range_max
)
self.viscous_coeff[env_ids] = torch.zeros_like(self.viscous_coeff[env_ids]).uniform_(
self.cfg.actuator_mu_range_min,
self.cfg.actuator_mu_range_max
)Update _apply_action():
def _apply_action(self) -> None:
# Compute friction torques
stiction = self.stiction_coeff * torch.tanh(self.robot.data.joint_vel / 0.1)
viscous = self.viscous_coeff * self.robot.data.joint_vel
actuator_friction_torque = stiction + viscous
# Compute PD torques with friction
self.torques = torch.clip(
(
self.Kp * (self.desired_joint_pos - self.robot.data.joint_pos)
- self.Kd * self.robot.data.joint_vel
- actuator_friction_torque
),
-self.torque_limits,
self.torque_limits,
)
self.robot.set_joint_effort_target(self.torques)Added physics material randomization using Isaac Lab's MDP events system in rob6323_go2_env_cfg.py:
from isaaclab.managers import EventTermCfg as EventTerm
from isaaclab.managers import SceneEntityCfg
@configclass
class EventCfg:
robot_physics_material = EventTerm(
func=mdp.randomize_rigid_body_material,
mode="reset",
params={
"asset_cfg": SceneEntityCfg("robot", body_names=".*"),
"static_friction_range": (0.7, 1.0),
"dynamic_friction_range": (0.6, 1.0),
"restitution_range": (1.0, 1.0),
"num_buckets": 250,
"make_consistent": True
},
)
@configclass
class Rob6323Go2EnvCfg(DirectRLEnvCfg):
...
events: EventCfg = EventCfg()This randomizes the friction properties of all robot body parts on each episode reset, making the policy more robust to varying ground conditions.
Create rob6323_go2_rough_env_cfg.py with generator. It has the following unique elements:
USE_FLAT = False # used for curriculum learning
terrain = TerrainImporterCfg(
prim_path="/World/ground",
terrain_type="generator",
terrain_generator=ROUGH_TERRAINS_CFG,
max_init_terrain_level=5,
collision_group=-1,
physics_material=sim_utils.RigidBodyMaterialCfg(
friction_combine_mode="multiply",
restitution_combine_mode="multiply",
static_friction=1.0,
dynamic_friction=1.0,
restitution=0.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,
)
def __post_init__(self):
"""Post initialization - scale down terrains for Go2."""
if not USE_FLAT and self.terrain.terrain_generator is not None:
# Scale down the terrains because the Go2 robot is small
self.terrain.terrain_generator.sub_terrains["boxes"].grid_height_range = (0.025, 0.1)
self.terrain.terrain_generator.sub_terrains["random_rough"].noise_range = (0.01, 0.06)
self.terrain.terrain_generator.sub_terrains["random_rough"].noise_step = 0.01
# Disable curriculum for now
self.terrain.terrain_generator.curriculum = FalseCreate rob6323_go2_rough_env.py in this case it is a mirror copy of the flat environment (we don't add sensors for our implementation)
class Rob6323Go2RoughEnv(DirectRLEnv):
cfg: Rob6323Go2RoughEnvCfg
def __init__(self, cfg: Rob6323Go2RoughEnvCfg, render_mode: str | None = None, **kwargs):
super().__init__(cfg, render_mode, **kwargs)
...update rsl_rl_ppo_cfg.py. Notice here we increase the network size to [512, 256, 128] for both actor and critic, following the recommendation from the Anymal C environment.
@configclass
class PPORunnerRoughCfg(RslRlOnPolicyRunnerCfg):
num_steps_per_env = 32
max_iterations = 500
save_interval = 50
experiment_name = "go2_rough_direct"
policy = RslRlPpoActorCriticCfg(
init_noise_std=1.0,
actor_obs_normalization=True,
critic_obs_normalization=True,
actor_hidden_dims=[512, 256, 128],
critic_hidden_dims=[512, 256, 128],
activation="elu",
)
algorithm = RslRlPpoAlgorithmCfg(
value_loss_coef=4.0,
use_clipped_value_loss=True,
clip_param=0.2,
entropy_coef=0.0,
num_learning_epochs=5,
num_mini_batches=4,
learning_rate=3.0e-4,
schedule="adaptive",
gamma=0.99,
lam=0.95,
desired_kl=0.008,
max_grad_norm=1.0,
)Update __init__.py:
gym.register(
id="Template-Rob6323-Go2-Direct-Rough-v0",
entry_point=f"{__name__}.rob6323_go2_rough_env:Rob6323Go2RoughEnv",
disable_env_checker=True,
kwargs={
"env_cfg_entry_point": f"{__name__}.rob6323_go2_rough_env_cfg:Rob6323Go2RoughEnvCfg",
"rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:PPORunnerCfg",
},
)To train the robot on uneven surfaces, we employ a two-stage curriculum learning approach: first, we train on rough terrain using the identical reward structure from the flat environment, allowing the policy to adapt to height variations while maintaining learned gait patterns. In the second stage, we remove the Raibert heuristic and gait phasing rewards, as these overly constrain foot placement and timing on irregular surfaces where strict trotting patterns are suboptimal. Training resumes from the final flat environment checkpoint, enabling the policy to develop more adaptive locomotion strategies.
In rob6323_go2_rough_env_cfg.py ensure USE_FLAT=True
cd $HOME/rob6323_go2_project
./train_rough.shWe want to launch the last completed checkpoint. You'll want to replace <last_run_name> with the timestamp of the last completed run, which will be located under the logs/[job_id]rsl_rl/go2_rough_direct folder
cd $HOME/rob6323_go2_project
./train_rough.sh --video --resume --load_run <last_run_name>After training completes, logs will be in logs/[job_id]/rsl_rl/go2_rough_direct/[timestamp]/. Download and view with TensorBoard:
rsync -avzP <netid>@dtn.hpc.nyu.edu:/home/<netid>/rob6323_go2_project/logs ./
tensorboard --logdir ./logsVideos are generated automatically in videos/play/ subdirectory of each run.
- Added orientation penalty to keep base level
- Added vertical velocity penalty to reduce bouncing
- Added joint velocity penalty for smoother motion
- Added angular velocity penalty to reduce roll and pitch
- Implemented foot clearance rewards for proper swing phase
- Implemented contact force tracking for stance phase control
- Implemented actuator friction model with randomization
- Added stiction and viscous friction components
- Randomized friction parameters per episode reset
- Created rough terrain environment with height scanning
- Added 160-point height map to observations
- Scaled terrain difficulty for Go2 dimensions
- PD gains: Kp=20.0, Kd=0.5
- Torque limits: 100.0 Nm
- Friction stiction range: 0.0 to 2.5
- Friction viscous range: 0.0 to 0.3
- Height scan: 1.6m x 1.0m grid at 0.1m resolution
- Observation space: 52 (flat) or 212 (rough terrain)