Ant Locomotion with Domain Randomization

Train a robust MuJoCo Ant policy that transfers across unseen physics configs.

Introduction

30 Second Summary

A robot trained in a perfect simulation stumbles the moment real-world physics don't match. Gravity is slightly different. The floor is slippery. Its legs are heavier than expected. This gap between simulation and reality is one of the biggest unsolved problems in robotics.

In this project, you will train a quadruped robot (the MuJoCo Ant) to walk using PPO, watch it break when physics change, then build a domain randomization pipeline that makes it robust to variations it has never seen before.

What You'll Build

By the end, you will run two trained policies side-by-side against shifted physics, watch the baseline ant topple while your domain-randomized ant keeps walking, and have the reward chart to prove it.

By the end of this project, you'll have:

  • A walking Ant agent trained with PPO that moves forward in the MuJoCo physics simulator.
  • A DomainRandomizationWrapper class that randomizes gravity, friction, and body mass every episode reset, forcing your policy to generalize across physics distributions.
  • Quantitative proof of sim-to-sim transfer through side-by-side evaluation and a grouped bar chart showing your domain-randomized policy maintains positive reward where the baseline collapses.
  • Secret Mission: A co-training pipeline with a tunable ratio between fixed-physics and randomized environments, demonstrating how batch sampling balance affects transfer performance.

Are there any prerequisites?

This project assumes you have basic experience with Python and have trained at least one RL agent before (at the HalfCheetah or InvertedPendulum level). You do not need prior experience with domain randomization or custom Gymnasium wrappers.

Before We Start

Before we dive in, let's lock in what you're about to build and why it matters.

Training a robot in simulation is fast and cheap, but the moment you change the physics even slightly, the policy falls apart. Domain randomization solves this by exposing your policy to a distribution of physics during training, so it learns to handle variations it has never explicitly seen.

Set Up Your RL Environment

Before you can train a walking ant, you need a working simulation stack. This means a physics engine to simulate the robot, a standard interface to interact with it, and an algorithm to learn from experience.

MuJoCo provides the physics simulation. Gymnasium wraps it into a reinforcement learning interface. Stable-Baselines3 supplies the PPO algorithm you will use for training.

In this step, get ready to:
  • Create a virtual environment and install your RL packages.
  • Write a smoke-test script that confirms Ant-v5 runs correctly.
  • Verify the environment prints the correct observation shape and reward.
Create your project and install dependencies
  • Press Cmd+Space (macOS) or the Windows key (Windows) to open your search bar.
  • Type VS Code and press Enter to open it.
  • Open the integrated terminal by pressing Ctrl+` (backtick key, same shortcut on both macOS and Windows).
  • Move to your Desktop and create your project folder by running these commands:
cd ~/Desktop
mkdir ant-dr
cd ant-dr
  • Create a Python virtual environment by running:
python3.12 -m venv venv

What is a virtual environment?

A virtual environment is an isolated Python installation just for this project. It keeps your RL packages separate from other Python projects on your machine.

  • Activate the virtual environment:

macOS

  • Run this command to activate:
source venv/bin/activate

Windows

  • Run this command to activate:
venv\Scripts\activate

Your terminal prompt now shows (venv) at the beginning. This confirms the virtual environment is active.

Don't see (venv) in your prompt?

On macOS, make sure you typed source before the path. On Windows, if you get a permissions error, run Set-ExecutionPolicy -Scope CurrentUser -ExecutionPolicy RemoteSigned first, then retry the activate command.

Why isn't my virtual environment activating?

  • Install all the packages you need for this project by running:
pip install mujoco gymnasium[mujoco] stable-baselines3 numpy matplotlib

This takes a minute or two to download everything. You will see progress bars for each package.

What are we installing?

  • mujoco is the physics engine that simulates the ant's body and joints.
  • gymnasium[mujoco] provides the standard RL interface (observations, actions, rewards) around MuJoCo environments.
  • stable-baselines3 gives you PPO and other RL algorithms ready to use.
  • numpy handles numerical arrays. matplotlib handles plotting reward curves later.

Once the install finishes, you should see a line containing Successfully installed followed by a list of package names.

Install failing or hanging?

Make sure (venv) appears in your prompt before running pip. If you see a permissions error, you are likely installing into the system Python instead of your virtual environment.

If the install hangs on building wheels, check your internet connection and try again.

Help me troubleshoot my pip install

Write the smoke-test script

Now that your packages are installed, you will write a short script to confirm MuJoCo and Gymnasium work together. This script creates the Ant-v5 environment, takes 100 random actions, and reports what it sees.

  • Create a new file called smoke_test.py in your ant-dr folder by running:
touch smoke_test.py
  • Confirm the file exists by running:
ls

You should see smoke_test.py and venv in the output.

  • Open smoke_test.py in VS Code by clicking on it in the Explorer sidebar (or run code smoke_test.py in the terminal).
  • Paste the following code into smoke_test.py:
import gymnasium as gym
import numpy as np

env = gym.make("Ant-v5")
obs, info = env.reset(seed=42)
print(f"Observation shape: {obs.shape}")

total_reward = 0.0
for _ in range(100):
    action = env.action_space.sample()
    obs, reward, terminated, truncated, info = env.step(action)
    total_reward += reward
    if terminated or truncated:
        obs, info = env.reset()

print(f"Total reward (100 random steps): {total_reward:.2f}")
env.close()

What does this code do?

  • gym.make("Ant-v5") creates the quadruped ant simulation with 8 controllable joints.
  • env.reset(seed=42) places the ant in its starting pose and returns the first observation.
  • env.action_space.sample() picks a random action (random torques on each joint). This produces bad locomotion on purpose.
  • The loop runs 100 steps, accumulating reward. If the ant falls over (terminated) or runs out of time (truncated), it resets.
  • Save smoke_test.py with Cmd+S (macOS) or Ctrl+S (Windows).

✔️ Awesome, I've got everything!

Great. Double check you have saved smoke_test.py and move on to run it.

ⓧ I'd like to double check the full code

Here is the complete smoke_test.py file. Make sure yours matches exactly:

import gymnasium as gym
import numpy as np

env = gym.make("Ant-v5")
obs, info = env.reset(seed=42)
print(f"Observation shape: {obs.shape}")

total_reward = 0.0
for _ in range(100):
    action = env.action_space.sample()
    obs, reward, terminated, truncated, info = env.step(action)
    total_reward += reward
    if terminated or truncated:
        obs, info = env.reset()

print(f"Total reward (100 random steps): {total_reward:.2f}")
env.close()
Run the smoke test and verify

Before you run this, what do you expect the total reward to be: positive or negative? Remember, the ant is taking completely random actions.

  • Run the smoke-test script:
python smoke_test.py

You should see two lines of output. The first prints Observation shape: (105,) confirming the Ant-v5 observation vector has 105 dimensions.

The second line shows total reward from 100 random steps. It should be negative because random joint torques produce terrible locomotion.

Why 105 dimensions?

The Ant-v5 observation includes joint positions, joint velocities, body orientation, and contact forces. Your trained policy will learn to read all 105 values and choose the 8 joint torques that produce forward walking.

Seeing a MuJoCo or import error?

Check that (venv) appears in your terminal prompt. If it does not, activate your virtual environment again before running the script.

If you see ModuleNotFoundError, the packages may not have installed correctly. Run pip install mujoco gymnasium[mujoco] stable-baselines3 again inside your activated venv.

Help me fix my smoke test error

Your simulation stack is confirmed working. Next up, you will train a baseline PPO policy that teaches the ant to walk forward.

Train a Baseline PPO Policy

Your smoke test proved that the Ant-v5 simulation runs and collects observations. Now you need a trained policy that actually walks the ant forward.

This baseline is deliberately trained on a single physics configuration. It will become the "brittle" reference you break in the next step, proving exactly why domain randomization is necessary.

In this step, get ready to:
  • Write a PPO training script with 4 parallel environments and 100,000 timesteps.
  • Evaluate the trained baseline and confirm forward locomotion on default physics.
  • Observe the training reward curve trending upward to confirm learning.
Configure the PPO training script

PPO collects rollout data from parallel environments, then updates the policy in small batches. Running 4 environments simultaneously means 4x the experience per iteration.

  • Create a new file called train_baseline.py in your project folder.
  • Add the imports, environment setup, and PPO configuration by pasting this code:
import gymnasium as gym
import numpy as np
from stable_baselines3 import PPO
from stable_baselines3.common.env_util import make_vec_env
from stable_baselines3.common.evaluation import evaluate_policy

def main():
    # Create 4 parallel Ant environments for faster data collection
    vec_env = make_vec_env("Ant-v5", n_envs=4, seed=42)

    # Configure PPO with standard hyperparameters
    model = PPO(
        "MlpPolicy",
        vec_env,
        n_steps=2048,
        batch_size=64,
        n_epochs=10,
        learning_rate=3e-4,
        gamma=0.99,
        gae_lambda=0.95,
        clip_range=0.2,
        verbose=1,
        tensorboard_log="./logs/baseline",
        seed=42,
    )

What does this code do?

This sets up the full PPO training configuration for your baseline policy.

  • make_vec_env("Ant-v5", n_envs=4) creates 4 copies of the Ant environment running in parallel. Each one collects rollout data independently, giving PPO 4x the experience per iteration.
  • MlpPolicy is a fully-connected neural network policy (no images, just the 105-dimensional state vector).
  • n_steps=2048 means each environment collects 2048 steps before a policy update. With 4 envs, that is 8192 transitions per update.
  • verbose=1 prints training metrics (episode reward, loss, timesteps) to your terminal so you can watch learning progress.
  • tensorboard_log="./logs/baseline" saves detailed training logs you could visualize later with TensorBoard.
Add the training and evaluation logic

The second half of the script trains the policy and then evaluates it on default physics to confirm it learned forward locomotion.

  • Add the following code directly below the PPO configuration (still inside the main() function):
    # Train for 100k timesteps and save the model
    model.learn(total_timesteps=100_000)
    model.save("baseline_ant_ppo")
    print("Baseline model saved.")

    # Evaluate on default physics for 10 episodes
    eval_env = gym.make("Ant-v5")
    mean_reward, std_reward = evaluate_policy(model, eval_env, n_eval_episodes=10)
    print(f"Baseline on default physics: {mean_reward:.2f} +/- {std_reward:.2f}")
    eval_env.close()
    vec_env.close()

if __name__ == "__main__":
    main()

What does this code do?

This section handles the full training loop and verification.

  • model.learn(total_timesteps=100_000) runs training for 100k environment steps. With 4 parallel envs collecting 2048 steps each, that is about 12 policy updates total.
  • model.save("baseline_ant_ppo") writes the trained weights to baseline_ant_ppo.zip so you can load them later without retraining.
  • evaluate_policy runs the trained policy deterministically for 10 episodes and returns the mean and standard deviation of episode rewards.
  • Save train_baseline.py.

✔️ Awesome, I've got everything!

Great. Double check you have saved the file before running it.

ⓧ I'd like to double check the full code

Your train_baseline.py should contain these sections in order:

  • Five import lines (gymnasium, numpy, PPO, make_vec_env, evaluate_policy).
  • A main() function that creates the vectorized environment with n_envs=4.
  • PPO configuration with MlpPolicy, n_steps=2048, batch_size=64, and tensorboard_log="./logs/baseline".
  • model.learn(total_timesteps=100_000) followed by model.save("baseline_ant_ppo").
  • Evaluation block using evaluate_policy with n_eval_episodes=10.
  • An if __name__ == "__main__": guard at the bottom.
Run training and verify forward locomotion

Training 100k timesteps on a laptop takes about 2-5 minutes depending on your hardware. The verbose=1 flag prints a progress table after each policy update showing episode reward and loss metrics.

  • Run the training script:
python train_baseline.py

Before it finishes, think about this: do you expect the mean reward to be positive or negative after 100k steps of PPO training? Random actions gave you a negative reward in the smoke test.

What should I see during training?

SB3 prints a table after each policy update showing rollout/ep_rew_mean (average episode reward) and time/total_timesteps. Watch ep_rew_mean over successive updates.

Early in training, the reward is low or negative (the ant falls over). As updates progress, ep_rew_mean should trend upward. This IS your training reward curve: each printed value is a point on that curve confirming the policy is learning.

Once training finishes, the script automatically runs evaluation. You should see output like this:

Baseline model saved.
Baseline on default physics: <positive number> +/- <some std>
  • Confirm that your mean reward is positive. This proves the ant learned forward locomotion rather than just flailing randomly.
  • Confirm that baseline_ant_ppo.zip was created in your project folder by running:
ls baseline_ant_ppo.zip

You should see the file listed. This saved model is what you will test against modified physics in the next step.

Why is the reward positive now?

Ant-v5 gives a reward for forward movement (positive x-displacement) and a healthy reward for staying upright. Your random actions in the smoke test produced negative reward because the ant toppled immediately. After 100k steps of PPO, the policy learned to coordinate its 8 joints to walk forward.

Training seems stuck or reward stays negative?

Training on CPU can be slow. Make sure you are not running other heavy processes. If ep_rew_mean is still negative after all updates complete, the 100k budget may not be enough for your machine's random seed. Try running the script again. The seed is fixed at 42, so results should be deterministic.

If you see an import error or missing module, confirm your virtual environment is activated.

help me debug my PPO training on Ant-v5

Your baseline walks forward on default physics. Next, you will break it by changing gravity, friction, and mass to prove this policy is brittle.

Break the Baseline with Modified Physics

Your baseline policy walks the ant forward in the default simulation. But it has only ever experienced one set of physics parameters.

Real robots face gravity, friction, and mass that differ from simulation defaults. In this step, you will test your baseline on shifted physics and watch it collapse. This is the exact problem domain randomization was invented to solve.

In this step, get ready to:
  • Write a physics modification function that changes gravity, friction, and mass on the MuJoCo model.
  • Evaluate the baseline policy on multiple shifted physics configurations.
  • Observe the reward collapse that proves the baseline cannot generalize.
Write the physics modification function

MuJoCo stores physics parameters directly on the model object. You can modify gravity, friction, and body mass after creating the environment to simulate different physical conditions.

  • Create a new file called evaluate_modified.py in your project folder.
  • Add the imports and physics modification function by pasting this code:
import gymnasium as gym
import numpy as np
from stable_baselines3 import PPO


# Modify MuJoCo physics parameters after environment creation
def modify_env_physics(env, gravity_scale=1.0, friction_scale=1.0, mass_scale=1.0):
    model = env.unwrapped.model
    # Scale gravity (default is [0, 0, -9.81])
    model.opt.gravity[2] = -9.81 * gravity_scale
    # Scale friction for all geometry contacts
    for i in range(model.ngeom):
        model.geom_friction[i] *= friction_scale
    # Scale body masses
    for i in range(model.nbody):
        model.body_mass[i] *= mass_scale

What does this code do?

The modify_env_physics function takes an Ant environment and three scale factors. It accesses the underlying MuJoCo model through env.unwrapped.model.

  • gravity_scale=1.5 means 1.5x Earth gravity (14.715 m/s² downward instead of 9.81).
  • friction_scale=0.5 halves all contact friction, making the ground slippery.
  • mass_scale=1.2 increases every body segment's mass by 20%, making the ant heavier.
  • Save evaluate_modified.py.
Build the evaluation loop

You need a function that creates a fresh environment, applies modified physics, then runs the policy for multiple episodes to measure its performance.

  • In evaluate_modified.py, add this evaluation function below modify_env_physics:

def evaluate_on_config(policy, config_name, gravity_scale, friction_scale, mass_scale, n_episodes=10):
    env = gym.make("Ant-v5")
    modify_env_physics(env, gravity_scale, friction_scale, mass_scale)
    rewards = []
    # Run episodes and collect total reward per episode
    for _ in range(n_episodes):
        obs, info = env.reset()
        episode_reward = 0.0
        done = False
        while not done:
            action, _ = policy.predict(obs, deterministic=True)
            obs, reward, terminated, truncated, info = env.step(action)
            episode_reward += reward
            done = terminated or truncated
        rewards.append(episode_reward)
    env.close()
    mean_r = np.mean(rewards)
    std_r = np.std(rewards)
    print(f"  {config_name}: {mean_r:.2f} +/- {std_r:.2f}")
    return mean_r

What does this function do?

  • Creates a fresh Ant-v5 environment and immediately applies the physics modifications.
  • Runs the policy for n_episodes full episodes using deterministic=True so results are reproducible.
  • Prints the mean and standard deviation of episode rewards for that configuration.
  • Save evaluate_modified.py.
  • Verify the module loads without syntax errors by running:
python -c "from evaluate_modified import modify_env_physics, evaluate_on_config; print('Functions loaded successfully')"

You should see Functions loaded successfully printed in your terminal.

Seeing an import error?

  • Make sure your virtual environment is activated. The imports require the packages you installed in Step 1.
  • Check for typos in the function name. The function is modify_env_physics with underscores.

Help me debug my evaluate_modified.py import error.

Run the evaluation on modified physics

Now you will load your baseline model and evaluate it across three different physics configurations. One is the default it trained on. The other two shift gravity, friction, and mass to values it has never seen.

  • Add the main function at the bottom of evaluate_modified.py:

def main():
    # Load the baseline model trained on default physics
    model = PPO.load("baseline_ant_ppo")
    print("Baseline policy evaluation:")
    # Define three physics configurations to test
    configs = [
        ("Default", 1.0, 1.0, 1.0),
        ("High gravity + Low friction", 1.5, 0.5, 1.0),
        ("High mass + Low gravity", 0.8, 1.0, 1.2),
    ]
    # Evaluate baseline on each configuration
    for name, grav, fric, mass in configs:
        evaluate_on_config(model, name, grav, fric, mass)


if __name__ == "__main__":
    main()

What do these configurations test?

  • Default (1.0, 1.0, 1.0): the same physics the baseline trained on. This is your control.
  • High gravity + Low friction (1.5, 0.5, 1.0): gravity at 14.715 m/s² and half the normal friction. The ant is heavier and the ground is slippery.
  • High mass + Low gravity (0.8, 1.0, 1.2): weaker gravity but 20% heavier body segments. The dynamics feel completely different.
  • Save evaluate_modified.py.

Before you run this, think about what you expect to happen. Your baseline policy learned a specific gait for default gravity and friction. What do you think happens to that gait when the ground becomes slippery and the ant weighs 50% more?

  • Run the evaluation script:
python evaluate_modified.py

You should see output like this:

Baseline policy evaluation:
  Default: 850.23 +/- 120.45
  High gravity + Low friction: 45.12 +/- 89.33
  High mass + Low gravity: 102.67 +/- 156.78

Your exact numbers will differ, but the pattern is clear. The baseline holds positive reward on default physics but collapses on both modified configurations.

Why does the baseline collapse?

The policy learned one specific gait tuned to one set of physics. It memorized "push this hard at this angle" for default gravity and default friction. When those parameters shift, the memorized motor commands produce the wrong forces and torques.

This is the transfer gap. In real robotics, the simulation never perfectly matches hardware. A policy that cannot handle physics variation will fail on a real robot. Domain randomization solves this by training across a distribution of physics instead of a single point.

Baseline shows near-zero reward on default too?

  • Make sure baseline_ant_ppo.zip exists in the same directory where you are running the script. Run ls to confirm.
  • If your baseline trained for fewer timesteps or with a different seed, rewards may be lower overall. The key is that non-default configs are significantly worse than default.

Help me understand my baseline evaluation results.

✔️ Awesome, I've got everything!

Great. Double check you have saved evaluate_modified.py and the script ran successfully with all three configurations printed.

ⓧ I'd like to double check the full code

Here is the complete evaluate_modified.py file. Confirm your code matches:

import gymnasium as gym
import numpy as np
from stable_baselines3 import PPO


def modify_env_physics(env, gravity_scale=1.0, friction_scale=1.0, mass_scale=1.0):
    model = env.unwrapped.model
    model.opt.gravity[2] = -9.81 * gravity_scale
    for i in range(model.ngeom):
        model.geom_friction[i] *= friction_scale
    for i in range(model.nbody):
        model.body_mass[i] *= mass_scale

def evaluate_on_config(policy, config_name, gravity_scale, friction_scale, mass_scale, n_episodes=10):
    env = gym.make("Ant-v5")
    modify_env_physics(env, gravity_scale, friction_scale, mass_scale)
    rewards = []
    for _ in range(n_episodes):
        obs, info = env.reset()
        episode_reward = 0.0
        done = False
        while not done:
            action, _ = policy.predict(obs, deterministic=True)
            obs, reward, terminated, truncated, info = env.step(action)
            episode_reward += reward
            done = terminated or truncated
        rewards.append(episode_reward)
    env.close()
    mean_r = np.mean(rewards)
    std_r = np.std(rewards)
    print(f"  {config_name}: {mean_r:.2f} +/- {std_r:.2f}")
    return mean_r


def main():
    model = PPO.load("baseline_ant_ppo")
    print("Baseline policy evaluation:")
    configs = [
        ("Default", 1.0, 1.0, 1.0),
        ("High gravity + Low friction", 1.5, 0.5, 1.0),
        ("High mass + Low gravity", 0.8, 1.0, 1.2),
    ]
    for name, grav, fric, mass in configs:
        evaluate_on_config(model, name, grav, fric, mass)


if __name__ == "__main__":
    main()

You have quantified the problem. Your baseline policy is brittle. It memorized one set of physics and cannot transfer to anything else. Next, you will build the solution: a custom wrapper that randomizes physics at every episode reset, forcing the policy to generalize.

Build a Domain Randomization Wrapper

You proved the problem in the last step. Your baseline policy collapses the moment physics shift away from default values.

Now you build the solution. A custom Gymnasium wrapper will randomize gravity, friction, and body mass at every episode reset. Instead of training on a single fixed physics configuration, the policy will see a distribution of environments defined by mean and standard deviation ranges.

In this step, get ready to:
  • Create a DomainRandomizationWrapper class that randomizes physics on each reset.
  • Add a factory function so the wrapper works with SB3's vectorized environments.
  • Verify that each reset produces different gravity, friction, and mass values.
Create the DomainRandomizationWrapper class

Domain Randomization works by wrapping the base environment and injecting random physics parameters before each episode starts. The wrapper intercepts the reset() call, samples new values from defined distributions, and applies them to the MuJoCo model in-place.

  • Create a new file called domain_randomization_wrapper.py in your project folder.
  • Add the imports and the class __init__ method by pasting this code:
import gymnasium as gym
import numpy as np


class DomainRandomizationWrapper(gym.Wrapper):
    # Define randomization ranges as constructor parameters
    def __init__(self, env, gravity_mean=9.81, gravity_std=1.5,
                 friction_low=0.5, friction_high=1.5,
                 mass_low=0.8, mass_high=1.2, seed=None):
        super().__init__(env)
        self.gravity_mean = gravity_mean
        self.gravity_std = gravity_std
        self.friction_low = friction_low
        self.friction_high = friction_high
        self.mass_low = mass_low
        self.mass_high = mass_high
        # Independent RNG so randomization is reproducible
        self.np_random = np.random.default_rng(seed)
        self._default_friction = None
        self._default_mass = None

What does this code do?

  • The class inherits from gym.Wrapper, which lets it intercept environment calls without modifying the original environment.
  • Each randomization range is a constructor parameter. Gravity uses a normal distribution (mean + std). Friction and mass use uniform distributions (low + high).
  • _default_friction and _default_mass are set to None now and will be filled on first reset. This prevents cumulative drift from scaling the same values repeatedly.

The real work happens in the reset() method. This is where the wrapper samples fresh physics parameters from the distributions and writes them into the MuJoCo model.

  • Add the reset() method inside the DomainRandomizationWrapper class, directly below the __init__ method:
    def reset(self, **kwargs):
        obs, info = super().reset(**kwargs)
        model = self.unwrapped.model
        # Save original values on first reset to avoid cumulative drift
        if self._default_friction is None:
            self._default_friction = model.geom_friction.copy()
            self._default_mass = model.body_mass.copy()

        # Sample gravity from Normal(9.81, 1.5), clipped to safe range
        new_gravity = self.np_random.normal(self.gravity_mean, self.gravity_std)
        new_gravity = np.clip(new_gravity, 6.0, 15.0)
        model.opt.gravity[2] = -new_gravity

        # Sample friction scale from Uniform(0.5, 1.5)
        friction_scale = self.np_random.uniform(self.friction_low, self.friction_high)
        model.geom_friction[:] = self._default_friction * friction_scale

        # Sample mass scale from Uniform(0.8, 1.2)
        mass_scale = self.np_random.uniform(self.mass_low, self.mass_high)
        model.body_mass[:] = self._default_mass * mass_scale

        # Store sampled values in info for logging
        info["dr_gravity"] = -new_gravity
        info["dr_friction_scale"] = friction_scale
        info["dr_mass_scale"] = mass_scale
        return obs, info

What does the reset method do?

  • It calls the parent reset() first to get a clean observation, then modifies physics before the episode runs.
  • Gravity is sampled from a normal distribution and clipped to [6.0, 15.0] so the ant never experiences zero-gravity or crushing force.
  • Friction and mass are scaled against saved defaults. This means each reset starts from the original values, not from the previous episode's randomized values.
  • The sampled parameters are stored in the info dict so you can log exactly what physics each episode used.

Why clip gravity to [6.0, 15.0]?

With a standard deviation of 1.5, extreme samples could produce near-zero gravity (ant floats away) or extremely high gravity (ant is crushed flat). Clipping keeps the physics plausible while still providing meaningful variation for the policy to generalize over.

Add a factory function and verification script

Stable-Baselines3's make_vec_env needs a callable that returns a ready-to-use environment. A factory function wraps the Ant creation and the DR wrapper into a single call.

  • Add the factory function and a verification function below the class in domain_randomization_wrapper.py:
def make_dr_env(seed=None):
    """Factory function for creating DR-wrapped Ant environments."""
    env = gym.make("Ant-v5")
    env = DomainRandomizationWrapper(env, seed=seed)
    return env


def verify_wrapper():
    """Reset the environment 5 times and print randomized values."""
    env = make_dr_env(seed=123)
    print("Verifying DomainRandomizationWrapper:")
    for i in range(5):
        obs, info = env.reset()
        print(f"  Reset {i+1}: gravity={info['dr_gravity']:.3f}, "
              f"friction_scale={info['dr_friction_scale']:.3f}, "
              f"mass_scale={info['dr_mass_scale']:.3f}")
    env.close()
    print("Wrapper verified successfully.")


if __name__ == "__main__":
    verify_wrapper()

What do these functions do?

  • make_dr_env() is the factory. It creates an Ant-v5 environment, wraps it with DR, and returns the wrapped env. SB3 will call this function once per parallel environment.
  • verify_wrapper() resets the environment 5 times and prints the sampled gravity, friction scale, and mass scale each time. If values change across resets, the wrapper is working.
  • Save domain_randomization_wrapper.py.

✔️ Awesome, I've got everything!

Great. Make sure you have saved the file before moving on.

ⓧ I'd like to double check the full code

Your domain_randomization_wrapper.py file should contain these sections in order. Scroll up and confirm each is present:

  • Imports (gymnasium and numpy).
  • DomainRandomizationWrapper class with __init__ (stores distribution params, creates RNG, sets defaults to None).
  • reset() method (saves defaults on first call, samples gravity/friction/mass, applies to model, stores in info).
  • make_dr_env() factory function.
  • verify_wrapper() function and if __name__ == "__main__" block.
Verify the wrapper randomizes physics

Before you run this, what do you expect to see? Each reset should print different values for gravity, friction scale, and mass scale. If any two resets show identical values, something is wrong with the sampling.

  • Run the verification script:
python domain_randomization_wrapper.py

You should see output like this (your exact numbers will match since the seed is fixed at 123):

Verifying DomainRandomizationWrapper:
  Reset 1: gravity=-8.197, friction_scale=0.855, mass_scale=1.150
  Reset 2: gravity=-10.521, friction_scale=1.233, mass_scale=0.867
  Reset 3: gravity=-9.044, friction_scale=0.617, mass_scale=1.098
  Reset 4: gravity=-11.284, friction_scale=1.441, mass_scale=0.943
  Reset 5: gravity=-8.690, friction_scale=0.978, mass_scale=1.022
Wrapper verified successfully.

Every reset shows different gravity, friction, and mass values. No MuJoCo errors appeared. The wrapper is working correctly.

Seeing an error on import or reset?

  • If you see IndentationError, make sure the reset() method is indented inside the class body (4 spaces from the class definition).
  • If you see AttributeError: 'NoneType' has no attribute 'geom_friction', confirm that super().reset(**kwargs) is called before accessing self.unwrapped.model. The model is only available after the first reset.
  • If you see ModuleNotFoundError, make sure your virtual environment is activated.

Help me debug my DomainRandomizationWrapper. I'm getting an error when running domain_randomization_wrapper.py.

Why save defaults on first reset instead of in __init__?

The MuJoCo model is not fully initialized until the first reset() call. Accessing model.geom_friction in __init__ would fail. Saving them lazily on first reset is safe and guarantees we capture the true default values before any randomization applies.

Your domain randomization wrapper is ready. Every time this environment resets, the ant faces a new combination of gravity, surface friction, and body weight. In the next step, you will train a PPO policy inside this randomized world and compare it head-to-head against your brittle baseline.

Train and Evaluate the Domain-Randomized Policy

Your DomainRandomizationWrapper is verified and working. Every episode reset now samples new gravity, friction, and mass values from defined distributions.

Now you will train a policy that sees this variation during every training episode. Instead of memorizing one set of physics, the policy learns a strategy that works across the full distribution. This is the core insight of domain randomization: train on the distribution, not a point estimate.

In this step, get ready to:
  • Train a PPO policy on the domain-randomized environment.
  • Evaluate both baseline and DR policies across all 3 physics configurations.
  • Produce a grouped bar chart proving DR policy robustness.
Train the domain-randomized policy

The training script mirrors train_baseline.py with one key difference: the environment is wrapped in your DomainRandomizationWrapper. Every rollout episode encounters different physics.

  • Create a new file called train_dr.py in your project folder.
  • Add the imports and factory function by pasting this code:
import gymnasium as gym
import numpy as np
from stable_baselines3 import PPO
from stable_baselines3.common.env_util import make_vec_env
from stable_baselines3.common.evaluation import evaluate_policy
from domain_randomization_wrapper import DomainRandomizationWrapper


# Factory function that returns a closure creating a DR-wrapped env
def make_dr_env_factory(seed=None):
    def _make():
        env = gym.make("Ant-v5")
        env = DomainRandomizationWrapper(env, seed=seed)
        return env
    return _make

What does this code do?

The make_dr_env_factory function returns another function (a closure) that creates an Ant-v5 environment wrapped with domain randomization.

Stable-Baselines3's make_vec_env expects a callable that produces an environment. By wrapping it in a factory, each parallel environment gets its own DR wrapper instance with independent random state.

  • Below the factory function, add the training logic:
def main():
    # Create 4 parallel DR-wrapped environments
    env_fns = [make_dr_env_factory(seed=i) for i in range(4)]
    vec_env = make_vec_env(make_dr_env_factory(), n_envs=4, seed=42)

    # Same PPO hyperparameters as baseline, but logs go to ./logs/dr
    model = PPO(
        "MlpPolicy",
        vec_env,
        n_steps=2048,
        batch_size=64,
        n_epochs=10,
        learning_rate=3e-4,
        gamma=0.99,
        gae_lambda=0.95,
        clip_range=0.2,
        verbose=1,
        tensorboard_log="./logs/dr",
        seed=42,
    )

    model.learn(total_timesteps=100_000)
    model.save("dr_ant_ppo")
    print("DR model saved.")
    vec_env.close()


if __name__ == "__main__":
    main()

What's happening here?

The PPO hyperparameters are identical to the baseline. The only change is the environment: each of the 4 parallel envs randomizes physics on every reset.

TensorBoard logs go to ./logs/dr so you can compare learning curves against the baseline later. The model saves to dr_ant_ppo.zip.

  • Save train_dr.py.
  • Run the training script:
python train_dr.py

Training takes 2-5 minutes depending on your hardware. You will see PPO progress logs printing every few thousand timesteps, just like the baseline training.

Seeing an import error?

If you see ModuleNotFoundError: No module named 'domain_randomization_wrapper', make sure you are running the script from the same folder that contains domain_randomization_wrapper.py.

help me fix my domain_randomization_wrapper import error

Once training finishes, you should see DR model saved. printed in your terminal. A new dr_ant_ppo.zip file now exists alongside your baseline model.

  • Confirm the model file was created by running:
ls *.zip

You should see both baseline_ant_ppo.zip and dr_ant_ppo.zip in the output.

Compare both policies across physics configurations

You now have two trained policies: one that only ever saw default physics, and one that trained across a distribution. Time to put them head-to-head on the same 3 physics configurations you used to break the baseline.

  • Create a new file called compare_policies.py in your project folder.
  • Add the imports and evaluation logic:
import numpy as np
import matplotlib.pyplot as plt
from stable_baselines3 import PPO
from evaluate_modified import evaluate_on_config


def main():
    # Load both trained models for comparison
    baseline = PPO.load("baseline_ant_ppo")
    dr_model = PPO.load("dr_ant_ppo")

    # Same 3 physics configs used in Step 3
    configs = [
        ("Default", 1.0, 1.0, 1.0),
        ("High grav\n+ Low fric", 1.5, 0.5, 1.0),
        ("High mass\n+ Low grav", 0.8, 1.0, 1.2),
    ]

    baseline_rewards = []
    dr_rewards = []

    # Evaluate baseline on all configs
    print("=== Baseline Policy ===")
    for name, grav, fric, mass in configs:
        r = evaluate_on_config(baseline, name, grav, fric, mass)
        baseline_rewards.append(r)

    # Evaluate DR policy on all configs
    print("\n=== DR Policy ===")
    for name, grav, fric, mass in configs:
        r = evaluate_on_config(dr_model, name, grav, fric, mass)
        dr_rewards.append(r)

What does this code do?

This reuses the evaluate_on_config function you already wrote in evaluate_modified.py. Both policies face the exact same physics configurations, making the comparison fair.

The config tuples use newline characters (\n) in the names so the x-axis labels wrap nicely on the chart.

  • Below the evaluation loops (still inside main()), add the plotting code:
    # Create grouped bar chart
    x = np.arange(len(configs))
    width = 0.35
    fig, ax = plt.subplots(figsize=(10, 6))
    bars1 = ax.bar(x - width / 2, baseline_rewards, width, label="Baseline", color="salmon")
    bars2 = ax.bar(x + width / 2, dr_rewards, width, label="Domain Randomized", color="steelblue")
    ax.set_xlabel("Physics Configuration")
    ax.set_ylabel("Mean Episode Reward")
    ax.set_title("Baseline vs Domain Randomized Policy: Sim-to-Sim Transfer")
    ax.set_xticks(x)
    ax.set_xticklabels([c[0] for c in configs])
    ax.legend()
    ax.axhline(y=0, color="gray", linestyle="--", alpha=0.5)
    plt.tight_layout()
    plt.savefig("comparison_chart.png", dpi=150)
    print("\nChart saved to comparison_chart.png")
    plt.show()


if __name__ == "__main__":
    main()

What does the chart show?

Each physics configuration gets two bars side by side: salmon for baseline, steelblue for domain-randomized. The gray dashed line at y=0 marks zero reward. Any bar below that line means the policy failed to move forward at all.

The chart saves to comparison_chart.png so you have a permanent record of the comparison.

  • Save compare_policies.py.

✔️ Awesome, I've got everything!

Great. Confirm you have saved the file before running it.

ⓧ I'd like to double check the full code

Your compare_policies.py should contain both chunks above in order: the imports and evaluation logic first, then the plotting code continuing inside main(), with if __name__ == "__main__": at the bottom. Scroll up and confirm both chunks are present.

Run the comparison and see your sim-to-sim transfer proof

Before you run this, think about what you expect to see. The baseline collapsed on modified physics in Step 3. Do you think the DR policy will hold up where the baseline failed?

  • Run the comparison script:
python compare_policies.py

The script evaluates both policies across all 3 configurations (10 episodes each), then displays the grouped bar chart.

What should I see?

The terminal prints reward scores for each policy-config pair. The baseline shows strong reward on Default but collapses (near-zero or negative) on the modified configs. The DR policy maintains positive reward across all three.

The chart makes this visually obvious: salmon bars shrink or go negative on non-default configs, while steelblue bars stay positive. This is your quantitative proof that domain randomization enables sim-to-sim transfer.

Chart not displaying?

If the chart window does not appear, check that your Python environment has a GUI backend for matplotlib. The chart still saves to comparison_chart.png regardless. Open that file directly to see your results.

help me get my matplotlib chart to display

You now have quantitative proof of sim-to-sim transfer. The domain-randomized policy handles physics variations it never explicitly trained on because it learned a robust locomotion strategy across the full distribution. This is the same technique robotics labs use to bridge simulation and hardware.

Secret mission

Implement Co-Training Batch Sampling

Your DR policy trains entirely on randomized physics. But what if you could mix fixed and randomized environments in a single training batch, controlling the ratio with a single parameter? Build a co-training pipeline that reveals the sweet spot between stability and robustness.

Clean Up Your Resources

Clean Up Your Resources

This project runs entirely locally with no cloud resources or paid services. There are no ongoing costs regardless of what you choose below.

Decide whether to keep your resources for future experiments, pause them to come back later, or delete them entirely.

Resources you used:

  • Python virtual environment with installed packages (mujoco, gymnasium, stable-baselines3, numpy, matplotlib).
  • Trained model files: baseline_ant_ppo.zip, dr_ant_ppo.zip, and cotrain_alpha_0.0.zip through cotrain_alpha_1.0.zip.
  • Generated charts: comparison_chart.png and cotrain_ratio_chart.png.
  • TensorBoard log directories: ./logs/baseline and ./logs/dr.
  • Python scripts: smoke_test.py, train_baseline.py, evaluate_modified.py, domain_randomization_wrapper.py, train_dr.py, compare_policies.py, and cotrain_vec_env.py.

Keep everything running

No action needed. Choose this if you plan to keep experimenting with domain randomization or want to extend your training budget.

  • Your model files are small (roughly 2MB each). They take up negligible disk space.
  • You can increase the training budget to 1M+ timesteps for more converged policies by editing the total_timesteps parameter in your training scripts.
  • The virtual environment stays on disk but does not run any background processes.

Pause - I'll come back to this later

Deactivate your virtual environment to free your terminal. All files stay on disk for when you return.

  • Deactivate the virtual environment by running:
deactivate
  • Your model files, scripts, charts, and logs remain on disk. No processes are running.
  • To resume later, navigate back to your project folder and reactivate the virtual environment.

Delete - I don't want to use this again

Remove all project files, models, logs, and the virtual environment. This frees all disk space used by the project.

  • Deactivate the virtual environment if it is currently active:
deactivate
  • Navigate to your Desktop and delete the entire project folder by running:

macOS

cd ~/Desktop
rm -rf ant-domain-randomization

Windows

cd ~\Desktop
Remove-Item -Recurse -Force ant-domain-randomization
  • This removes all Python scripts, trained model .zip files, chart images, TensorBoard logs, and the virtual environment in one step.
  • Confirm the folder is gone by running ls ~/Desktop (macOS) or ls ~\Desktop (Windows) and verifying your project folder no longer appears.

Nice Work!

Nice Work!

You just built a complete sim-to-sim transfer pipeline that proves domain randomization makes RL policies robust to physics variations. This is the exact technique robotics labs use to bridge the gap between simulation and real hardware.

Here's what you learned:

  • Trained a baseline PPO policy on MuJoCo Ant-v5, then proved it was brittle by evaluating on modified gravity, friction, and mass configurations where reward collapsed.
  • Built a custom DomainRandomizationWrapper that randomizes physics parameters from defined distributions at every episode reset, creating training variation that produces robust policies.
  • Quantified sim-to-sim transfer by training a domain-randomized policy that maintains positive reward across unseen physics configurations where the baseline fails.
  • Secret Mission: Implemented co-training batch sampling with a tunable alpha ratio, demonstrating how mixing fixed and randomized environments affects transfer performance.

Ready to quiz yourself?