World model in code · Part 3 of 3

Learning dynamics and planning in dreams

August 24, 2026 · approx. 8 min read

A state and an action turn into a possible future. This part covers the dynamics network, multi-step predictions, and the planner that pits 100 action sequences against each other.

1 · DataEnvironment and transitions 2 · Seeing49 pixels to 8 numbers 3 · PlanningDynamics and dreams

Four actions become four switches

The latent state has eight values. The action is appended as a one-hot vector: exactly one of four values is set to 1. That gives the dynamics network twelve input values.

def dynamik_eingabe(z, aktion):
    einheit = [0.0] * 4
    einheit[aktion] = 1.0
    return list(z) + einheit

# Example for "left":
# [z0, z1, ... z7, 0, 0, 1, 0]

A single action number would also be a number, but it would imply a false ordering. "Right" isn't twice as much as "down". One-hot encoding treats all four directions as separate categories.

Ten outputs describe the next moment

The network outputs eight values for the next state, one for the reward, and one for whether the episode ends.

Dynamiknetz: 12 → 32 → 10

Output 0 to 7: next latent state
Output 8:      expected reward
Output 9:      probability of episode end

State, reward, and end-of-episode each need a different value range. tanh bounds the state between −1 and 1. A sigmoid turns the last two raw values into probabilities between 0 and 1.

def dynamik_ausgabe(roh):
    z = [math.tanh(w) for w in roh[:8]]
    belohnung = sigmoid(roh[8])
    fertig = sigmoid(roh[9])
    return z, belohnung, fertig

Training knows the true future

The transitions from part 1 contain the image before, the action, and the image after. The trained encoder from part 2 translates both images into latent states. The dynamics network is trained to predict the second state from the first state and the action.

eingabe = dynamik_eingabe(
    kodiere(modell, uebergang["vorher"]),
    uebergang["aktion"])

ziel_z = kodiere(modell, uebergang["nachher"])
ziel_r = uebergang["belohnung"]
ziel_f = uebergang["fertig"]

The reward appears in only 1.3 percent of training rows. Its error is therefore multiplied by 10. Without this counterweight, "always zero" would be a convenient solution for the network.

Weighting changes the learning objective. The factor of 10 is a decision made for this experiment. It makes rare rewards more visible, but it can also amplify false rewards. Exactly these errors show up later in the planner.

A dream step doesn't ask the environment

Once trained, a prediction needs only the dynamics network:

def traumschritt(modell, z, aktion):
    x = dynamik_eingabe(z, aktion)
    roh, _ = vorwaerts(modell["dynamik"], x)
    return dynamik_ausgabe(roh)

For a longer dream, the predicted state becomes the next input again. After the first image, the model no longer receives any real observation.

z = kodiere(modell, erstes_bild)

for aktion in aktionsfolge:
    z, belohnung, fertig = traumschritt(modell, z, aktion)

As a result, errors grow with every step. In the measurement, the predicted position was correct in 95.0 percent of cases after one step, and in 71.3 percent after ten steps.

The planner tries 100 sequences

Random shooting is a short, readable planning method. It generates random action sequences, plays each one out in the model, and keeps the sequence with the highest predicted reward.

def plane(modell, z, horizont, versuche, zufall):
    beste_folge = None
    bester_wert = -1e9

    for _ in range(versuche):
        folge = [zufall.randrange(4)
                 for _ in range(horizont)]
        zt = z
        wert = 0.0

        for k, aktion in enumerate(folge):
            zt, r, fertig = traumschritt(modell, zt, aktion)
            wert += (0.95 ** k) * r

        if wert > bester_wert:
            beste_folge, bester_wert = folge, wert

    return beste_folge[0], bester_wert

The factor 0.95 ** k weighs early rewards a bit higher than later ones. The best sequence is not played out in full in the real environment. If the planner picked, say, [right, right, down, left, …], return beste_folge[0] only returns right. The remaining steps exist solely to score that first move.

After every real move, planning starts over

The returned action goes into aktion and is passed to schritt(pos, aktion). The agent carries out this single move in the real environment, reads in the new image, and restarts the search. The rest of the sequence is discarded. This limits the damage from model errors that grow over longer horizons.

while not fertig:
    z = kodiere(modell, beobachtung(pos))
    aktion, wert = plane(modell, z, 6, 100, zufall)
    pos, belohnung, fertig = schritt(pos, aktion)

This approach is called model predictive control. The horizon here is six steps, but the plan is corrected after every single step.

When an error helps the planner

The world model sometimes predicted small rewards at places where the environment gives none. In this run, these errors often lay along paths toward the goal. That gave the planner a gradient to follow, even though the real reward exists on only a single cell.

A good planning result doesn't prove an accurate model. With a six-step horizon, the fabricated reward landscape helped. With a three-step horizon, it led the agent in the wrong direction. Model error and planning performance have to be measured separately.

The complete chain

Image
  → Encoder: 49 values to 8
  → Dynamics: state + action to next state
  → Rollout: multiple dynamics steps
  → Planner: compare 100 rollouts
  → first action in the real world
  → new image, and start over

The environment supplies transitions. The autoencoder compresses images, the dynamics network predicts their consequences, and the planner picks the next move from that.