Training CartPole on WebGPU

This page trains a neural network to balance a pole on a cart. Training starts when the page loads and runs in WebGPU compute shaders on your GPU. If your browser doesn't support WebGPU, the same algorithm runs on the CPU.

Live TrainingLoading…

t = 0

Updates: 0
Environment Steps: 0
Episodes: 0
Mean Length: –

Mean Episode Length, Last 20 Episodes–

Waiting for the first finished episode…

Policy Loss–

Waiting for the first update…

Value Loss–

Waiting for the first update…

Starting the trainer…

At first, the policy picks left and right with equal probability, and the pole falls after about 20 steps. In my tests, the mean episode length passed 475 steps after 230 to 800 updates. That is the threshold Gymnasium uses to call CartPole solved. The rest of this post explains the math behind each number in the figures.

CartPole as a Markov decision process

CartPole is a control task from Barto, Sutton, and Anderson (1983). A cart moves on a frictionless track, and a hinge attaches a pole to the cart. At each time step, you push the cart left or right with a fixed force. The goal is to keep the pole upright for as long as possible.

A Markov decision process (MDP) describes this kind of problem. Markov means that the next state depends only on the current state and action, not on earlier steps. The CartPole MDP has these parts:

State: $s_t = (x_t, \dot x_t, \theta_t, \dot\theta_t)$ . These are the cart position in meters, the cart velocity, the pole angle from vertical in radians, and the angular velocity of the pole.
Action: $a_t \in \{\text{left}, \text{right}\}$ . The action applies a force of $F_t = \pm 10\,\text{N}$ to the cart.
Transition: $s_{t+1} = f(s_t, a_t)$ . The next section gives $f$ , which is deterministic.
Reward: $r_t = 1$ for each step.
Termination: The episode ends when $|\theta_{t+1}| > 12^\circ$ or $|x_{t+1}| > 2.4\,\text{m}$ . The flag $d_{t+1}$ is 1 at that step and 0 otherwise. Episodes also stop after 500 steps.

Each step earns the same reward, so the return of an episode equals its length. More reward means more steps with the pole up. The trainer discounts later rewards with $\gamma = 0.99$ , so the return from step $t$ is $G_t = \sum_{k \ge 0} \gamma^k r_{t+k}$ .

One Transition

s_{t} = 0.000 0.000 0.000 0.000 a_{t} s_{t + 1} = 0.000 0.000 0.000 0.000 π_{θ} (right ∣ s_{t}) a_{t} r_{t} = 0.000 = right, F_{t} = 00 N = 1, d_{t + 1} = 0

The cart in the first figure takes one of these steps every 20 ms. Pause, then step to read each value.

How the cart and pole move

The transition function $f$ comes from the equations of motion for a pole on a cart. The simulator uses the same form as Gymnasium's CartPole. It computes the angular acceleration of the pole first, then uses it to compute the acceleration of the cart:

\ddot\theta_t = \frac{g\sin\theta_t - \cos\theta_t \cdot \dfrac{F_t + m_p l\,\dot\theta_t^2\sin\theta_t}{m_c + m_p}}{l\left(\dfrac{4}{3} - \dfrac{m_p\cos^2\theta_t}{m_c + m_p}\right)} \qquad \ddot x_t = \frac{F_t + m_p l\left(\dot\theta_t^2\sin\theta_t - \ddot\theta_t\cos\theta_t\right)}{m_c + m_p}

The constants are the Gymnasium defaults:

$g = 9.8\,\text{m/s}^2$ is gravity.
$m_c = 1\,\text{kg}$ is the mass of the cart.
$m_p = 0.1\,\text{kg}$ is the mass of the pole.
$l = 0.5\,\text{m}$ is half the length of the pole.

The simulator then moves each state variable forward by $\tau = 0.02\,\text{s}$ with an explicit Euler step. Each position changes by its old velocity, and each velocity changes by the new acceleration:

x_{t+1} = x_t + \tau\dot x_t \qquad \dot x_{t+1} = \dot x_t + \tau\ddot x_t \qquad \theta_{t+1} = \theta_t + \tau\dot\theta_t \qquad \dot\theta_{t+1} = \dot\theta_t + \tau\ddot\theta_t

The agent never sees these equations. It observes the four numbers in $s_t$ and the reward, and it learns the effect of each push from experience.

Accelerations for This Step

θ ¨_{t} x ¨_{t} θ_{t} = \frac{g sin θ _{t} - cos θ _{t} \cdot \frac{F _{t} + m _{p} l θ ˙ _{t}^{2} sin θ _{t}}{m _{c} + m _{p}}}{l ( \frac{4}{3} - \frac{m _{p} cos ^{2} θ _{t}}{m _{c} + m _{p}} )} = 00.00 rad/s^{2} = \frac{F _{t} + m _{p} l ( θ ˙ _{t}^{2} sin θ _{t} - θ ¨ _{t} cos θ _{t} )}{m _{c} + m _{p}} = 00.00 m/s^{2} = 0.000, θ ˙_{t} = 0.000, F_{t} = 00 N

The simulator evaluates both equations once per step, then applies Euler updates.

The policy and value networks

The trainer uses two small neural networks. The policy network $\pi_\theta$ maps a state to a probability for each action. The value network $V_\phi$ maps a state to an estimate of the discounted return from that state.

Each network has one hidden layer of 32 units with a tanh activation:

\pi_\theta(\cdot \mid s) = \operatorname{softmax}\left(W_2 \tanh(W_1 s + b_1) + b_2\right) \qquad V_\phi(s) = w_2^\top \tanh(U_1 s + c_1) + c_2

The policy network has 226 parameters and the value network has 193, so each update adjusts 419 weights. The output weights of the policy start near zero, so both actions start with a probability near 0.5.

The agent samples each action from $\pi_\theta$ instead of always taking the more likely action. Sampling makes the agent try both actions, and it needs both to learn which one works better.

Network Outputs

z_{t} π_{θ} (\cdot ∣ s_{t}) V_{ϕ} (s_{t}) = W_{2} tanh (W_{1} s_{t} + b_{1}) + b_{2} = (00.00, 00.00) = softmax (z_{t}) = (0.000, 0.000) = 00.00

Both networks read the same four numbers. Only the policy chooses the action.

Learning from advantages

The policy improves when actions that did better than expected become more likely. The advantage $\hat A_t$ measures how much better an action did. It is the return after $a_t$ minus the prediction $V_\phi(s_t)$ .

The trainer estimates the advantage with generalized advantage estimation (GAE). GAE starts from the temporal difference error $\delta_t$ , which compares one real reward plus the predicted value of the next state with the predicted value of the current state:

\delta_t = r_t + \gamma (1 - d_{t+1}) V_\phi(s_{t+1}) - V_\phi(s_t) \qquad \hat A_t = \delta_t + \gamma\lambda (1 - d_{t+1}) \hat A_{t+1}

With $\lambda = 0.95$ , $\hat A_t$ is a weighted sum of the errors over the next steps. The factor $1 - d_{t+1}$ stops the sum at the end of an episode. The target for the value network is $\hat R_t = \hat A_t + V_\phi(s_t)$ .

Each update minimizes two losses over a batch of $B = 512$ samples:

L_\pi(\theta) = -\frac{1}{B}\sum_t \hat A_t \log \pi_\theta(a_t \mid s_t) - \beta \mathcal{H} \qquad L_V(\phi) = \frac{1}{2B}\sum_t \left(V_\phi(s_t) - \hat R_t\right)^2

The gradient of $L_\pi$ increases $\log \pi_\theta(a_t \mid s_t)$ when $\hat A_t$ is positive and decreases it when $\hat A_t$ is negative. The entropy term $\mathcal{H}$ , with $\beta = 0.01$ , keeps the policy from becoming certain too early. Adam updates the policy with a learning rate of 0.003 and the value network with 0.01.

The policy loss doesn't measure progress the way a supervised loss does. Its scale depends on the advantages, which change as the value network improves. Use the episode length chart to measure progress. The value loss shows how well $V_\phi$ predicts returns. In my runs, it rose as episodes got longer, because the discounted returns grew from under 20 to almost 100.

Latest Update

A^_{t} L_{π} L_{V} H = 00.00, lo g π_{θ} (a_{t} ∣ s_{t}) = 0.000 = - \frac{1}{B} t \sum A^_{t} lo g π_{θ} (a_{t} ∣ s_{t}) - β H = 0.000 = \frac{1}{2 B} t \sum (V_{ϕ} (s_{t}) - R^_{t})^{2} = 00.00 = 0.000

The advantage and log probability belong to the first sample of the batch. The losses average all 512 samples.

Running the trainer on the GPU

An update has two kinds of independent work. Each environment steps on its own, and the gradient of each parameter is a separate sum over the batch. The trainer runs one WGSL compute shader for each kind:

The rollout kernel runs one thread per environment. Each of the 32 threads runs 16 steps. At each step, it evaluates both networks, samples an action, steps the physics, and resets the environment if the episode ended. Then it walks back through its 16 steps to compute $\hat A_t$ and the gradients for each sample.
The update kernel runs one thread per parameter. Each of the 419 threads sums its gradient over the 512 samples and applies the Adam step to its parameter.

The environment states, the random number generator states, and the Adam moments stay in GPU buffers between updates. After the updates for each animation frame, the page reads back the weights, the lengths of finished episodes, and the losses. The cart and the equations on this page use those weights on the CPU.

The GPU and CPU versions use the same PCG random number generator, so I can check one against the other. From the same weights and seeds, they produce the same episodes. After 3 updates, the weights differ by at most $7.6 \times 10^{-7}$ because of float32 rounding on the GPU. On an Apple M5 Max, a GPU update took 0.3 ms in Chrome, including one readback for every 4 updates. A CPU update took 2.1 ms in Node.

The page limits training to 30 updates per second so you can watch the policy learn. At that rate, a run that solves in 300 updates takes 10 s.

What this trainer leaves out

The core loop has three steps:

Act with the current policy.
Measure how much better each action did than the value network predicted.
Move probability toward the better actions.

CartPole is small enough that this loop fits in two compute shaders.

The trainer leaves out these standard techniques:

It treats the 500 step limit as a failure. A more accurate trainer uses $V_\phi$ at the cutoff, because the pole did not fall.
It doesn't normalize the observations or the advantages.
It takes one gradient step per batch. Proximal Policy Optimization (PPO) uses each batch for several steps and limits how far the policy moves.
It stops when the last 20 training episodes average 475 steps. Gymnasium's threshold for CartPole uses 100 episodes.