Model-based RL

What it does

Model-based RL learns how the robot responds to its actions, then chooses actions by trying them out on that model instead of on the robot. Here cartpole.linearize() is switched off: the program pushes the cart at random for 40 short episodes, and you learn the dynamics from what happened.

Learning the model

The model predicts the change in state over one 20 ms period as a weighted sum of 10 given features: $s_{k+1} = s_k + \phi(s_k, u_k)^T W$. The features include $\sin\theta$, $\dot\theta^2 \sin\theta$ and $u \cos\theta$, the shapes of the real equations of motion. Stack the transitions and solve by least squares:

Planning with it: CEM

Every period, the cross-entropy method looks for H = 20 pushes (0.4 s) that keep the pole up. It keeps a Gaussian over sequences, with a mean and a std for each step:

1. Sample N sequences from it and clip them to ±u_max.
2. Roll each through the model and add up step_cost after every push.
3. Keep the E cheapest (the elites) and refit the mean and std to them.

After a few iterations, push with the mean's first step. Then shift the plan one step for next time (MPC).

Tools

np.linalg.lstsq(X, Y, rcond=None)[0] solves least squares. rng.standard_normal((N, H)) samples, np.clip clips, np.argsort(c)[:E] picks the cheapest, and .mean(axis=0) and .std(axis=0) refit.

Your task

Implement fit_model(S, U, S2), which returns W (10 × 4); sequence_costs(W, s0, U), the total cost of each row of U (N × H) from state s0; and cem(costs, mean, std, iters, samples, elites, u_max, rng), which returns the final mean. Call costs(U) once per iteration, on all the samples together. Planning on your model must catch the pole from 15–18° and keep it within 4° of upright from 1.5 s on, without hitting the stops.