Gymnasium basics

Catch, the grid arm and ACT’s simulator are very different worlds. Yet a learner talks to each of them the same way.

What is that common language?

Overview

Gymnasium is a Python package that defines one shape for every task. LeRobot’s simulators use it, so learning it once is enough.

The mental model: a task is a box with two buttons. reset starts an episode; step takes an action and says what happened.

The loop

Every Gymnasium learner runs this loop:

gymnasium_basics.py · the loopOpen source ↗

observation, info = env.reset(seed=seed)
steps = []
while True:
    action = policy(observation)
    observation, reward, terminated, truncated, info = env.step(action)
    steps.append((action, observation, reward, terminated, truncated))
    if terminated or truncated:
        return steps

Full file · gymnasium_basics.py · 183 lines

"""Gymnasium basics: the Robot basics arm in MuJoCo, wrapped as a Gymnasium task.

Run: python3 site/ml/foundations/examples/gymnasium_basics.py

Gymnasium is the standard way to package a task for a learner: reset() starts an episode, step(action) moves it on.
LeRobot's simulators (gym-aloha, gym-hil) use exactly these calls.
"""

import time

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

L1, L2 = 0.10, 0.10                                  # link lengths in metres, as in Robot basics
COMMANDS = np.array([[10, 0], [-10, 0], [0, 10], [0, -10]])   # degrees added to the shoulder or elbow target
COMMAND_NAMES = ["shoulder +10°", "shoulder −10°", "elbow +10°", "elbow −10°"]
PERIOD = 0.2                                         # seconds of physics per command
START = np.array([90.0, -90.0])                      # shoulder up, forearm pointing forward

XML = f"""
<mujoco model="two-link arm">
  <option timestep="0.001" gravity="0 0 -9.81"/>
  <worldbody>
    <light pos="0 -0.3 0.4" dir="0 0.6 -1"/>
    <camera name="side" pos="0 -0.6 0.03" xyaxes="1 0 0 0 0 1"/>
    <site name="target" type="sphere" size="0.015" rgba="0.95 0.45 0.1 0.6"/>
    <geom type="plane" size="0.4 0.3 0.01" pos="0 0 -0.21" rgba="0.82 0.83 0.8 1" contype="0" conaffinity="0"/>
    <body name="upper_arm">
      <joint name="shoulder" type="hinge" axis="0 -1 0" armature="0.02" range="0 180"/>
      <geom type="capsule" fromto="0 0 0 {L1} 0 0" size="0.004" mass="0.20"/>
      <body name="forearm" pos="{L1} 0 0">
        <joint name="elbow" type="hinge" axis="0 -1 0" armature="0.02" range="-180 0"/>
        <geom type="capsule" fromto="0 0 0 {L2} 0 0" size="0.004" mass="0.15"/>
        <site name="hand" pos="{L2} 0 0"/>
      </body>
    </body>
  </worldbody>
  <actuator>
    <position joint="shoulder" kp="6" kv="0.4"/>
    <position joint="elbow" kp="6" kv="0.4"/>
  </actuator>
</mujoco>
"""


def hand_position(shoulder, elbow):
    """Where the hand is for two joint angles in degrees, in centimetres (x forward, y up)."""
    a, b = np.radians(shoulder), np.radians(shoulder + elbow)
    return 100 * np.array([L1 * np.cos(a) + L2 * np.cos(b), L1 * np.sin(a) + L2 * np.sin(b)])


class ArmReach(gym.Env):
    """Bring the hand within 1.5 cm of a target. One step is one command, then 0.2 s of physics."""
    metadata = {"render_modes": ["human", "rgb_array"], "render_fps": 5}

    def __init__(self, render_mode=None):
        self.render_mode = render_mode               # "human": MuJoCo's viewer window; "rgb_array": pictures
        self.viewer = self.renderer = None
        self.model = mujoco.MjModel.from_xml_string(XML)
        self.data = mujoco.MjData(self.model)
        # What the learner sees: 2 joint angles (°), 2 joint speeds (°/s), the target's x and y (cm).
        self.observation_space = spaces.Box(low=np.array([-5, -185, -360, -360, -20, -20], dtype=np.float32),
                                            high=np.array([185, 5, 360, 360, 20, 20], dtype=np.float32))
        # What the learner may do: one of the 4 commands.
        self.action_space = spaces.Discrete(4)

    def reset(self, seed=None, options=None):
        super().reset(seed=seed)                     # seeds self.np_random: the same seed, the same target
        shoulder = self.np_random.integers(1, 18) * 10            # a target the arm can reach: the hand of a random pose
        elbow = -self.np_random.integers(1, 18) * 10
        self.target = hand_position(shoulder, elbow)
        self.model.site("target").pos = [self.target[0] / 100, 0, self.target[1] / 100]   # the orange ball, for pictures
        mujoco.mj_resetData(self.model, self.data)
        self.command = START.copy()
        self.data.qpos[:] = self.data.ctrl[:] = np.radians(self.command)
        mujoco.mj_forward(self.model, self.data)
        if self.render_mode == "human":
            self.render()
        return self.observe(), {}

    def step(self, action):
        self.command = np.clip(self.command + COMMANDS[action], [0, -180], [180, 0])
        self.data.ctrl[:] = np.radians(self.command)                 # the new target angles for the motors
        mujoco.mj_step(self.model, self.data, nstep=round(PERIOD / self.model.opt.timestep))
        distance = np.linalg.norm(self.hand() - self.target)
        terminated = bool(distance < 1.5)            # the task is done
        reward = 1.0 if terminated else 0.0
        if self.render_mode == "human":
            self.render()
        return self.observe(), reward, terminated, False, {"distance": distance}

    def render(self):
        """The picture comes from MuJoCo: Gymnasium only defines when and how to ask for it."""
        if self.render_mode == "rgb_array":
            if self.renderer is None:
                self.renderer = mujoco.Renderer(self.model, 240, 320)
            self.renderer.update_scene(self.data, camera="side")
            return self.renderer.render()
        if self.render_mode == "human":
            if self.viewer is None:
                self.viewer = mujoco.viewer.launch_passive(self.model, self.data)
                self.viewer.cam.lookat[:], self.viewer.cam.distance = [0, 0, 0.02], 0.65     # look at the arm
                self.viewer.cam.azimuth, self.viewer.cam.elevation = 90, -8
            self.viewer.sync()
            time.sleep(1 / self.metadata["render_fps"])

    def close(self):
        if self.viewer is not None:
            self.viewer.close()
            time.sleep(0.5)                          # let the viewer's window thread finish before Python exits
        if self.renderer is not None:
            self.renderer.close()

    def observe(self):
        angles, speeds = np.degrees(self.data.qpos), np.degrees(self.data.qvel)
        return np.concatenate([angles, speeds, self.target]).astype(np.float32)

    def hand(self):
        x, _, z = self.data.site("hand").xpos
        return 100 * np.array([x, z])


gym.register("ArmReach-v0", entry_point=ArmReach, max_episode_steps=40)   # gym.make adds a 40-command time limit


def closest_policy(observation):
    """A hand-written policy: the command whose pose puts the hand nearest the target."""
    angles, target = observation[:2], observation[4:]
    return int(np.argmin([np.linalg.norm(hand_position(*(angles + c)) - target) for c in COMMANDS]))


def random_policy(env):
    return lambda observation: int(env.action_space.sample())


def run_episode(env, policy, seed):
    """The loop every Gymnasium learner uses."""
    observation, info = env.reset(seed=seed)
    steps = []
    while True:
        action = policy(observation)
        observation, reward, terminated, truncated, info = env.step(action)
        steps.append((action, observation, reward, terminated, truncated))
        if terminated or truncated:
            return steps


def success_rate(policy_for, seeds):
    env = gym.make("ArmReach-v0")
    env.action_space.seed(0)                         # the random policy's draws repeat too
    results = [run_episode(env, policy_for(env), seed) for seed in seeds]
    return sum(r[-1][3] for r in results), [len(r) for r in results]


if __name__ == "__main__":
    env = gym.make("ArmReach-v0")
    print("observation space:", env.observation_space, "| action space:", env.action_space)

    observation, info = env.reset(seed=3)
    print("reset(seed=3):", observation.round(2))
    print("reset(seed=3) again:", env.reset(seed=3)[0].round(2), "| reset(seed=4):", env.reset(seed=4)[0].round(2))

    env.reset(seed=3)
    observation, reward, terminated, truncated, info = env.step(0)
    print("step(0):", observation.round(2), reward, terminated, truncated, f"distance {info['distance']:.2f} cm")

    from gymnasium.utils.env_checker import check_env
    check_env(gym.make("ArmReach-v0").unwrapped)
    print("check_env: passed")

    for name, policy_for in (("random", random_policy), ("closest", lambda env: closest_policy)):
        reached, lengths = success_rate(policy_for, range(100))
        print(f"{name:8s} policy: {reached}/100 targets reached; commands per episode: median {np.median(lengths):.0f}")

    envs = gym.make_vec("ArmReach-v0", num_envs=10, vectorization_mode="sync")
    observations, infos = envs.reset(seed=1000)                  # copy i gets seed 1000 + i
    print("10 copies: observations", observations.shape, "| first target of copy 1:", observations[1, 4:].round(2),
          "| alone with seed 1001:", env.reset(seed=1001)[0][4:].round(2))
    observations, rewards, terminated, truncated, infos = envs.step(np.zeros(10, dtype=int))
    print("one step of all 10:", observations.shape, rewards.shape, terminated.shape)

observation is what the learner sees now.
reward is the score for the last action.
terminated says the task itself ended: success, or a crash.
truncated says time ran out.
info holds extras for us, never for the learner.

Our task

The arm from Robot basics, now asked to bring its hand within 1.5 cm of a target. Wrapping it takes one class with two methods:

gymnasium_basics.py · the taskOpen source ↗

def step(self, action):
    self.command = np.clip(self.command + COMMANDS[action], [0, -180], [180, 0])
    self.data.ctrl[:] = np.radians(self.command)                 # the new target angles for the motors
    mujoco.mj_step(self.model, self.data, nstep=round(PERIOD / self.model.opt.timestep))
    distance = np.linalg.norm(self.hand() - self.target)
    terminated = bool(distance < 1.5)            # the task is done
    reward = 1.0 if terminated else 0.0
    return self.observe(), reward, terminated, False, {"distance": distance}

Full file: gymnasium_basics.py, in its first card above ↑

One step is one command, then 0.2 s of MuJoCo physics. The motors chase the new target angles, as in Robot basics.

The class also declares its spaces, the shape of what goes in and out:

Action space: Discrete(4), one of four commands: shoulder or elbow, +10° or −10°.
Observation space: Box of 6 numbers: two joint angles (°), two joint speeds (°/s), the target’s x and y (cm).

Why declare them? A learner can then size its network without knowing anything else about the task: 6 numbers in, 4 choices out.

Seeds

reset(seed=3) picks a random target, but the same one every time for seed 3:

That is how an evaluation repeats exactly: ACT’s 200 test episodes were seeds 1000 to 1199, the same cube positions for every policy.

Episode end

An episode ends in one of two ways, and Gymnasium keeps them apart:

Terminated: the hand reached the target. Nothing comes after this, so a learner counts no future reward.
Truncated: 40 commands passed. The task was not over; we simply stopped watching.

The 40-command limit is added when the task is registered, outside the class:

gymnasium_basics.py · registering the taskOpen source ↗

gym.register("ArmReach-v0", entry_point=ArmReach, max_episode_steps=40)   # gym.make adds a 40-command time limit

Full file: gymnasium_basics.py, in its first card above ↑

Then gym.make("ArmReach-v0") builds it, exactly as LeRobot builds gym_aloha/AlohaTransferCube-v0.

Two policies

A policy turns an observation into an action. Here are two, run through the same loop:

Random: env.action_space.sample(), any of the four commands.
Closest: a hand-written rule. Try each command in your head, and pick the one that puts the hand nearest the target.

Pick a policy and a seed, then play. Predict first: how often will random commands find the target in 40 tries?

ClosestRandom

Step0

How to read the lab:

The orange ring is the target; the grey line is the hand’s path so far.
The table is exactly what step returned at this point.

Render

Gymnasium itself has no window and draws nothing. It only fixes one more method, render(), and the task fills it with its simulator’s picture:

gymnasium_basics.py · renderOpen source ↗

if self.render_mode == "rgb_array":
    if self.renderer is None:
        self.renderer = mujoco.Renderer(self.model, 240, 320)
    self.renderer.update_scene(self.data, camera="side")
    return self.renderer.render()
if self.render_mode == "human":
    if self.viewer is None:
        self.viewer = mujoco.viewer.launch_passive(self.model, self.data)

Full file: gymnasium_basics.py, in its first card above ↑

render_mode="rgb_array": each render() returns a picture, a 240 × 320 × 3 array, made by MuJoCo’s renderer.
render_mode="human": the task opens MuJoCo’s viewer window and shows every step live.

Seed 3 with the closest rule: one render() picture per step.

MuJoCo’s viewer window in "human" mode, seeds 0 to 4.

To open that window on your computer, run gymnasium_watch.py:

python3 site/ml/foundations/examples/gymnasium_watch.py

gym-aloha works the same way: its camera picture, and the test videos in the ACT workshop, come from MuJoCo through render().

Ten at once

Gymnasium can run several copies of a task side by side, stepping them together:

gymnasium_basics.py · ten copiesOpen source ↗

envs = gym.make_vec("ArmReach-v0", num_envs=10, vectorization_mode="sync")
observations, infos = envs.reset(seed=1000)                  # copy i gets seed 1000 + i

Full file: gymnasium_basics.py, in its first card above ↑

A policy network can then compute all ten actions in one batch. That is what lerobot-eval --eval.batch_size=10 did in the ACT workshop.

gym-aloha

LeRobot’s simulator for ACT is the same box with the same two buttons. Only the contents differ:

The loop above would run it unchanged. Only the policy needs to change, and that is what ACT is.

Recap

gym.register(..., max_episode_steps=40) names the task and adds its time limit; gym.make builds it.
env.reset(seed=3) starts an episode; the seed fixes the random parts.
env.step(action) returns observation, reward, terminated, truncated and info.
observation_space and action_space tell a learner the sizes: 6 numbers in, 4 choices out.
gym.make("ArmReach-v0", render_mode=...) chooses pictures ("rgb_array") or MuJoCo’s window ("human"); MuJoCo draws both.
gym.make_vec(..., num_envs=10) steps ten copies together, seeds 1000 + i.

Next, the same arm gets a wall to reach past: Robot arm in MuJoCo.

Key concepts

Environment: A task packaged for a learner, with reset and step.
Observation · action: What the learner sees · what it does.
Space: The type and size of observations or actions, such as Discrete(4) or Box.
Episode: One run from reset until it ends.
Terminated · truncated: Ended by the task itself · stopped by the time limit.
Seed: A number that fixes the random choices, so an episode can be repeated exactly.
Render: Asking the task for a picture or a window; the simulator draws it.
Vector environment: Several copies stepped together.

Run it yourself

From the repository root:

python3 site/ml/foundations/examples/gymnasium_basics.py
python3 site/ml/foundations/examples/gymnasium_watch.py

The first prints the spaces, the seeded resets, one step, Gymnasium’s own checker result, both policies over 100 seeds and the ten copies (about 3 seconds). The second opens MuJoCo’s viewer and plays five targets. tools/build_gymnasium_data.py saves the episodes this page plays.

Reference: Gymnasium basic usage. The lesson works offline; reference links need a connection.

Gymnasium basics

Overview

The loop

Our task

Seeds

Episode end

Two policies

The loop, step by step

Render

Ten at once

gym-aloha

Recap

Key concepts