Learning a model to plan with

Learn the model, then plan

Q-learning and REINFORCE learned what to do straight from many trials. Model-based RL learns how the robot responds to its actions instead, and then plans on that model. Here the cart-pole's equations are hidden (cartpole.linearize() is switched off). The program pushes the cart at random for 30 short episodes and records every transition $(s, u, s')$: the state $s = [x, \dot x, \theta, \dot\theta]$ (cart position in m, its speed, the pole's angle from upright in rad, its rate), the push $u$ (N), and the state $s'$ one control period (0.02 s) later.

A model by least squares

The model predicts the change, $s' - s \approx \phi(s, u)^\top W$: features times weights, step 4's regression with four outputs. The features decide what it can express. features_linear is $[1, s, u]$, a straight-line model that is good near upright. features_physics adds $\sin\theta$, $\dot\theta^2 \sin\theta$, $u \cos\theta$ and $\sin\theta \cos\theta$, the terms in the cart-pole's equations of motion. The Plots tab shows each model's one-step error in $\dot\theta$ against $\theta$: the linear model's grows as the pole leans over.

Planning by sampling

The planner looks $H = 20$ periods ahead. sequence_costs rolls a sequence of pushes $u_0 \dots u_{H-1}$ through the model from the current state and adds up step_cost, which charges for angle, position, speed and effort. The cross-entropy method (CEM) searches over sequences: sample many from a Gaussian with a mean and standard deviation for each step, clip them to the force limit, cost them all, keep the cheapest few (the elites) and refit the mean and standard deviation to them. A few rounds shrink the Gaussian onto a good sequence. The program applies its first push, steps once and plans again from the new state: model predictive control.

Your task

1. Write fit_model(S, U, S2, features), least squares of S2 - S on features(S, U), with least_squares.
2. Write elite_refit(U, costs, n_elites).
3. Write cem(costs, mean, std, iters, samples, n_elites, u_max, rng), sampling with rng.standard_normal.

The program then catches the pole from 12–16° with the physics model (afterwards, try MODEL = "linear": the cart usually runs into a stop). From 1.5 s on it must stay within 5° of upright, and the cart must not hit a stop.