How to Train a Robot with Reinforcement Learning: A Complete Beginner's Guide

rl-reach/
├── requirements.txt
├── check_setup.py
├── assets/
│   └── reach_arm.xml
├── reach_env.py
├── test_env.py
├── train_ppo.py
├── plot_training.py
├── evaluate.py
├── view_policy.py
├── README.md
├── runs/
│   └── ppo_seed0/        created by train_ppo.py
└── results/              created by evaluate.py

mkdir rl-reach
cd rl-reach
python3.12 -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip

mkdir rl-reach
cd rl-reach
py -3.12 -m venv .venv
.venv\Scripts\Activate.ps1
python -m pip install --upgrade pip

# Python 3.12 or 3.13. NumPy 2.5 requires Python 3.12 or newer.
# Versions checked against PyPI and the projects' release tags on 21 October 2026.
mujoco==3.16.0
gymnasium==1.4.0
stable-baselines3==2.9.0
torch==2.14.1
numpy==2.5.4
matplotlib==3.11.2

pip install torch==2.14.1 --index-url https://download.pytorch.org/whl/cpu

pip install -r requirements.txt

"""check_setup.py: print the computer and package details that matter for this project.

Run it before installing (to see Python, CPU and disk) and again after installing (to see versions).
"""

import importlib
import os
import platform
import shutil
import sys

print(f"Operating system : {platform.platform()}")
print(f"Machine          : {platform.machine()}")
print(f"Python           : {sys.version.split()[0]} ({sys.executable})")
print(f"Logical CPUs     : {os.cpu_count()}")
free_gb = shutil.disk_usage(".").free / 1e9
print(f"Free disk here   : {free_gb:.1f} GB")

if sys.version_info < (3, 12):
    print("WARNING: requirements.txt pins NumPy 2.5, which needs Python 3.12 or newer.")

for name in ["mujoco", "gymnasium", "stable_baselines3", "torch", "numpy", "matplotlib"]:
    try:
        module = importlib.import_module(name)
        print(f"{name:<17}: {getattr(module, '__version__', 'unknown version')}")
    except ImportError:
        print(f"{name:<17}: not installed")

try:
    import torch

    print(f"CUDA available   : {torch.cuda.is_available()}")
    if torch.cuda.is_available():
        print(f"CUDA device      : {torch.cuda.get_device_name(0)}")
    mps = getattr(torch.backends, "mps", None)
    print(f"Apple MPS        : {bool(mps and mps.is_available())}")
    print("This project trains on the CPU either way (device='cpu' in train_ppo.py).")
except ImportError:
    pass

python check_setup.py

python -c "import mujoco; m = mujoco.MjModel.from_xml_string('<mujoco/>'); print(mujoco.__version__)"

<!--
  reach_arm.xml: a three-joint robot arm for the TechniaHQRobot reinforcement learning tutorial.

  Frames: world z points up, x points forward. Lengths are in metres, masses in kilograms,
  angles in radians. All joint angles are zero when the arm points straight up.

  The arm is a teaching model, not a copy of a commercial robot. Its proportions are
  roughly those of a small desktop arm.
-->
<mujoco model="reach_arm">
  <compiler angle="radian" autolimits="true"/>

  <!-- 2 ms physics step. The environment applies one action every 20 physics steps (25 Hz). -->
  <option timestep="0.002" gravity="0 0 -9.81" integrator="implicitfast"/>

  <default>
    <!-- Passive joint damping and rotor inertia (armature) keep the joints numerically well behaved. -->
    <joint type="hinge" damping="0.5" armature="0.01"/>
    <!-- Arm geoms (contype 2) collide with the floor (conaffinity 2) but not with each other. -->
    <geom contype="2" conaffinity="1" density="0" rgba="0.25 0.45 0.65 1"/>
    <!-- Position servos: the policy moves each servo's setpoint; the servo produces torque. -->
    <position kp="100" dampratio="1" inheritrange="1"/>
  </default>

  <visual>
    <global offwidth="640" offheight="480"/>
  </visual>

  <worldbody>
    <light pos="0 -0.6 1.5" dir="0 0.4 -1"/>
    <geom name="floor" type="plane" size="1 1 0.05" rgba="0.9 0.9 0.88 1" contype="1" conaffinity="2"/>

    <!-- Fixed pedestal. A body without a joint is welded to the world. -->
    <body name="pedestal">
      <geom name="pedestal" type="cylinder" pos="0 0 0.05" size="0.06 0.05" rgba="0.3 0.3 0.3 1"/>

      <!-- Joint 1: base yaw turns the whole arm about the vertical axis. -->
      <body name="turntable" pos="0 0 0.1">
        <joint name="base_yaw" axis="0 0 1" range="-2.6 2.6"/>
        <geom name="turntable" type="cylinder" pos="0 0 0.025" size="0.045 0.025" mass="0.5"/>

        <!-- Joint 2: shoulder pitch tilts the upper arm forward and back. -->
        <body name="upper_arm" pos="0 0 0.05">
          <joint name="shoulder_pitch" axis="0 1 0" range="-1.6 1.6"/>
          <geom name="upper_arm" type="capsule" fromto="0 0 0 0 0 0.3" size="0.03" mass="1.0"/>

          <!-- Joint 3: elbow pitch bends the forearm relative to the upper arm. -->
          <body name="forearm" pos="0 0 0.3">
            <joint name="elbow_pitch" axis="0 1 0" range="-2.4 2.4"/>
            <geom name="forearm" type="capsule" fromto="0 0 0 0 0 0.25" size="0.025" mass="0.6"/>
            <!-- The end effector is a site: a named point with no mass or collision. -->
            <site name="end_effector" pos="0 0 0.25" size="0.015" rgba="1 0.8 0.1 1"/>
          </body>
        </body>
      </body>
    </body>

    <!-- The target is a mocap body: the code places it, physics never moves it, and it collides with nothing. -->
    <body name="target" mocap="true" pos="0.3 0 0.3">
      <geom name="target" type="sphere" size="0.02" contype="0" conaffinity="0" rgba="0.9 0.2 0.2 0.6"/>
    </body>
  </worldbody>

  <actuator>
    <!-- inheritrange="1" copies each joint range into the servo's control range. -->
    <position name="base_yaw" joint="base_yaw" forcerange="-10 10"/>
    <position name="shoulder_pitch" joint="shoulder_pitch" forcerange="-25 25"/>
    <position name="elbow_pitch" joint="elbow_pitch" forcerange="-12 12"/>
  </actuator>
</mujoco>

python -m mujoco.viewer --mjcf=assets/reach_arm.xml

"""reach_env.py: a Gymnasium environment in which a three-joint MuJoCo arm reaches a random target.

The policy controls position servos. Each action moves the three servo setpoints by at most
`max_target_step` radians. The servos then produce the joint torques during 20 physics steps.
"""

from pathlib import Path

import gymnasium as gym
import mujoco
import numpy as np
from gymnasium import spaces

MODEL_PATH = Path(__file__).resolve().parent / "assets" / "reach_arm.xml"

# Joint ranges used only to place targets. They sit inside the hard joint limits so that every
# target can be reached without pressing a joint against its stop.
TARGET_JOINT_LOW = np.array([-2.0, -1.2, -2.0])
TARGET_JOINT_HIGH = np.array([2.0, 1.2, 2.0])


class ReachEnv(gym.Env):
    """Move the end effector to within `success_threshold` metres of the target and stay there."""

    metadata = {"render_modes": ["human", "rgb_array"], "render_fps": 25}

    def __init__(
        self,
        render_mode=None,
        max_episode_steps=100,
        frame_skip=20,
        max_target_step=0.05,
        success_threshold=0.02,
        action_cost=0.01,
        success_bonus=0.1,
    ):
        if render_mode is not None and render_mode not in self.metadata["render_modes"]:
            raise ValueError(f"Unsupported render_mode {render_mode!r}")
        self.render_mode = render_mode

        # The model and its simulation state. `_scratch` is a second state used only to compute
        # target positions, so target sampling never disturbs the running simulation.
        self.model = mujoco.MjModel.from_xml_path(str(MODEL_PATH))
        self.data = mujoco.MjData(self.model)
        self._scratch = mujoco.MjData(self.model)
        assert self.model.nq == 3 and self.model.nu == 3, "The model must have three joints and three actuators."

        # Timing: 20 physics steps of 2 ms give one control step of 40 ms (25 Hz).
        self.frame_skip = frame_skip
        self.control_dt = self.model.opt.timestep * frame_skip
        self.max_episode_steps = max_episode_steps

        # Task and reward settings.
        self.max_target_step = max_target_step
        self.success_threshold = success_threshold
        self.action_cost = action_cost
        self.success_bonus = success_bonus

        # Look up names once. Integer ids are faster than names inside the step loop.
        self.ee_site = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_SITE, "end_effector")
        target_body = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_BODY, "target")
        self.target_mocap = int(self.model.body_mocapid[target_body])
        self.ctrl_low = self.model.actuator_ctrlrange[:, 0].copy()
        self.ctrl_high = self.model.actuator_ctrlrange[:, 1].copy()

        # Actions: one number per joint in [-1, 1], scaled to a setpoint change in step().
        self.action_space = spaces.Box(low=-1.0, high=1.0, shape=(3,), dtype=np.float32)
        # Observations: 18 scaled numbers, described in _get_obs(), clipped to [-10, 10].
        self.observation_space = spaces.Box(low=-10.0, high=10.0, shape=(18,), dtype=np.float32)

        self.target_pos = np.zeros(3)
        self.steps = 0
        self.first_reach_step = None
        self._renderer = None
        self._viewer = None

    # ----------------------------------------------------------------- Gymnasium API

    def reset(self, *, seed=None, options=None):
        # Seeds self.np_random. Every random choice below must use self.np_random.
        super().reset(seed=seed)
        options = options or {}

        mujoco.mj_resetData(self.model, self.data)
        start = self.np_random.uniform(-0.1, 0.1, size=3)
        self.data.qpos[:] = start
        self.data.qvel[:] = 0.0
        self.data.ctrl[:] = start  # The servos begin by holding the starting pose.

        if "target" in options:
            self.target_pos = np.asarray(options["target"], dtype=np.float64).copy()
        else:
            self.target_pos = self._sample_target()
        self.data.mocap_pos[self.target_mocap] = self.target_pos

        # mj_forward computes positions such as site_xpos from qpos without advancing time.
        mujoco.mj_forward(self.model, self.data)
        self.steps = 0
        self.first_reach_step = None

        distance = self._distance()
        if self.render_mode == "human":
            self.render()
        return self._get_obs(), self._get_info(distance, {})

    def step(self, action):
        action = np.clip(np.asarray(action, dtype=np.float64), -1.0, 1.0)

        # 1. Apply the action: move each servo setpoint, then keep it inside the joint range.
        self.data.ctrl[:] = np.clip(
            self.data.ctrl + action * self.max_target_step, self.ctrl_low, self.ctrl_high
        )

        # 2. Advance the physics by one control period, then refresh positions for the new state.
        bad_before = self.data.warning[int(mujoco.mjtWarning.mjWARN_BADQACC)].number
        mujoco.mj_step(self.model, self.data, nstep=self.frame_skip)
        mujoco.mj_forward(self.model, self.data)
        self.steps += 1

        # 3. Score the new state.
        distance = self._distance()
        reward, parts = self._reward(distance, action)
        if self.first_reach_step is None and distance < self.success_threshold:
            self.first_reach_step = self.steps

        # 4. Decide whether the episode ends.
        # terminated: the episode reached a terminal state. Here that only happens when the
        # simulation became numerically unstable (MuJoCo resets the state and counts a warning).
        unstable = (
            self.data.warning[int(mujoco.mjtWarning.mjWARN_BADQACC)].number > bad_before
            or not np.all(np.isfinite(self.data.qpos))
            or not np.all(np.isfinite(self.data.qvel))
        )
        terminated = bool(unstable)
        # truncated: the time limit ended the episode, not the task.
        truncated = bool(self.steps >= self.max_episode_steps)

        info = self._get_info(distance, parts)
        if terminated or truncated:
            # Success is judged on the final state, so passing through the target does not count.
            info["is_success"] = bool(distance < self.success_threshold and not terminated)

        if self.render_mode == "human":
            self.render()
        return self._get_obs(), float(reward), terminated, truncated, info

    def render(self):
        if self.render_mode == "rgb_array":
            if self._renderer is None:
                self._renderer = mujoco.Renderer(self.model, height=480, width=640)
            self._renderer.update_scene(self.data)
            return self._renderer.render()
        if self.render_mode == "human":
            if self._viewer is None:
                import mujoco.viewer

                self._viewer = mujoco.viewer.launch_passive(self.model, self.data)
            if self._viewer.is_running():
                self._viewer.sync()
        return None

    def close(self):
        if self._renderer is not None:
            self._renderer.close()
            self._renderer = None
        if self._viewer is not None:
            self._viewer.close()
            self._viewer = None

    # ----------------------------------------------------------------- task details

    def _sample_target(self):
        """Pick a random joint configuration and use its end-effector position as the target.

        Targets made this way are reachable by construction. Points too close to the floor or to
        the pedestal are rejected.
        """
        for _ in range(100):
            q = self.np_random.uniform(TARGET_JOINT_LOW, TARGET_JOINT_HIGH)
            self._scratch.qpos[:] = q
            mujoco.mj_kinematics(self.model, self._scratch)
            point = self._scratch.site_xpos[self.ee_site].copy()
            if point[2] > 0.08 and np.linalg.norm(point[:2]) > 0.10:
                return point
        raise RuntimeError("Could not sample a valid target in 100 attempts.")

    def _distance(self):
        return float(np.linalg.norm(self.data.site_xpos[self.ee_site] - self.target_pos))

    def _reward(self, distance, action):
        parts = {
            "reward_distance": -distance,
            "reward_action": -self.action_cost * float(np.sum(np.square(action))),
            "reward_success": self.success_bonus if distance < self.success_threshold else 0.0,
        }
        return sum(parts.values()), parts

    def _get_obs(self):
        """18 numbers, each divided by a fixed scale so most values fall roughly within [-1, 1].

        0-2   joint angles / pi                  (rad)
        3-5   joint velocities / 10              (rad/s)
        6-8   servo setpoints / pi               (rad)
        9-11  end-effector position / 0.6        (m)
        12-14 target position / 0.6              (m)
        15-17 (target - end effector) / 0.6      (m)
        """
        ee = self.data.site_xpos[self.ee_site]
        obs = np.concatenate(
            [
                self.data.qpos / np.pi,
                self.data.qvel / 10.0,
                self.data.ctrl / np.pi,
                ee / 0.6,
                self.target_pos / 0.6,
                (self.target_pos - ee) / 0.6,
            ]
        )
        # Clipping only matters if something goes badly wrong, for example a violent collision.
        return np.clip(obs, -10.0, 10.0).astype(np.float32)

    def _get_info(self, distance, parts):
        return {
            "distance": distance,
            "steps_to_reach": -1 if self.first_reach_step is None else self.first_reach_step,
            **parts,
        }

"""test_env.py: check ReachEnv before any training.

Run:  python test_env.py
      python test_env.py --image arm.png   (also saves a rendered frame; needs an OpenGL backend)

The script stops with an AssertionError at the first failed check.
"""

import argparse

import numpy as np
from gymnasium.utils.env_checker import check_env

from reach_env import ReachEnv

parser = argparse.ArgumentParser()
parser.add_argument("--image", help="save one rendered frame to this PNG file")
args = parser.parse_args()

# 1. Gymnasium's own API checks: spaces, reset/step return types, seeding and determinism.
env = ReachEnv()
check_env(env, skip_render_check=True)
print("Gymnasium API check passed.")

# 2. Seeding: the same seed must give the same start and target.
obs_a, _ = env.reset(seed=42)
obs_b, _ = env.reset(seed=42)
assert np.array_equal(obs_a, obs_b), "reset(seed=42) twice gave different observations"
obs, info = env.reset(seed=42)
assert env.observation_space.contains(obs), "first observation is outside the observation space"
print(f"Observation shape {obs.shape}, dtype {obs.dtype}. Target {np.round(env.target_pos, 3)} m.")
print(f"Initial distance to target: {info['distance']:.3f} m")

# 3. One episode of random actions. Print every 10th step and check every value.
low = env.model.jnt_range[:, 0] - 0.05  # MuJoCo joint limits are soft, so allow a small overshoot.
high = env.model.jnt_range[:, 1] + 0.05
print(f"{'step':>4} {'distance':>9} {'r_dist':>8} {'r_action':>9} {'r_success':>9} {'reward':>8}")
for t in range(1, 1000):
    action = env.action_space.sample()
    obs, reward, terminated, truncated, info = env.step(action)
    assert np.all(np.isfinite(obs)) and np.isfinite(reward), f"non-finite value at step {t}"
    assert np.all(env.data.qpos >= low) and np.all(env.data.qpos <= high), f"joint outside its range at step {t}"
    parts = info["reward_distance"] + info["reward_action"] + info["reward_success"]
    assert abs(parts - reward) < 1e-9, "reward components do not add up to the reward"
    if t % 10 == 0 or terminated or truncated:
        print(f"{t:>4} {info['distance']:9.3f} {info['reward_distance']:8.3f} {info['reward_action']:9.4f} "
              f"{info['reward_success']:9.2f} {reward:8.3f}")
    if terminated or truncated:
        break

# 4. The episode must end by truncation at exactly max_episode_steps and report success.
assert truncated and not terminated, "episode should end by time limit, not termination"
assert t == env.max_episode_steps, f"episode ended after {t} steps, expected {env.max_episode_steps}"
assert "is_success" in info, "final info must contain is_success"
print(f"Episode truncated after {t} steps. is_success={info['is_success']}.")

# 5. Zero action: the servos should hold the pose, so the arm barely moves.
env.reset(seed=7)
start = env.data.site_xpos[env.ee_site].copy()
for _ in range(25):
    env.step(np.zeros(3, dtype=np.float32))
drift = np.linalg.norm(env.data.site_xpos[env.ee_site] - start)
print(f"End-effector drift after 1 s of zero actions: {drift * 1000:.1f} mm")
assert drift < 0.02, "the arm drifts with zero action; check actuator gains and gravity"

# 6. Extreme actions must keep the servo setpoints inside the actuator control range.
env.reset(seed=3)
for _ in range(100):
    env.step(np.ones(3, dtype=np.float32))
assert np.all(env.data.ctrl <= env.ctrl_high + 1e-12), "setpoint exceeded the control range"
print(f"Setpoints after 100 steps of +1 actions: {np.round(env.data.ctrl, 3)} (limits {env.ctrl_high})")
env.close()

# 7. Optional rendering check.
if args.image:
    import matplotlib.pyplot as plt

    render_env = ReachEnv(render_mode="rgb_array")
    render_env.reset(seed=42)
    frame = render_env.render()
    plt.imsave(args.image, frame)
    render_env.close()
    print(f"Saved a {frame.shape[1]}x{frame.shape[0]} frame to {args.image}")

print("All environment checks passed.")

python test_env.py

Gymnasium API check passed.
Observation shape (18,), dtype float32. Target [ 0.152 -0.03   0.479] m.
Initial distance to target: 0.263 m
step  distance   r_dist  r_action r_success   reward
  10     0.266   -0.266   -0.0095      0.00   -0.276
  20     0.299   -0.299   -0.0097      0.00   -0.309
  30     0.286   -0.286   -0.0084      0.00   -0.295
  40     0.305   -0.305   -0.0113      0.00   -0.317
  50     0.322   -0.322   -0.0104      0.00   -0.333
  60     0.365   -0.365   -0.0122      0.00   -0.377
  70     0.400   -0.400   -0.0050      0.00   -0.405
  80     0.368   -0.368   -0.0113      0.00   -0.379
  90     0.389   -0.389   -0.0180      0.00   -0.407
 100     0.428   -0.428   -0.0011      0.00   -0.429
Episode truncated after 100 steps. is_success=False.
End-effector drift after 1 s of zero actions: 2.3 mm
Setpoints after 100 steps of +1 actions: [2.6 1.6 2.4] (limits [2.6 1.6 2.4])
All environment checks passed.

python test_env.py --image arm.png
MUJOCO_GL=egl python test_env.py --image arm.png   # Linux without a display

r = − d  −  0.01 × (a₁² + a₂² + a₃²)  +  0.1 × [d < 0.02 m]

d   distance from the tip to the target, in metres
a   the three actions, each between −1 and 1
[ ] 1 when the condition is true, otherwise 0

"""train_ppo.py: train a PPO policy for ReachEnv with Stable-Baselines3.

Run:     python train_ppo.py --timesteps 500000 --seed 0
Resume:  python train_ppo.py --resume runs/ppo_seed0/checkpoints/ppo_reach_200000_steps.zip --timesteps 300000

Outputs, under runs/ppo_seed<seed>/ unless --run-dir is given:
  config.json        settings, seeds and package versions (config_from_<steps>.json after a resume)
  logs/progress.csv  training metrics written by SB3's logger (logs_from_<steps>/ after a resume)
  monitor/*.csv      one row per finished training episode
  eval/evaluations.npz   periodic evaluation results (eval_from_<steps>/ after a resume)
  checkpoints/       intermediate models
  best_model/best_model.zip   highest mean evaluation reward so far
  final_model.zip    model at the end of training
"""

import argparse
import json
import platform
import sys
from pathlib import Path

import gymnasium
import mujoco
import numpy as np
import stable_baselines3
import torch
from stable_baselines3 import PPO
from stable_baselines3.common.callbacks import CallbackList, CheckpointCallback, EvalCallback
from stable_baselines3.common.env_checker import check_env
from stable_baselines3.common.env_util import make_vec_env
from stable_baselines3.common.logger import configure
from stable_baselines3.common.utils import set_random_seed

from reach_env import ReachEnv

# Starting values for this task, not universal settings. Change one at a time and re-evaluate.
PPO_CONFIG = {
    "learning_rate": 3e-4,  # step size of each network update
    "n_steps": 1024,        # steps collected per environment before each update
    "batch_size": 256,      # samples per gradient step; 1024 x 4 envs = 4096 = 16 batches
    "n_epochs": 10,         # passes over each batch of collected experience
    "gamma": 0.99,          # discount: how much later rewards count
    "gae_lambda": 0.95,     # bias/variance trade-off when estimating advantages
    "clip_range": 0.2,      # how far one update may move the policy
    "ent_coef": 0.0,        # bonus for keeping actions random; 0 relies on the policy's own noise
    "vf_coef": 0.5,         # weight of the value-function loss
    "max_grad_norm": 0.5,   # cap on gradient size, for stability
    "policy_kwargs": {"net_arch": {"pi": [64, 64], "vf": [64, 64]}},
}


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--timesteps", type=int, default=500_000,
                        help="environment steps to train (when resuming: additional steps)")
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--n-envs", type=int, default=4, help="parallel copies of the environment")
    parser.add_argument("--run-dir", default=None)
    parser.add_argument("--resume", default=None, help="path to a saved .zip model to continue training")
    args = parser.parse_args()

    run_dir = Path(args.run_dir or f"runs/ppo_seed{args.seed}")
    run_dir.mkdir(parents=True, exist_ok=True)

    # Seed Python, NumPy and PyTorch. Results can still differ across computers and versions.
    set_random_seed(args.seed)

    # Check the environment against the API Stable-Baselines3 expects.
    check = ReachEnv()
    check_env(check, warn=True)
    check.close()

    # Training environments: n_envs copies stepped one after another in this process.
    # Monitor records each finished episode, including is_success and the final distance.
    train_env = make_vec_env(
        ReachEnv,
        n_envs=args.n_envs,
        seed=args.seed,
        monitor_dir=str(run_dir / "monitor"),
        # When resuming, append to the existing monitor files instead of overwriting them.
        monitor_kwargs={"info_keywords": ("is_success", "distance"), "override_existing": args.resume is None},
    )
    # A separate evaluation environment with a different seed, so evaluation targets differ.
    eval_env = make_vec_env(ReachEnv, n_envs=1, seed=args.seed + 10_000)

    if args.resume:
        model = PPO.load(args.resume, env=train_env, device="cpu")
        print(f"Resumed from {args.resume} at {model.num_timesteps} steps.")
    else:
        model = PPO("MlpPolicy", train_env, seed=args.seed, device="cpu", verbose=1, **PPO_CONFIG)

    # Write metrics to the screen and to a CSV file. A resumed run writes to new folders
    # (logs_from_<steps>, eval_from_<steps>) so the earlier files are kept.
    suffix = "" if not args.resume else f"_from_{model.num_timesteps}"
    model.set_logger(configure(str(run_dir / f"logs{suffix}"), ["stdout", "csv"]))

    # Callback frequencies count calls to env.step() on the vectorized env, and each call
    # advances all n_envs copies. Dividing by n_envs makes the numbers below environment steps.
    checkpoint = CheckpointCallback(
        save_freq=max(50_000 // args.n_envs, 1),
        save_path=str(run_dir / "checkpoints"),
        name_prefix="ppo_reach",
    )
    evaluation = EvalCallback(
        eval_env,
        n_eval_episodes=20,
        eval_freq=max(25_000 // args.n_envs, 1),
        log_path=str(run_dir / f"eval{suffix}"),
        best_model_save_path=str(run_dir / "best_model"),  # after a resume, "best" restarts from the resume point
        deterministic=True,
    )

    config = {
        "args": vars(args),
        "ppo": PPO_CONFIG,
        "env": {
            "max_episode_steps": 100, "frame_skip": 20, "max_target_step": 0.05,
            "success_threshold": 0.02, "action_cost": 0.01, "success_bonus": 0.1,
        },
        "versions": {
            "python": sys.version.split()[0], "platform": platform.platform(),
            "mujoco": mujoco.__version__, "gymnasium": gymnasium.__version__,
            "stable_baselines3": stable_baselines3.__version__, "torch": torch.__version__,
            "numpy": np.__version__,
        },
    }
    (run_dir / f"config{suffix}.json").write_text(json.dumps(config, indent=2))

    try:
        model.learn(
            total_timesteps=args.timesteps,
            callback=CallbackList([checkpoint, evaluation]),
            reset_num_timesteps=args.resume is None,
        )
    except KeyboardInterrupt:
        # Ctrl+C stops training; the model trained so far is still saved below.
        print("Training interrupted.")
    finally:
        model.save(run_dir / "final_model")
        print(f"Saved {run_dir / 'final_model.zip'} after {model.num_timesteps} steps.")
        train_env.close()
        eval_env.close()


if __name__ == "__main__":
    main()

python train_ppo.py --timesteps 500000 --seed 0

python train_ppo.py --resume runs/ppo_seed0/checkpoints/ppo_reach_200000_steps.zip --timesteps 300000

"""plot_training.py: plot the metrics recorded during training.

Run:  python plot_training.py --run-dir runs/ppo_seed0
Writes <run-dir>/training_curves.png and prints the last logged values.
"""

import argparse
import csv
from pathlib import Path

import matplotlib

matplotlib.use("Agg")  # write files without opening a window
import matplotlib.pyplot as plt
import numpy as np


def read_progress(run_dir):
    """Merge logs*/progress.csv files (one per run or resume) into columns keyed by metric name."""
    rows = []
    for path in sorted(run_dir.glob("logs*/progress.csv")):
        with open(path, newline="") as f:
            rows.extend(csv.DictReader(f))
    rows.sort(key=lambda r: float(r.get("time/total_timesteps") or 0))
    return rows


def column(rows, key):
    xs, ys = [], []
    for r in rows:
        if r.get(key) not in (None, "") and r.get("time/total_timesteps") not in (None, ""):
            xs.append(float(r["time/total_timesteps"]))
            ys.append(float(r[key]))
    return np.array(xs), np.array(ys)


def read_monitor_distances(run_dir):
    """Final distance of every training episode, ordered by wall-clock end time."""
    episodes = []
    for path in sorted((run_dir / "monitor").glob("*.monitor.csv")):
        with open(path, newline="") as f:
            next(f)  # first line is a JSON comment written by Monitor
            for r in csv.DictReader(f):
                episodes.append((float(r["t"]), float(r["distance"])))
    episodes.sort()
    return np.array([d for _, d in episodes])


def read_evaluations(run_dir):
    steps, success, reward = [], [], []
    for path in sorted(run_dir.glob("eval*/evaluations.npz")):
        data = np.load(path)
        steps.extend(data["timesteps"])
        reward.extend(data["results"].mean(axis=1))
        if "successes" in data.files:
            success.extend(data["successes"].mean(axis=1))
    return np.array(steps), np.array(reward), np.array(success)


def moving_average(values, window=100):
    if len(values) < window:
        return values
    return np.convolve(values, np.ones(window) / window, mode="valid")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--run-dir", default="runs/ppo_seed0")
    args = parser.parse_args()
    run_dir = Path(args.run_dir)

    rows = read_progress(run_dir)
    if not rows:
        raise SystemExit(f"No progress.csv found under {run_dir}. Has training started?")
    eval_steps, eval_reward, eval_success = read_evaluations(run_dir)
    distances = read_monitor_distances(run_dir)

    fig, axes = plt.subplots(2, 3, figsize=(15, 8))
    ax = axes[0, 0]
    ax.plot(*column(rows, "rollout/ep_rew_mean"), label="training (last 100 episodes)")
    if len(eval_steps):
        ax.plot(eval_steps, eval_reward, "o-", label="evaluation (deterministic)")
    ax.set(title="Episode reward", xlabel="environment steps", ylabel="return")
    ax.legend()

    ax = axes[0, 1]
    ax.plot(*column(rows, "rollout/success_rate"), label="training (stochastic actions)")
    if len(eval_success):
        ax.plot(eval_steps, eval_success, "o-", label="evaluation (deterministic)")
    ax.set(title="Success rate (final distance < 2 cm)", xlabel="environment steps", ylim=(-0.05, 1.05))
    ax.legend()

    ax = axes[0, 2]
    if len(distances):
        ax.plot(moving_average(distances) * 1000)
        ax.axhline(20, color="grey", linestyle="--", label="success threshold")
        ax.legend()
    ax.set(title="Final distance, training episodes (100-episode average)", xlabel="episode", ylabel="mm")

    ax = axes[1, 0]
    ax.plot(*column(rows, "train/value_loss"))
    ax.set(title="Value loss", xlabel="environment steps", yscale="log")

    ax = axes[1, 1]
    ax.plot(*column(rows, "train/policy_gradient_loss"), label="policy gradient loss")
    ax.plot(*column(rows, "train/approx_kl"), label="approximate KL")
    ax.plot(*column(rows, "train/clip_fraction"), label="clip fraction")
    ax.set(title="Policy update size", xlabel="environment steps")
    ax.legend()

    ax = axes[1, 2]
    ax.plot(*column(rows, "time/fps"))
    ax.set(title="Training speed", xlabel="environment steps", ylabel="steps per second")

    fig.tight_layout()
    out = run_dir / "training_curves.png"
    fig.savefig(out, dpi=120)
    print(f"Saved {out}")

    last = rows[-1]
    for key in ["time/total_timesteps", "rollout/ep_rew_mean", "rollout/success_rate",
                "train/value_loss", "train/approx_kl", "train/std", "time/fps"]:
        print(f"{key:<24} {last.get(key, '')}")


if __name__ == "__main__":
    main()

python plot_training.py --run-dir runs/ppo_seed0

"""evaluate.py: measure a trained policy and two baselines on the same targets.

Run:  python evaluate.py --model runs/ppo_seed0/best_model/best_model.zip --episodes 200
Writes results/episodes.csv (one row per episode and policy) and results/summary.json.

Success is fixed before evaluation: the end effector is within 0.02 m of the target at the
final step of the 100-step episode.
"""

import argparse
import csv
import json
import math
from pathlib import Path

import mujoco
import numpy as np

from reach_env import ReachEnv

EVAL_SEED = 2026  # different from the training seeds; episode i uses seed EVAL_SEED + i


class RandomPolicy:
    """Uniform random actions: the floor any learned policy must clear."""

    def __init__(self, env, seed):
        self.env = env
        self.env.action_space.seed(seed)

    def __call__(self, obs):
        return self.env.action_space.sample()


class JacobianController:
    """Hand-designed baseline: damped least-squares inverse kinematics.

    Uses the simulator's exact Jacobian, which the learned policy never sees.
    dq = J^T (J J^T + lambda^2 I)^-1 * error, then scaled to the [-1, 1] action range.
    """

    def __init__(self, env, gain=0.5, damping=0.05):
        self.env, self.gain, self.damping = env, gain, damping
        self.jacp = np.zeros((3, env.model.nv))
        self.jacr = np.zeros((3, env.model.nv))

    def __call__(self, obs):
        env = self.env
        mujoco.mj_jacSite(env.model, env.data, self.jacp, self.jacr, env.ee_site)
        error = env.target_pos - env.data.site_xpos[env.ee_site]
        jac = self.jacp
        dq = jac.T @ np.linalg.solve(jac @ jac.T + self.damping**2 * np.eye(3), error)
        return np.clip(self.gain * dq / env.max_target_step, -1.0, 1.0).astype(np.float32)


class PPOPolicy:
    def __init__(self, path, deterministic):
        from stable_baselines3 import PPO  # imported here so the baselines run without PyTorch

        self.model = PPO.load(path, device="cpu")
        self.deterministic = deterministic

    def __call__(self, obs):
        action, _ = self.model.predict(obs, deterministic=self.deterministic)
        return action


def run_episodes(name, policy_factory, episodes):
    env = ReachEnv()
    policy = policy_factory(env)
    rows = []
    for i in range(episodes):
        obs, info = env.reset(seed=EVAL_SEED + i)  # same seed -> same start pose and target for every policy
        episode_return, action_sq, terminated, truncated = 0.0, 0.0, False, False
        while not (terminated or truncated):
            action = policy(obs)
            obs, reward, terminated, truncated, info = env.step(action)
            episode_return += reward
            action_sq += float(np.sum(np.square(action)))
        rows.append({
            "policy": name,
            "episode": i,
            "success": int(info["is_success"]),
            "final_distance_m": info["distance"],
            "steps_to_reach": info["steps_to_reach"],
            "episode_length": env.steps,
            "return": episode_return,
            "mean_squared_action": action_sq / env.steps,
        })
    env.close()
    return rows


def wilson_interval(successes, n, z=1.96):
    """95% confidence interval for a success rate measured on n episodes."""
    if n == 0:
        return (0.0, 0.0)
    p = successes / n
    centre = (p + z * z / (2 * n)) / (1 + z * z / n)
    half = z * math.sqrt(p * (1 - p) / n + z * z / (4 * n * n)) / (1 + z * z / n)
    return (max(0.0, centre - half), min(1.0, centre + half))


def summarise(rows):
    n = len(rows)
    successes = sum(r["success"] for r in rows)
    final = np.array([r["final_distance_m"] for r in rows])
    reached = [r["steps_to_reach"] for r in rows if r["steps_to_reach"] > 0]
    low, high = wilson_interval(successes, n)
    return {
        "episodes": n,
        "success_rate": successes / n,
        "success_rate_95ci": [low, high],
        "final_distance_mean_m": float(final.mean()),
        "final_distance_std_m": float(final.std(ddof=1)) if n > 1 else 0.0,
        "final_distance_median_m": float(np.median(final)),
        "episode_length_mean": float(np.mean([r["episode_length"] for r in rows])),
        "steps_to_reach_mean": float(np.mean(reached)) if reached else None,
        "return_mean": float(np.mean([r["return"] for r in rows])),
    }


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", help="trained PPO .zip; omit to evaluate only the baselines")
    parser.add_argument("--episodes", type=int, default=200)
    parser.add_argument("--stochastic", action="store_true", help="sample PPO actions instead of using the mean")
    parser.add_argument("--out", default="results")
    args = parser.parse_args()

    policies = {
        "random": lambda env: RandomPolicy(env, seed=EVAL_SEED),
        "jacobian_controller": lambda env: JacobianController(env),
    }
    if args.model:
        policies["ppo"] = lambda env: PPOPolicy(args.model, deterministic=not args.stochastic)

    all_rows, summary = [], {}
    for name, factory in policies.items():
        rows = run_episodes(name, factory, args.episodes)
        all_rows.extend(rows)
        summary[name] = summarise(rows)

    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)
    with open(out / "episodes.csv", "w", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=list(all_rows[0].keys()))
        writer.writeheader()
        writer.writerows(all_rows)
    meta = {"eval_seed": EVAL_SEED, "model": args.model, "deterministic_ppo": not args.stochastic}
    (out / "summary.json").write_text(json.dumps({"settings": meta, "policies": summary}, indent=2))

    print(f"{'policy':<20} {'success':>8} {'95% CI':>15} {'final dist (mm)':>18} {'steps to reach':>15}")
    for name, s in summary.items():
        ci = f"{s['success_rate_95ci'][0]:.2f}-{s['success_rate_95ci'][1]:.2f}"
        dist = f"{s['final_distance_mean_m'] * 1000:.1f} ± {s['final_distance_std_m'] * 1000:.1f}"
        steps = "-" if s["steps_to_reach_mean"] is None else f"{s['steps_to_reach_mean']:.1f}"
        print(f"{name:<20} {s['success_rate']:>8.2f} {ci:>15} {dist:>18} {steps:>15}")
    print(f"Wrote {out / 'episodes.csv'} and {out / 'summary.json'}")


if __name__ == "__main__":
    main()

python evaluate.py --model runs/ppo_seed0/best_model/best_model.zip --episodes 200

"""view_policy.py: watch a trained policy in MuJoCo's interactive viewer.

Run (Windows, Linux):  python view_policy.py --model runs/ppo_seed0/best_model/best_model.zip
Run (macOS):           mjpython view_policy.py --model runs/ppo_seed0/best_model/best_model.zip

Close the viewer window to stop. Needs a display; it will not work over a plain SSH session.
"""

import argparse
import time

from stable_baselines3 import PPO

from reach_env import ReachEnv

parser = argparse.ArgumentParser()
parser.add_argument("--model", required=True)
parser.add_argument("--episodes", type=int, default=5)
args = parser.parse_args()

model = PPO.load(args.model, device="cpu")
env = ReachEnv(render_mode="human")
try:
    for episode in range(args.episodes):
        obs, info = env.reset(seed=500 + episode)
        terminated = truncated = False
        while not (terminated or truncated):
            start = time.time()
            action, _ = model.predict(obs, deterministic=True)
            obs, reward, terminated, truncated, info = env.step(action)
            # Sleep so one control step takes about 40 ms of wall time (real-time playback).
            time.sleep(max(0.0, env.control_dt - (time.time() - start)))
            if env._viewer is not None and not env._viewer.is_running():
                raise SystemExit
        print(f"Episode {episode}: final distance {info['distance'] * 1000:.1f} mm, success {info['is_success']}")
finally:
    env.close()

python view_policy.py --model runs/ppo_seed0/best_model/best_model.zip
mjpython view_policy.py --model runs/ppo_seed0/best_model/best_model.zip   # macOS

for seed in 0 1 2; do python train_ppo.py --seed $seed; done
for seed in 0 1 2; do python evaluate.py --model runs/ppo_seed$seed/best_model/best_model.zip --out results/seed$seed; done

# rl-reach: train a robot arm to reach with reinforcement learning

Companion project for the TechniaHQRobot guide
"How to Train a Robot with Reinforcement Learning: A Complete Beginner's Guide"
(https://techniahqrobot.com/guides/train-robot-reinforcement-learning/).

A three-joint MuJoCo arm learns, with PPO from Stable-Baselines3, to move its tip within 2 cm
of a random target and stay there. Everything runs on a CPU, in simulation.

## Files

| File | Purpose |
| --- | --- |
| requirements.txt | Pinned package versions (Python 3.12 or 3.13) |
| check_setup.py | Prints Python, CPU, disk and package versions |
| assets/reach_arm.xml | MuJoCo model of the arm, servos and target |
| reach_env.py | Gymnasium environment |
| test_env.py | Environment checks; run before training |
| train_ppo.py | PPO training with logs, checkpoints and evaluation |
| plot_training.py | Training curves from the logs |
| evaluate.py | Success rate and distance for PPO, a random policy and a Jacobian controller |
| view_policy.py | Watch a trained policy in the MuJoCo viewer |
| executed/ | Outputs recorded when this project was checked (see below) |

## Commands

```bash
python3.12 -m venv .venv
source .venv/bin/activate            # Windows: .venv\Scripts\Activate.ps1
python -m pip install --upgrade pip
pip install torch==2.14.1 --index-url https://download.pytorch.org/whl/cpu   # Linux only
pip install -r requirements.txt
python check_setup.py
python test_env.py
python train_ppo.py --timesteps 500000 --seed 0
python plot_training.py --run-dir runs/ppo_seed0
python evaluate.py --model runs/ppo_seed0/best_model/best_model.zip --episodes 200
python view_policy.py --model runs/ppo_seed0/best_model/best_model.zip   # macOS: mjpython
```

## What was executed

See executed/execution.json for the exact environment. In short: the model and environment
were run with MuJoCo 3.16.0 and Gymnasium 1.4.0, test_env.py passed, and evaluate.py was run
for the two baselines. plot_training.py was run only on synthetic log files to test its parsing.
train_ppo.py, view_policy.py and the PPO part of evaluate.py were not executed, because PyTorch
and Stable-Baselines3 could not be installed in that environment. No PPO result is claimed
anywhere in this project.