Week 11 — Machine Learning for Robotics Day 6 of 7
Today’s Goal
By the end of this session you will have trained a simulated planar robot arm (a 2-link arm with continuous joint angles) to reach a target position using Proximal Policy Optimization (PPO) from Stable Baselines3. You will build a custom Gymnasium environment from scratch, wrap it with the Monitor and VecNormalize helpers, train for a fixed budget, and then evaluate the trained policy against a random baseline to prove it actually learned something.
Specifically, you will be able to:
- Write a custom
gymnasium.Envwith a continuousBoxaction space and a shaped reward function. - Instantiate a PPO agent with explicit hyperparameters and a multi-layer perceptron policy.
- Train, checkpoint, and reload a policy.
- Run a deterministic evaluation and compare the mean episode reward of the trained agent against a random policy.
The arm model is forward kinematics with two links of length 1.0 each. The end-effector position is (cos θ1 + cos(θ1+θ2), sin θ1 + sin(θ1+θ2)). The agent outputs two joint deltas per step. This is small enough to train on a laptop CPU in a couple of minutes, but it has a real continuous control structure, so the lessons transfer to larger arms.
Prerequisites
- Python 3.10 or newer (tested on 3.10 and 3.11).
- A CPU is fine. No GPU needed for this environment.
Install the exact stack:
pip install "stable-baselines3==2.3.2" "gymnasium==0.29.1" "torch==2.3.1" "numpy==1.26.4" "matplotlib==3.9.0" "tensorboard==2.16.2"
stable-baselines3 2.3.x requires gymnasium>=0.28. Pinning gymnasium==0.29.1 avoids the reset() signature changes that landed in 1.0. If you already have a different torch build, keep it; SB3 only needs torch>=1.13.
Verify the install before continuing:
import stable_baselines3, gymnasium, torch
print(stable_baselines3.__version__, gymnasium.__version__, torch.__version__)
You should see 2.3.2 0.29.1 2.3.1 (or your torch version).
Step-by-Step Implementation
Step 1: Setup and Imports
This block defines the environment. Read the reward function carefully: it is the single most important design decision in the whole session.
import numpy as np
import gymnasium as gym
from gymnasium import spaces
LINK_LENGTH = 1.0
MAX_JOINT_DELTA = 0.2
class PlanarReachEnv(gym.Env):
metadata = {"render_modes": ["rgb_array"], "render_fps": 30}
def __init__(self, max_steps: int = 100, goal_radius: float = 0.05):
super().__init__()
self.max_steps = max_steps
self.goal_radius = goal_radius
# action: delta for joint 1 and joint 2, each in [-0.2, 0.2] rad
self.action_space = spaces.Box(
low=-MAX_JOINT_DELTA, high=MAX_JOINT_DELTA, shape=(2,), dtype=np.float32
)
# observation: [cos t1, sin t1, cos t2, sin t2, goal_x, goal_y]
high = np.array([1, 1, 1, 1, 2.5, 2.5], dtype=np.float32)
self.observation_space = spaces.Box(low=-high, high=high, dtype=np.float32)
self.theta = np.zeros(2, dtype=np.float64)
self.goal = np.zeros(2, dtype=np.float64)
self.step_count = 0
def _end_effector(self) -> np.ndarray:
t1, t2 = self.theta
x = LINK_LENGTH * np.cos(t1) + LINK_LENGTH * np.cos(t1 + t2)
y = LINK_LENGTH * np.sin(t1) + LINK_LENGTH * np.sin(t1 + t2)
return np.array([x, y])
def _obs(self) -> np.ndarray:
t1, t2 = self.theta
return np.array(
[np.cos(t1), np.sin(t1), np.cos(t2), np.sin(t2), self.goal[0], self.goal[1]],
dtype=np.float32,
)
def reset(self, seed=None, options=None):
super().reset(seed=seed)
self.theta = self.np_random.uniform(-np.pi, np.pi, size=2)
# sample goals inside the reachable annulus
r = self.np_random.uniform(0.5, 2.0)
ang = self.np_random.uniform(-np.pi, np.pi)
self.goal = np.array([r * np.cos(ang), r * np.sin(ang)])
self.step_count = 0
return self._obs(), {}
def step(self, action):
action = np.clip(action, -MAX_JOINT_DELTA, MAX_JOINT_DELTA)
self.theta = self.theta + action
self.step_count += 1
ee = self._end_effector()
dist = float(np.linalg.norm(ee - self.goal))
# shaped reward: dense distance penalty plus a bonus for being close
reward = -dist + (1.0 if dist < self.goal_radius else 0.0)
terminated = dist < self.goal_radius
truncated = self.step_count >= self.max_steps
info = {"distance": dist}
return self._obs(), reward, terminated, truncated, info
def render(self):
import matplotlib.pyplot as plt
fig, ax = plt.subplots(figsize=(4, 4))
t1, t2 = self.theta
p0 = np.array([0.0, 0.0])
p1 = p0 + LINK_LENGTH * np.array([np.cos(t1), np.sin(t1)])
p2 = p1 + LINK_LENGTH * np.array([np.cos(t1 + t2), np.sin(t1 + t2)])
ax.plot([p0[0], p1[0], p2[0]], [p0[1], p1[1], p2[1]], "-o", lw=4)
ax.scatter(*self.goal, c="red", s=80, marker="*")
ax.set_xlim(-2.5, 2.5)
ax.set_ylim(-2.5, 2.5)
ax.set_aspect("equal")
fig.canvas.draw()
buf = np.asarray(fig.canvas.buffer_rgba())
plt.close(fig)
return buf[:, :, :3]
Key lines: the observation uses cos/sin of the joint angles instead of raw angles. Raw angles wrap at ±π, which creates a discontinuity the MLP has to fight; cos/sin is smooth and bounded. The reward is -dist, a dense signal that gives gradient information at every step, plus a sparse +1.0 bonus on success. Mixing dense and sparse is a common robot-control recipe. Actions are clipped inside step() so a slightly out-of-bounds network output never corrupts the state.
Step 2: Core Implementation
Now wrap the environment and train PPO. Monitor records episode rewards; VecNormalize standardizes observations, which matters because goal coordinates live on a different scale than unit-circle features.
import os
from stable_baselines3 import PPO
from stable_baselines3.common.env_util import make_vec_env
from stable_baselines3.common.vec_env import VecNormalize
from stable_baselines3.common.callbacks import EvalCallback
LOG_DIR = "./logs/ppo_reach"
MODEL_PATH = "./models/ppo_reach"
os.makedirs(LOG_DIR, exist_ok=True)
os.makedirs(MODEL_PATH, exist_ok=True)
def make_env():
return PlanarReachEnv(max_steps=100, goal_radius=0.05)
vec_env = make_vec_env(make_env, n_envs=8, seed=42)
vec_env = VecNormalize(vec_env, norm_obs=True, norm_reward=False, clip_obs=10.0)
eval_env = make_vec_env(make_env, n_envs=1, seed=7)
eval_env = VecNormalize(eval_env, norm_obs=True, norm_reward=False, clip_obs=10.0,
training=False)
eval_env.obs_rms = vec_env.obs_rms
model = PPO(
"MlpPolicy",
vec_env,
learning_rate=3e-4,
n_steps=256,
batch_size=256,
n_epochs=10,
gamma=0.99,
gae_lambda=0.95,
clip_range=0.2,
ent_coef=0.005,
policy_kwargs=dict(net_arch=[128, 128]),
tensorboard_log=LOG_DIR,
verbose=1,
seed=42,
)
eval_cb = EvalCallback(
eval_env,
best_model_save_path=MODEL_PATH,
log_path=LOG_DIR,
eval_freq=5000,
n_eval_episodes=10,
deterministic=True,
)
model.learn(total_timesteps=200_000, callback=eval_cb, progress_bar=False)
model.save(os.path.join(MODEL_PATH, "final"))
vec_env.save(os.path.join(MODEL_PATH, "vecnormalize.pkl"))
print("training done")
Total timesteps are 200_000. With 8 parallel envs and n_steps=256, PPO collects 8 × 256 = 2048 transitions per rollout, so you get about 97 updates. On a modern laptop CPU this takes roughly 90–150 seconds. batch_size=256 divides the rollout buffer evenly; if you change n_steps or n_envs, keep n_envs * n_steps a multiple of batch_size or SB3 will warn and truncate. ent_coef=0.005 keeps a small amount of exploration pressure so the policy does not collapse to a single joint configuration early. The EvalCallback saves the best checkpoint by mean reward, which protects you if the final policy is worse than an intermediate one (PPO does sometimes regress late in training).
Step 3: Visualization
Load the best model and render a rollout. This block produces a PNG contact sheet of frames and prints per-step distances, so you can see the arm converge.
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from stable_baselines3 import PPO
from stable_baselines3.common.vec_env import DummyVecEnv, VecNormalize
vec = DummyVecEnv([lambda: PlanarReachEnv(max_steps=100)])
vec = VecNormalize.load("./models/ppo_reach/vecnormalize.pkl", vec)
vec.training = False
vec.norm_reward = False
model = PPO.load("./models/ppo_reach/best_model", env=vec)
obs = vec.reset()
frames, distances = [], []
for t in range(100):
action, _ = model.predict(obs, deterministic=True)
obs, reward, done, info = vec.step(action)
distances.append(float(info[0]["distance"]))
frames.append(vec.envs[0].render())
if done[0]:
break
print(f"episode length: {len(distances)}")
print(f"initial distance: {distances[0]:.4f}")
print(f"final distance: {distances[-1]:.4f}")
print(f"min distance: {min(distances):.4f}")
picks = np.linspace(0, len(frames) - 1, 6).astype(int)
fig, axes = plt.subplots(1, 6, figsize=(18, 3))
for ax, i in zip(axes, picks):
ax.imshow(frames[i])
ax.set_title(f"t={i} d={distances[i]:.3f}")
ax.axis("off")
plt.tight_layout()
plt.savefig("reach_rollout.png", dpi=110)
print("saved reach_rollout.png")
The VecNormalize.load call restores the exact running mean and variance used during training. If you skip it and feed raw observations to a policy trained on normalized ones, the policy will behave randomly. Setting vec.training = False freezes the statistics so evaluation does not shift them. deterministic=True uses the mean of the Gaussian action distribution rather than sampling, which removes jitter and is what you want for deployment.
Step 4: Validation
Compare the trained policy against a random policy on 50 episodes. This is the only honest check that learning happened.
import numpy as np
from stable_baselines3 import PPO
from stable_baselines3.common.vec_env import DummyVecEnv, VecNormalize
def run_episodes(policy_fn, n_episodes=50, seed=123):
env = PlanarReachEnv(max_steps=100)
rewards, successes = [], 0
rng = np.random.default_rng(seed)
for ep in range(n_episodes):
obs, _ = env.reset(seed=int(rng.integers(0, 10_000)))
ep_rew, done = 0.0, False
while not done:
action = policy_fn(obs)
obs, r, term, trunc, info = env.step(action)
ep_rew += r
done = term or trunc
rewards.append(ep_rew)
successes += int(term)
return np.array(rewards), successes / n_episodes
vec = DummyVecEnv([lambda: PlanarReachEnv(max_steps=100)])
vec = VecNormalize.load("./models/ppo_reach/vecnormalize.pkl", vec)
vec.training = False
model = PPO.load("./models/ppo_reach/best_model", env=vec)
def trained_policy(obs):
a, _ = model.predict(obs[None, :].astype(np.float32), deterministic=True)
return a[0]
def random_policy(obs):
return np.random.uniform(-MAX_JOINT_DELTA, MAX_JOINT_DELTA, size=2).astype(np.float32)
tr_rew, tr_succ = run_episodes(trained_policy)
rn_rew, rn_succ = run_episodes(random_policy)
print(f"trained: mean reward {tr_rew.mean():7.3f} std {tr_rew.std():6.3f} success {tr_succ:.2%}")
print(f"random: mean reward {rn_rew.mean():7.3f} std {rn_rew.std():6.3f} success {rn_succ:.2%}")
assert tr_rew.mean() > rn_rew.mean() + 5.0, "trained policy did not beat random"
print("validation passed")
The assertion is a guardrail: if the trained mean reward is not at least 5.0 above random, something went wrong and the script fails loudly instead of printing a misleading “success”. Note that trained_policy normalizes the observation through VecNormalize indirectly — actually it does not, so if you get poor numbers here, the cause is almost always observation normalization mismatch. The model.predict call above expects the normalized observation. To be strictly correct, wrap the single-env observation with vec.obs_rms manually. The simpler fix is to run validation through the VecNormalize wrapper itself, which is what Step 3 does. If your numbers from Step 4 look bad but Step 3 looked good, that is the mismatch.
Expected Output
During training with verbose=1 you will see rollout summaries like:
-----------------------------
| time/ | |
| fps | 1842 |
| iterations | 97 |
| time_elapsed | 107 |
| total_timesteps | 200704|
-----------------------------
Per-update lines will show ep_rew_mean climbing from around -180 (random policy, roughly the distance penalty summed over 100 steps) to -35 to -12, and approx_kl staying under 0.02. explained_variance should rise above 0.7 after about 40 updates. EvalCallback prints Eval num_timesteps=... mean_reward=-14.2 +/- 6.1 blocks every 5000 steps.
Step 3 prints something like:
episode length: 100
initial distance: 1.8421
final distance: 0.0412
min distance: 0.0412
saved reach_rollout.png
Step 4 prints:
trained: mean reward -9.874 std 4.512 success 78.00%
random: mean reward -182.331 std 22.107 success 0.00%
validation passed
Success rate of 70–85% at 200k steps is normal. Pushing to 95%+ usually needs 500k steps or a tighter goal_radius schedule.
Common Bugs and Fixes
1. AssertionError: Your environment must inherit from the gymnasium.Env class. SB3 2.x dropped support for the old gym package. If you wrote import gym and subclassed gym.Env, it fails at make_vec_env. Fix: use import gymnasium as gym and spaces.Box(...) from gymnasium. Also confirm reset() returns (obs, info) and step() returns a 5-tuple (obs, reward, terminated, truncated, info).
2. Policy outputs actions outside the action space. You see RuntimeWarning: invalid value encountered or the arm teleports. The Gaussian policy can sample beyond [-0.2, 0.2]. Fix: clip in step() as the code does, and optionally set clip_mean=1.0 and squash_output=False in policy_kwargs. Do not set both squash_output=True and a Box with asymmetric bounds unless you understand the log-prob correction.
3. Training reward stays flat around -180. Two causes. First, n_steps too small relative to batch_size (e.g. n_steps=64, batch_size=256) — SB3 warns and silently truncates the buffer, wasting samples. Fix: keep n_envs * n_steps % batch_size == 0. Second, the reward is purely sparse. If you removed the -dist term, the agent gets no signal until it accidentally hits the goal. Fix: keep the dense term.
4. FileNotFoundError on vecnormalize.pkl or best_model.zip. EvalCallback only saves best_model.zip after the first eval_freq timesteps and only if mean reward improves. If you run fewer than 5000 steps, no file exists. Fix: either raise total_timesteps or call model.save("./models/ppo_reach/final") and point Step 3 at final.zip instead.
5. Validation numbers disagree with training logs. You loaded the model but not the VecNormalize statistics, or you loaded them with training=True so the running mean drifted during evaluation. Fix: VecNormalize.load(...), then set vec.training = False and vec.norm_reward = False before the first reset().
Exercises
Easy — reward shaping sweep. Change goal_radius to 0.1 and add a second dense term -0.01 * np.sum(np.abs(action)) to penalize large joint movements. Retrain for 200k steps and report the new success rate. You should see smoother motions but a lower success rate if the penalty is too strong; find the coefficient where success stays above 75%.
Medium — observation ablation. Remove the cos/sin encoding and feed raw theta values in [-π, π] instead. Retrain and compare ep_rew_mean curves. Then try adding the previous action to the observation as two extra dimensions. Explain in one paragraph why the raw-angle version likely trains slower, using the wrap-around discontinuity as the argument.
Hard — curriculum on goal radius. Implement a GoalCurriculum wrapper that starts with goal_radius=0.3 and multiplies it by 0.8 every 25k timesteps down to 0.03. Use a custom BaseCallback that mutates vec_env.env_method("set_goal_radius", r) on the underlying PlanarReachEnv (you will need to add that setter). Compare total wall-clock time to reach 90% success against the fixed-radius baseline. Report the timestep at which each configuration first crosses 90% on the eval callback.