← All lessonsWorkshop · lesson 1

Bouncing ball

The MuJoCo lesson showed a simulator that knows physics. The PyTorch lessons showed models that learn from examples.

Here we put them together. MuJoCo throws a ball hundreds of times, and a small network learns to predict where the ball will be 10 milliseconds later, bounce included, without ever being told about gravity or the floor.

Overview

Write down what a ball does in the next 10 ms and you need gravity, and a special rule for the floor. Instead of writing those rules, we let a network find them from examples.

The simulator plays the role of the world: it produces the right answers, and the network only ever sees those answers.

Throw the ball600 throws in MuJoCo.Cut into pairsState now → change 10 ms later.Train a networkThe PyTorch loop, at scale.Test itOn 100 throws it never saw.

Three things are new compared with the foundation lessons: where the examples come from, training on 100,000 examples, and a test where the network must live with its own guesses. Everything else is what you already know.

World

The world is the MuJoCo lesson’s ball with one addition: a second slide joint, so a throw can travel sideways as well as up and down. Friction is set to zero, so the sideways speed never changes and the interesting physics stays vertical.

ball.py · recording one throwOpen source ↗
for t in range(n_steps):
    for _ in range(SUBSTEPS):
        mujoco.mj_step(model, data)
        mujoco.mj_forward(model, data)      # refresh contacts for the new position
        contact[t] |= data.ncon > 0
    states[t + 1] = [*data.qpos, *data.qvel]
Full file · ball.py · 60 lines
"""
ball.py — the bouncing-ball world. Everything here is fixed physics; nothing is learned.

It is the ball from the MuJoCo basics lesson (radius 5 cm, 100 g, the same contact spring, 1 ms steps)
with one addition: a sideways joint, so a throw can travel across the floor.

State of the ball, used by every other script:
    s = [x, z, vx, vz]   (m, m, m/s, m/s)   vz > 0 means moving up
The velocities are MuJoCo's own state (data.qvel), not estimated from positions.
"""

import numpy as np
import mujoco

SIM_DT = 0.001                  # one MuJoCo step: 1 ms (the bounce needs small steps)
SUBSTEPS = 10                   # we keep every 10th state ...
MODEL_DT = SIM_DT * SUBSTEPS    # ... so the network sees the ball every 10 ms
RADIUS = 0.05
CONTACT = 'solref="-100000 -20" solimp="0.99 0.99 0.001"'   # a stiff, lightly damped spring

SCENE_XML = f"""
<mujoco model="bouncing ball">
  <option timestep="{SIM_DT}" gravity="0 0 -9.81"/>
  <default>
    <geom friction="0 0 0" {CONTACT}/>
  </default>
  <worldbody>
    <geom name="floor" type="plane" size="50 5 0.1"/>
    <body name="ball">
      <joint name="x" type="slide" axis="1 0 0"/>
      <joint name="z" type="slide" axis="0 0 1"/>
      <geom type="sphere" size="{RADIUS}" mass="0.1"/>
    </body>
  </worldbody>
</mujoco>
"""


def simulate(s0, n_steps):
    """Throw the ball from state s0 and record n_steps states, 10 ms apart.

    Returns states (n_steps + 1, 4) and contact (n_steps,): True if the ball touched the floor
    at any moment between two recorded states.
    """
    model = mujoco.MjModel.from_xml_string(SCENE_XML)
    data = mujoco.MjData(model)
    data.qpos[:] = s0[:2]           # x, z
    data.qvel[:] = s0[2:]           # vx, vz
    mujoco.mj_forward(model, data)

    states = np.zeros((n_steps + 1, 4))
    contact = np.zeros(n_steps, dtype=bool)
    states[0] = s0
    for t in range(n_steps):
        for _ in range(SUBSTEPS):
            mujoco.mj_step(model, data)
            mujoco.mj_forward(model, data)      # refresh contacts for the new position
            contact[t] |= data.ncon > 0
        states[t + 1] = [*data.qpos, *data.qvel]
    return states, contact

MuJoCo still steps every 1 ms, because the bounce needs small steps. But we keep only every 10th state (SUBSTEPS = 10), so the network sees the ball once every 10 ms.

What gets saved:

  • each state is four numbers, [x, z, vx, vz]: sideways position, height, sideways speed, vertical speed;
  • contact[t] remembers whether the ball touched the floor at any moment between two saved states, even if the bounce began and ended in between.
One recorded throw
State now [x, z, vx, vz]—
Change in the next 10 ms—

Scrub along the throw. In the air, the dots are evenly spread sideways and bunch up near the top of each arc.

At the floor, two orange intervals turn the fall around.

Examples

One example for the network is one pair: a saved state, and how it changed by the next saved state. Every throw is 200 intervals of 10 ms, so every throw gives 200 pairs:

03_train.py · from throws to pairsOpen source ↗
throws = data["states"][data["split"] == "train"]             # (500, 201, 4)
now, after = throws[:, :-1].reshape(-1, 4), throws[:, 1:].reshape(-1, 4)
inputs = torch.tensor(now, dtype=torch.float32)               # (100000, 4)
targets = torch.tensor(after - now, dtype=torch.float32)      # the change in 10 ms
Full file · 03_train.py · 60 lines
"""
Step 3 — Train the network to predict the change in 10 ms.

    input  : s_t = [x, z, vx, vz]            output : s_(t+1) - s_t
    model  : 4 -> 16 ReLU neurons -> 4        loss   : mean squared error, in raw units (m, m/s)
    update : plain SGD, learning rate 0.1, batches of 256, 240 passes over the data

Run:  python 03_train.py                      (the lesson's network)
      python 03_train.py --hidden 2           (change one thing; see 05_compare.py)
      python 03_train.py --hidden 8 8         (two hidden layers of 8)
Out:  results/<tag>/network.pt, results/<tag>/losses.json
"""

import argparse, json, os
import numpy as np
import torch
from torch import nn
from model import make_network

parser = argparse.ArgumentParser()
parser.add_argument("--hidden", type=int, nargs="+", default=[16])
parser.add_argument("--lr", type=float, default=0.1)
parser.add_argument("--seed", type=int, default=0)
parser.add_argument("--epochs", type=int, default=240)
parser.add_argument("--batch", type=int, default=256)
args = parser.parse_args()
EPOCHS, BATCH = args.epochs, args.batch
tag = f"hidden{'x'.join(map(str, args.hidden))}_lr{args.lr:g}_seed{args.seed}"
tag += (f"_e{EPOCHS}" if EPOCHS != 240 else "") + (f"_b{BATCH}" if BATCH != 256 else "")

# Data: every pair from the training throws.
data = np.load("data/throws.npz")
throws = data["states"][data["split"] == "train"]             # (500, 201, 4)
now, after = throws[:, :-1].reshape(-1, 4), throws[:, 1:].reshape(-1, 4)
inputs = torch.tensor(now, dtype=torch.float32)               # (100000, 4)
targets = torch.tensor(after - now, dtype=torch.float32)      # the change in 10 ms

torch.manual_seed(args.seed)
network = make_network(args.hidden[0] if len(args.hidden) == 1 else args.hidden)
loss_fn = nn.MSELoss()
optimizer = torch.optim.SGD(network.parameters(), lr=args.lr)

losses = []
for epoch in range(EPOCHS):
    order = torch.randperm(len(inputs))                       # a new shuffle every pass
    for i in range(0, len(order), BATCH):
        batch = order[i:i + BATCH]
        loss = loss_fn(network(inputs[batch]), targets[batch])
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    with torch.no_grad():
        losses.append(loss_fn(network(inputs), targets).item())
    if (epoch + 1) % 40 == 0:
        print(f"epoch {epoch + 1:3d}  loss {losses[-1]:.5f}")

os.makedirs(f"results/{tag}", exist_ok=True)
torch.save({"hidden": args.hidden[0] if len(args.hidden) == 1 else args.hidden, "weights": network.state_dict()}, f"results/{tag}/network.pt")
json.dump({"args": vars(args), "losses": losses}, open(f"results/{tag}/losses.json", "w"))
print(f"saved results/{tag}/")
  • throws[:, :-1] is every state except each throw’s last; throws[:, 1:] is every state except its first. Row by row, they line up as “now” and “10 ms later”.
  • reshape(-1, 4) stacks all 500 throws into one long table of 100,000 rows, four numbers each: the batch shape from PyTorch basics, at scale.

The target is the change, after - now, not the next state itself. In 10 ms the ball barely moves, so the next state is almost a copy of the current one. Asking for the change puts all the network’s effort on the part that carries the physics:

  • In the air, the vertical speed always changes by the same −0.0981 m/s: gravity.
  • At the floor, it jumps by more than 6 m/s in a single interval.

Only 2.6% of all pairs touch the floor, so the hardest thing to learn is also the rarest.

The 600 throws start from random heights, positions and speeds. Whole throws are split before training: 500 to learn from, 100 kept aside for the test.

Why whole throws? If pairs from one throw landed on both sides, a test pair would sit 10 ms from a training pair, and the test would reward remembering instead of predicting.

One square per throw.

Blue throws train the network; orange throws are never shown to it until the test.

Network

The model needs a bend: the vertical speed must change by \(-0.0981\) almost everywhere, then suddenly by \(+6\) when the ball is low and falling.

So we use a network like the one in Neural network: one input, with four numbers in, four out, and 16 hidden neurons:

model.py · the networkOpen source ↗
def make_network(hidden=16):
    """4 numbers in (the state) -> `hidden` ReLU neurons -> 4 numbers out (the change)."""
    return nn.Sequential(nn.Linear(4, hidden), nn.ReLU(), nn.Linear(hidden, 4))
Full file · model.py · 34 lines
"""
model.py — the learned part: a network that maps a state to its change over the next 10 ms.
One hidden layer of ReLU neurons gives the rule the bends it needs to turn a fall into a bounce.
"""

import torch
from torch import nn


def make_network(hidden=16):
    """4 numbers in (the state) -> hidden layer(s) of ReLU neurons -> 4 numbers out (the change).

    hidden=16 gives one hidden layer of 16 neurons; hidden=[8, 8] gives two hidden layers of 8.
    """
    sizes = [hidden] if isinstance(hidden, int) else list(hidden)
    layers, width = [], 4
    for size in sizes:
        layers += [nn.Linear(width, size), nn.ReLU()]
        width = size
    return nn.Sequential(*layers, nn.Linear(width, 4))


def load_network(path):
    saved = torch.load(path)
    network = make_network(saved["hidden"])
    network.load_state_dict(saved["weights"])
    return network


def predict_next(network, states):
    """Next state = current state + the network's guessed change. Works on one state or many."""
    with torch.no_grad():
        change = network(torch.tensor(states, dtype=torch.float32)).numpy()
    return states + change
The same network the classic way: 4 inputs, a hidden layer of 16 ReLU neurons, 4 linear output neurons.

Every line is one weight.

It has the same shape as that lesson’s curve-fitter, just wider at the ends:

  • each of the 16 hidden neurons reads all four state numbers and bends once;
  • each of the 4 output neurons adds the 16 ramps into one predicted change.

That is \(4 \times 16 + 16 = 80\) numbers in the first layer and \(16 \times 4 + 4 = 68\) in the second: 148 adjustable numbers in all.

Why 16 neurons, and not 2 or 4 or 64? Not by habit: we compared sizes. Small networks either get stuck or cannot do both jobs, flying and bouncing.

The Network design lesson shows that comparison, with learning curves and test results for every size.

Training

The loop is the one from PyTorch basics. Two things are new, both because we now have 100,000 examples:

03_train.py · the training loopOpen source ↗
network = make_network(16)
loss_fn = nn.MSELoss()
optimizer = torch.optim.SGD(network.parameters(), lr=0.1)

for epoch in range(240):
    order = torch.randperm(len(inputs))                       # a new shuffle every pass
    for i in range(0, len(order), 256):
        batch = order[i:i + 256]
        loss = loss_fn(network(inputs[batch]), targets[batch])
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

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

Batches. Each update looks at 256 pairs, not all 100,000. A small random sample gives a good enough direction, and it is far cheaper than scoring everything.

torch.randperm shuffles the pairs, and the inner loop walks through them 256 at a time: 391 updates.

Epoch
one walk through all the training pairs.

We do 240 epochs, about 94,000 updates in all, which takes under 20 seconds on a laptop CPU.

Everything inside the inner loop is the update you already know.

Why learning rate 0.1, batches of 256 and 240 epochs? Each was chosen by comparison:

  • learning rate: 0.01 learns too slowly, and 0.3 overshoots;
  • batch size: batches of 1 are too noisy, and the whole dataset at once gives too few updates;
  • epochs: more keep helping, with diminishing returns.

Network design runs each of those experiments, one change at a time.

Test

Now the 100 throws the network never saw. We ask two different questions.

First, one step: start from a true state and guess 10 ms ahead.

As a yardstick we use “gravity only”: the exact free-fall formula a physics student would write, with no floor.

  • In the air, the formula wins by a mile: it is the true law, and the network only approximates it.
  • At the floor, the formula has no idea a floor exists and misses by the whole bounce, while the network gets most of it.

The network learned something the formula lacks.

The second question is harder: a rollout. Give the network only the first state of a throw, then feed each guess back in as the next input, for 200 steps.

Nothing corrects it.

04_evaluate.py · rolloutOpen source ↗
path = np.zeros_like(truth)
path[:, 0] = truth[:, 0]                      # only the starting state is real
for t in range(truth.shape[1] - 1):
    path[:, t + 1] = step(path[:, t])         # each guess becomes the next input
Full file · 04_evaluate.py · 73 lines
"""
Step 4 — Test on the 100 throws the network never saw.

    A. One step: from the TRUE state, how far off is the guess 10 ms later?
    B. Rollout:  from the true start only, feed each guess back in for 200 steps (2 s). Nothing corrects it.

The yardstick is "gravity only": exact free-fall physics with no floor.

Run:  python 04_evaluate.py [--tag hidden16_lr0.1_seed0]
Out:  results/<tag>/metrics.json, rollouts.npz, 04_rollouts.png
"""

import argparse, json
import numpy as np
import matplotlib.pyplot as plt
import ball
from model import load_network, predict_next

parser = argparse.ArgumentParser()
parser.add_argument("--tag", default="hidden16_lr0.1_seed0")
tag = parser.parse_args().tag
network = load_network(f"results/{tag}/network.pt")
data = np.load("data/throws.npz")
truth = data["states"][data["split"] == "test"]               # (100, 201, 4)
contact = data["contact"][data["split"] == "test"]            # (100, 200)


def gravity_only(s):
    """Newton with no floor: z += vz dt - g dt^2 / 2, vz -= g dt. x moves at constant vx."""
    dt, g = ball.MODEL_DT, 9.81
    return s + np.stack([s[..., 2] * dt, s[..., 3] * dt - g * dt**2 / 2, 0 * s[..., 2], -g * dt + 0 * s[..., 3]], -1)


def network_step(s):
    return predict_next(network, s)


# A. One step from the true state, scored separately in free flight and at the floor.
now, after = truth[:, :-1].reshape(-1, 4), truth[:, 1:].reshape(-1, 4)
floor = contact.reshape(-1)
one_step = {}
for name, step in [("network", network_step), ("gravity", gravity_only)]:
    miss = step(now) - after
    position_mm = np.linalg.norm(miss[:, :2], axis=1) * 1000
    velocity = np.linalg.norm(miss[:, 2:], axis=1)
    one_step[name] = {"flight_position_mm": float(np.sqrt(np.mean(position_mm[~floor] ** 2))),
                      "floor_velocity": float(np.sqrt(np.mean(velocity[floor] ** 2)))}
    print(f"{name:8s} free flight {one_step[name]['flight_position_mm']:7.3f} mm   "
          f"at the floor {one_step[name]['floor_velocity']:6.3f} m/s")

# B. Rollouts: each prediction becomes the next input.
rollouts = {}
for name, step in [("network", network_step), ("gravity", gravity_only)]:
    path = np.zeros_like(truth)
    path[:, 0] = truth[:, 0]
    for t in range(truth.shape[1] - 1):
        path[:, t + 1] = step(path[:, t])
    rollouts[name] = path
error_cm = {k: np.linalg.norm(v[:, :, :2] - truth[:, :, :2], axis=2) * 100 for k, v in rollouts.items()}
for k, e in error_cm.items():
    print(f"{k:8s} median rollout error: {np.median(e[:, 20]):6.1f} cm at 0.2 s, {np.median(e[:, 100]):6.1f} cm at 1 s")

json.dump({"one_step": one_step, "median_error_cm": {k: np.median(e, axis=0).round(3).tolist() for k, e in error_cm.items()}},
          open(f"results/{tag}/metrics.json", "w"))
np.savez_compressed(f"results/{tag}/rollouts.npz", truth=truth[:8], **{k: v[:8] for k, v in rollouts.items()})

fig, axes = plt.subplots(2, 2, figsize=(10, 6))
for i, ax in enumerate(axes.flat):
    ax.plot(truth[i, :, 0], truth[i, :, 1], lw=2.5, label="MuJoCo")
    ax.plot(rollouts["network"][i, :, 0], rollouts["network"][i, :, 1], "--", label="network rollout")
    ax.axhline(0, color="grey", lw=1); ax.set_ylim(-0.3, 1.8); ax.set_title(f"test throw {i}")
axes[0, 0].legend()
fig.tight_layout(); fig.savefig(f"results/{tag}/04_rollouts.png", dpi=120)
Rollouts on test throws
Test throw

  • Gravity only is nearly perfect until the first bounce; then its ball falls through the floor forever.
  • The network does bounce, but its bounces come out too weak. Each one loses more height than the real one, so on all four throws above, its ball ends up skimming the floor while the real ball is still bouncing.

Every small error becomes the next starting point, and the gap grows. Across all 100 test throws, the median position error is about 2.5 cm after 0.2 s and 27 cm after one second.

Being good at one step does not guarantee being good for two seconds. That is the real lesson of any learned simulator.

Recap

The whole workshop, in the order the code runs it:

  • ball.simulate(s0, 200) throws the ball in MuJoCo from a starting state and keeps every 10th 1 ms state, plus whether the floor was touched in between.
  • 02_make_dataset.py makes 600 random throws and splits whole throws: 500 to train on, 100 for the test.
  • inputs and targets are 100,000 pairs: a state, and how it changed 10 ms later.
  • network = make_network(16) builds 4 inputs → a hidden layer of 16 ReLU neurons → 4 linear output neurons: 148 adjustable numbers with the bends a bounce needs.
  • for epoch in range(240) walks through the shuffled pairs 240 times, 256 at a time. Each batch runs the update from PyTorch basics: loss_fn, zero_grad(), backward(), step().
  • 04_evaluate.py tests on the 100 unseen throws: one step from the truth, then a rollout that feeds on its own guesses.

The simulator supplied every answer. The network supplied none of the physics, and still learned gravity and a floor from examples.

What it did not learn is how to stay accurate on its own for two seconds.

The next lesson, Network design, tests the choices this one made: size, learning rate, batch size, training length, and a second hidden layer.

Key concepts

State
The four numbers that describe the ball at one moment: x, z, vx, vz.
Pair
One example: a state, and its change 10 ms later. The change is what the network predicts.
Train / test split
Whole throws are divided: the network learns from 500 and is judged on 100 it never saw.
Hidden layer
Neurons between input and output. Each ReLU neuron adds one bend the network can use.
Batch
The 256 pairs one update looks at.
Epoch
One pass through all training pairs, in shuffled batches.
Baseline
A simple fixed rule used as a yardstick. Here: gravity only, no floor.
Rollout
Predicting a whole throw by feeding each guess back in as the next input, with no correction.

Every number and curve on this page comes from the lab’s saved runs (results/); the page only plays them back.