{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "\n# Train a World Model and Plan with It\n\nA world model predicts what an agent will see after it does an action. In this\ntutorial, we train LeWM from\n[LeWorldModel](https://arxiv.org/abs/2603.19312) and use it to move an agent\nto a goal.\n\nLeWM learns from frames and actions only. It does not use rewards or labels. An\nencoder turns each frame into an embedding. A predictor predicts the embedding\nof the next frame from the current embedding and the action. The SIGReg term of\nthe loss prevents the collapse of all embeddings to one point.\n\nThe environment is a small room with a wall. We write it in this tutorial, so\nthere is no dataset to download. Training takes about 5 minutes on one GPU.\n\nIn this tutorial you will learn:\n\n- How to collect episodes from an environment\n\n- How to train LeWM with lightly\n\n- How to make sure that the model uses the actions\n\n- How to read the predictions with a probe and a decoder\n\n- How to plan with the model to reach a goal\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Imports\n\nThis tutorial needs timm for the encoder and matplotlib for the plots.\n\n```console\npip install \"lightly[timm]\" matplotlib\n```\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import copy\nimport math\n\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport timm\nimport torch\nfrom matplotlib.patches import Rectangle\nfrom torch import nn\nfrom torch.utils.data import DataLoader, TensorDataset\n\nfrom lightly.loss import LeWMLoss\nfrom lightly.models.modules import (\n ActionEncoder,\n LatentDynamicsPredictor,\n LeWMProjectionHead,\n)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Configuration\n\nThe default configuration trains for 10 epochs on 1000 episodes. Training\nneeds about 2 GB of GPU memory. On a CPU, use fewer episodes and epochs to make\nthe run shorter.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "seed = 0\nnum_episodes = 1000\nepisode_length = 32\nclip_length = 4\nbatch_size = 128\nepochs = 10\nembed_dim = 192\n\nif torch.cuda.is_available():\n device = \"cuda\"\nelif torch.backends.mps.is_available():\n device = \"mps\"\nelse:\n device = \"cpu\"\n\ntorch.manual_seed(seed)\nrng = np.random.default_rng(seed)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## The environment\n\nThe environment is a room of 64x64 pixels. A vertical wall divides the room in\ntwo halves, and the agent can only go through the door in the wall. The state\nis the position of the top-left corner of the agent. An action is a vector\n``(dx, dy)`` in ``[-1, 1]``. It moves the agent by up to 6 pixels on each axis.\nThe agent stops when it touches the wall or the border.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "class WallWorld:\n size = 64 # Width and height of a frame in pixels.\n agent = 8 # Side length of the square agent.\n wall = (30.0, 34.0) # Columns of the wall.\n door = (24.0, 40.0) # Rows of the door in the wall.\n max_step = 6.0 # Largest movement on each axis in one step.\n\n def blocked(self, pos):\n x, y = pos\n in_wall_columns = x < self.wall[1] and x + self.agent > self.wall[0]\n in_door_rows = y >= self.door[0] and y + self.agent <= self.door[1]\n return in_wall_columns and not in_door_rows\n\n def random_position(self, rng):\n while True:\n pos = rng.uniform(0, self.size - self.agent, size=2)\n if not self.blocked(pos):\n return pos\n\n def step(self, pos, action):\n # Move along x, then along y. Stop at the wall on each axis.\n delta = np.clip(action, -1, 1) * self.max_step\n pos = np.array(pos, dtype=np.float64)\n pos[0] = np.clip(pos[0] + delta[0], 0, self.size - self.agent)\n if self.blocked(pos):\n pos[0] = self.wall[0] - self.agent if delta[0] > 0 else self.wall[1]\n pos[1] = np.clip(pos[1] + delta[1], 0, self.size - self.agent)\n if self.blocked(pos):\n pos[1] = self.door[1] - self.agent if delta[1] > 0 else self.door[0]\n return pos\n\n def render(self, positions):\n # Draw frames with shape (..., 3, 64, 64) for positions (..., 2).\n centers = np.arange(self.size) + 0.5\n x = positions[..., 0, None, None]\n y = positions[..., 1, None, None]\n agent = (centers[None, :] >= x) & (centers[None, :] < x + self.agent)\n agent = agent & (centers[:, None] >= y) & (centers[:, None] < y + self.agent)\n wall = (centers >= self.wall[0]) & (centers < self.wall[1])\n door = (centers >= self.door[0]) & (centers < self.door[1])\n wall = wall[None, :] & ~door[:, None]\n background = (0.08, 0.08, 0.12)\n wall_color = (0.55, 0.55, 0.6)\n agent_color = (1.0, 0.6, 0.1)\n frames = np.empty(positions.shape[:-1] + (3, self.size, self.size), np.float32)\n for c in range(3):\n frames[..., c, :, :] = np.where(\n agent, agent_color[c], np.where(wall, wall_color[c], background[c])\n )\n return frames\n\n\nworld = WallWorld()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Collect episodes\n\nLeWM learns from recorded episodes. It does not need a good policy, so we use a\nrandom one. The policy keeps its last action with a probability of 0.8. As a\nresult, the agent moves in straight lines for some steps and often goes through\nthe door.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "def collect(num_episodes, num_steps, rng, keep_prob=0.8):\n positions = np.empty((num_episodes, num_steps + 1, 2), np.float32)\n actions = np.empty((num_episodes, num_steps, 2), np.float32)\n for episode in range(num_episodes):\n pos = world.random_position(rng)\n action = rng.uniform(-1, 1, size=2)\n positions[episode, 0] = pos\n for t in range(num_steps):\n if rng.random() > keep_prob:\n action = rng.uniform(-1, 1, size=2)\n pos = world.step(pos=pos, action=action)\n actions[episode, t] = action\n positions[episode, t + 1] = pos\n return positions, actions\n\n\ntrain_positions, train_actions = collect(\n num_episodes=num_episodes, num_steps=episode_length, rng=rng\n)\ntest_positions, test_actions = collect(\n num_episodes=64, num_steps=episode_length, rng=rng\n)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "The plot shows one of the episodes. The agent slides along the wall and then\ngoes through the door.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "wall_center = sum(world.wall) / 2\nleft = train_positions[..., 0] + world.agent / 2 < wall_center\ncrossing = np.argmax(left.any(axis=1) & ~left.all(axis=1)) # 0 if none crosses\nframes = world.render(train_positions[crossing, ::4])\nfig, axes = plt.subplots(1, len(frames), figsize=(2 * len(frames), 2.4))\nfor t, (ax, frame) in enumerate(zip(axes, frames)):\n ax.imshow(frame.transpose(1, 2, 0))\n ax.set_title(f\"step {4 * t}\")\n ax.axis(\"off\")\nplt.show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "The training data are clips of ``clip_length`` consecutive frames and the\nactions between them. Each episode gives 30 clips that overlap. We store only\nthe positions and draw the frames when we need them.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "clip_starts = range(episode_length + 2 - clip_length)\nclip_positions = np.concatenate(\n [train_positions[:, s : s + clip_length] for s in clip_starts]\n)\nclip_actions = np.concatenate(\n [train_actions[:, s : s + clip_length - 1] for s in clip_starts]\n)\ndataset = TensorDataset(\n torch.from_numpy(clip_positions), torch.from_numpy(clip_actions)\n)\ndataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, drop_last=True)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## The model\n\nLeWM has four parts:\n\n- The encoder is a ViT-tiny from timm. It turns a frame into a vector of 192\n numbers.\n- The projection head maps this vector to the embedding.\n- The action encoder maps an action to a vector with the width of the\n embedding.\n- The predictor is a causal transformer. It predicts the embedding of the next\n frame from the embeddings and the actions so far.\n\nThe predictor sees at most 3 frames, which is the context of the training\nclips.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "class LeWM(nn.Module):\n def __init__(self, embed_dim, history):\n super().__init__()\n self.backbone = timm.create_model(\n \"vit_tiny_patch16_224\",\n pretrained=False,\n img_size=64,\n num_classes=0,\n dynamic_img_size=True,\n )\n self.projection_head = LeWMProjectionHead(\n input_dim=self.backbone.num_features, output_dim=embed_dim\n )\n self.action_encoder = ActionEncoder(action_dim=2, output_dim=embed_dim)\n self.predictor = LatentDynamicsPredictor(\n num_frames=history,\n input_dim=embed_dim,\n hidden_dim=384,\n depth=4,\n num_heads=6,\n )\n\n def encode(self, frames):\n # Encode frames (B, T, 3, H, W) to embeddings (B, T, D).\n batch_size, num_frames = frames.shape[:2]\n features = self.backbone(frames.flatten(0, 1))\n embeddings = self.projection_head(features)\n return embeddings.unflatten(0, (batch_size, num_frames))\n\n def predict(self, embeddings, actions):\n # Predict the next embedding for every frame (B, T, D).\n return self.predictor(embeddings, action_emb=self.action_encoder(actions))\n\n\nmodel = LeWM(embed_dim=embed_dim, history=clip_length - 1).to(device)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Train the model\n\nFor each clip, the predictor gets the embeddings of frames 1 to 3 and the\nthree actions. It predicts the embeddings of frames 2 to 4.\n:class:`lightly.loss.LeWMLoss` adds two terms. The prediction term is the mean\nsquared error between the predicted and the real embeddings. The SIGReg term\npushes the distribution of the embeddings to a standard Gaussian, so that the\nencoder cannot map all frames to one point.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "criterion = LeWMLoss(lambda_param=0.1).to(device) # Moves the SIGReg buffers.\noptimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.05)\ntotal_steps = epochs * len(dataloader)\nwarmup_steps = max(1, total_steps // 20)\nscheduler = torch.optim.lr_scheduler.LambdaLR(\n optimizer,\n lambda step: min(1.0, (step + 1) / warmup_steps)\n * 0.5\n * (1 + math.cos(math.pi * step / total_steps)),\n)\n\nhistory = []\nfor epoch in range(epochs):\n model.train()\n term_sums = torch.zeros(2, device=device)\n for positions, actions in dataloader:\n frames = torch.from_numpy(world.render(positions.numpy())).to(device)\n actions = actions.to(device)\n embeddings = model.encode(frames)\n predicted = model.predict(embeddings[:, :-1], actions)\n loss = criterion(\n predicted=predicted, target=embeddings[:, 1:], embeddings=embeddings\n )\n optimizer.zero_grad()\n loss.backward()\n optimizer.step()\n scheduler.step()\n # Log both terms for the plot. SIGReg expects the batch in dimension -2.\n with torch.no_grad():\n embeddings = embeddings.detach()\n prediction = (predicted - embeddings[:, 1:]).square().mean()\n sigreg = criterion.sigreg(embeddings.transpose(0, 1))\n term_sums += torch.stack([prediction, sigreg])\n history.append((term_sums / len(dataloader)).tolist())\n print(\n f\"epoch {epoch}: prediction {history[-1][0]:.4f}, SIGReg {history[-1][1]:.3f}\"\n )\n\nfig, axes = plt.subplots(1, 2, figsize=(10, 3.5))\nfor ax, values, title in zip(axes, zip(*history), [\"prediction term\", \"SIGReg term\"]):\n ax.plot(range(1, epochs + 1), values, marker=\"o\")\n ax.set_xlabel(\"epoch\")\n ax.set_title(title)\nfig.tight_layout()\nplt.show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Check the batch norm\n\nThe projection head has a batch norm layer. In training mode, it uses the\nstatistics of the batch. In eval mode, it uses running statistics. We plan in\neval mode, so the two modes must give almost the same embeddings. We compare\nthem on a batch of test frames. A copy of the head runs in training mode, so\nthat the running statistics do not change.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "model.eval()\nwith torch.no_grad():\n frames = torch.from_numpy(world.render(test_positions[:, 0])).to(device)\n features = model.backbone(frames)\n eval_embeddings = model.projection_head(features)\n train_embeddings = copy.deepcopy(model.projection_head).train()(features)\n gap = (train_embeddings - eval_embeddings).square().mean()\n gap = gap / eval_embeddings.var(dim=0).mean()\nprint(f\"difference between train and eval mode: {gap:.3f} of the variance\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Does the model use the actions?\n\nA low prediction loss does not prove that the model uses the actions. In many\nenvironments, the motion in the past frames predicts most of the next frame.\nTo make sure that the model uses the actions, we do rollouts on test episodes.\nA rollout starts from the embedding of one frame. The predictor predicts the\nnext embedding and then uses this prediction as its next input.\n\nWe compare three rollouts of 10 steps:\n\n- with the true actions of the episode\n- with the actions of a different episode\n- without the model: the embedding of the first frame, repeated\n\nThe plot shows the mean squared error to the real embeddings, divided by the\nvariance of the embeddings. A model that uses the actions has a much lower error\nwith the true actions.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "@torch.no_grad()\ndef encode_positions(positions, chunk_size=512):\n # Draw and encode frames for positions (..., 2). Returns (..., D) on the CPU.\n flat = positions.reshape(-1, 2)\n embeddings = []\n for i in range(0, len(flat), chunk_size):\n frames = torch.from_numpy(world.render(flat[i : i + chunk_size])).to(device)\n embeddings.append(model.encode(frames[None])[0].cpu())\n return torch.cat(embeddings).reshape(*positions.shape[:-1], -1)\n\n\n@torch.no_grad()\ndef rollout(context, actions):\n # Predict one embedding per action. context: (B, 1, D), actions: (B, H, 2).\n action_emb = model.action_encoder(actions)\n return model.predictor.rollout(\n context, action_emb=action_emb, steps=actions.shape[1]\n )\n\n\nhorizon = 10\ntest_embeddings = encode_positions(test_positions).to(device)\ntest_actions_t = torch.from_numpy(test_actions).to(device)\ncontext = test_embeddings[:, :1]\ntarget = test_embeddings[:, 1 : horizon + 1]\nvariance = target.var(dim=(0, 1)).mean()\nrollouts = {\n \"true actions\": rollout(context=context, actions=test_actions_t[:, :horizon]),\n \"actions of a different episode\": rollout(\n context=context, actions=test_actions_t[:, :horizon].roll(1, dims=0)\n ),\n \"first embedding, repeated\": context.expand_as(target),\n}\nplt.figure(figsize=(6, 4))\nfor label, predicted in rollouts.items():\n error = (predicted - target).square().mean(dim=(0, 2)) / variance\n plt.plot(range(1, horizon + 1), error.cpu(), marker=\"o\", label=label)\nplt.xlabel(\"rollout step\")\nplt.ylabel(\"normalized error\")\nplt.legend()\nplt.show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Read the predictions with a probe\n\nAn embedding is a vector of 192 numbers, and we do not know what each number\nmeans. A probe is a small model that reads one property from the embedding,\nhere the position of the agent. We know the position for every frame because\nwe made the frames. Thus, we can train the probe with this position as the\ntarget. The probe is not part of the world model. We train it after the world\nmodel, and the encoder does not change.\n\nThe LeWM paper uses linear probes and MLP probes. Here, we use an MLP with one\nhidden layer, because it reads the position more exactly. If the probe gives\nthe position of new frames with a small error, the embedding contains the\nposition.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "probe_positions = train_positions[:200].reshape(-1, 2)\nprobe_embeddings = encode_positions(probe_positions)\nprobe_targets = torch.from_numpy(probe_positions) / world.size\nprobe = nn.Sequential(nn.Linear(embed_dim, 256), nn.ReLU(), nn.Linear(256, 2))\nprobe_optimizer = torch.optim.Adam(probe.parameters(), lr=1e-3)\nfor step in range(2000):\n loss = (probe(probe_embeddings) - probe_targets).square().mean()\n probe_optimizer.zero_grad()\n loss.backward()\n probe_optimizer.step()\n\n\n@torch.no_grad()\ndef read_position(embeddings):\n # Apply the probe to embeddings (..., D). Returns positions (..., 2) in pixels.\n return probe(embeddings.cpu()) * world.size\n\n\nerror = read_position(test_embeddings) - torch.from_numpy(test_positions)\nprint(f\"probe error on test frames: {error.norm(dim=-1).mean():.1f} pixels\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Now we apply the same probe to the embeddings of a rollout. These embeddings\ncome from the predictor, not from a frame. Thus, the probe shows where the\nmodel expects the agent to be.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "imagined = read_position(rollouts[\"true actions\"])\nerror = imagined - torch.from_numpy(test_positions[:, 1 : horizon + 1])\nprint(f\"probe error on predicted embeddings: {error.norm(dim=-1).mean():.1f} pixels\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "The wall is a good test for the model. In the plot, the agent moves right for 10\nsteps from two start positions. On the left, the wall blocks the agent. On the\nright, the agent goes through the door. The dashed line is the rollout, read\nwith the probe.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "def push_right(starts, num_steps=10):\n actions = np.tile(np.float32([1.0, 0.0]), (len(starts), num_steps, 1))\n real = [starts]\n for t in range(num_steps):\n real.append(\n np.stack(\n [world.step(pos=p, action=a) for p, a in zip(real[-1], actions[:, t])]\n )\n )\n context = encode_positions(starts[:, None]).to(device)\n predicted = rollout(context=context, actions=torch.from_numpy(actions).to(device))\n real = np.stack(real, axis=1)\n imagined = read_position(predicted).numpy()\n imagined = np.concatenate([starts[:, None], imagined], axis=1)\n return real, imagined\n\n\nreal, imagined = push_right(np.array([[10.0, 8.0], [10.0, 28.0]]))\nfig, axes = plt.subplots(1, 2, figsize=(8, 4))\nfor ax, real_path, imagined_path in zip(axes, real, imagined):\n ax.imshow(world.render(real_path[0]).transpose(1, 2, 0), extent=(0, 64, 64, 0))\n center = world.agent / 2\n ax.plot(*(real_path + center).T, \"w-o\", markersize=3, label=\"real\")\n ax.plot(*(imagined_path + center).T, \"c--o\", markersize=3, label=\"rollout\")\n ax.axis(\"off\")\naxes[0].legend(loc=\"lower left\")\nplt.show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "The same test with 200 random start positions in the left half of the room\nshows how often the rollout agrees with the environment on one question: does\nthe agent get through the wall?\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "starts = np.stack([rng.uniform(0, 22, 200), rng.uniform(0, 56, 200)], axis=1)\nreal, imagined = push_right(starts)\nthrough_real = real[:, -1, 0] + world.agent / 2 > wall_center\nthrough_imagined = imagined[:, -1, 0] + world.agent / 2 > wall_center\nprint(\n f\"agreement with the environment: {np.mean(through_real == through_imagined):.0%}\"\n)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Look at the predictions with a decoder\n\nA decoder turns an embedding back into a frame. LeWM does not need a decoder,\nand the decoder does not change the world model. We train it after the world\nmodel, on frozen embeddings, only to look at the predictions.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "class FrameDecoder(nn.Module):\n def __init__(self, embed_dim):\n super().__init__()\n self.project = nn.Linear(embed_dim, 512 * 4 * 4)\n layers = []\n for in_channels, out_channels in [(512, 256), (256, 128), (128, 64), (64, 64)]:\n layers += [\n nn.Upsample(scale_factor=2),\n nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),\n nn.GroupNorm(32, out_channels),\n nn.SiLU(),\n ]\n layers += [nn.Conv2d(64, 3, kernel_size=3, padding=1), nn.Sigmoid()]\n self.layers = nn.Sequential(*layers)\n\n def forward(self, embeddings):\n return self.layers(self.project(embeddings).view(-1, 512, 4, 4))\n\n\ndecoder = FrameDecoder(embed_dim).to(device)\ndecoder_optimizer = torch.optim.AdamW(decoder.parameters(), lr=1e-3)\ndecoder_positions = train_positions.reshape(-1, 2)\ndecoder_embeddings = encode_positions(decoder_positions)\nfor step in range(1000):\n index = rng.integers(0, len(decoder_positions), 256)\n frames = torch.from_numpy(world.render(decoder_positions[index])).to(device)\n reconstructed = decoder(decoder_embeddings[index].to(device))\n loss = nn.functional.mse_loss(reconstructed, frames)\n decoder_optimizer.zero_grad()\n loss.backward()\n decoder_optimizer.step()\nprint(f\"decoder error: {loss.item():.5f}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "The first row shows a test episode. The second row decodes a rollout with the\ntrue actions. The third row decodes a rollout with the opposite actions. All\nrows start from the same frame.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "opposite = rollout(context=context[:1], actions=-test_actions_t[:1, :horizon])\nwith torch.no_grad():\n first = torch.from_numpy(world.render(test_positions[0, :1])).to(device)\n rows = {\n \"real\": torch.from_numpy(world.render(test_positions[0, : horizon + 1])),\n \"true actions\": torch.cat([first, decoder(rollouts[\"true actions\"][0])]),\n \"opposite actions\": torch.cat([first, decoder(opposite[0])]),\n }\nfig, axes = plt.subplots(3, horizon + 1, figsize=(horizon + 1, 3.6))\nfor row, (label, frames) in enumerate(rows.items()):\n for t in range(horizon + 1):\n axes[row, t].imshow(frames[t].cpu().numpy().transpose(1, 2, 0).clip(0, 1))\n axes[row, t].axis(\"off\")\n axes[row, 0].set_title(label, fontsize=8, loc=\"left\")\nplt.show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Plan with the model\n\nPlanning uses the model to select actions. The planner gets the current frame\nand a goal frame. It tries many action sequences in the model and keeps the\nsequence whose predicted embeddings are nearest to the embedding of the goal\nframe. The model does not change during planning.\n\nWe use the cross-entropy method (CEM). A plan is a sequence of 8 actions. CEM\nkeeps a Gaussian distribution over plans: ``mean`` is the best plan so far,\nand ``std`` is the range around it that CEM still tries. Each iteration does\nthese steps:\n\n1. Sample 300 plans from the distribution.\n2. Do a rollout for each plan from the current embedding.\n3. Calculate the cost of each plan: the mean squared distance between the\n predicted embeddings and the goal embedding, averaged over the 8 steps.\n4. Set ``mean`` and ``std`` to the mean and standard deviation of the 30 plans\n with the lowest cost.\n\nAfter 10 iterations, ``mean`` is the plan. ``plan`` makes plans for many\nepisodes at the same time, so every tensor has a batch dimension ``B``.\n\n

Note

The LeWM paper uses the distance at the last step of the plan only. The\n average over all steps makes the agent go to the goal early and stay\n there. On this environment, it gives a much higher success rate.

\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "plan_horizon = 8\n\n\n@torch.no_grad()\ndef plan(\n current, goal, horizon=plan_horizon, num_samples=300, num_iters=10, num_elites=30\n):\n # current, goal: embeddings (B, D). Returns plans (B, horizon, 2).\n batch, dim = current.shape\n mean = torch.zeros(batch, horizon, 2, device=device)\n std = torch.ones(batch, horizon, 2, device=device)\n context = current.repeat_interleave(num_samples, dim=0).unsqueeze(1)\n for _ in range(num_iters):\n noise = torch.randn(batch, num_samples, horizon, 2, device=device)\n candidates = (mean[:, None] + std[:, None] * noise).clamp(-1, 1)\n predicted = rollout(context=context, actions=candidates.flatten(0, 1))\n predicted = predicted.view(batch, num_samples, horizon, dim)\n costs = (predicted - goal[:, None, None]).square().mean(dim=(2, 3))\n # Refit the distribution to the plans with the lowest cost.\n elite_idx = torch.topk(costs, num_elites, dim=1, largest=False).indices\n elites = torch.take_along_dim(candidates, elite_idx[:, :, None, None], dim=1)\n mean = elites.mean(dim=1)\n std = elites.std(dim=1)\n return mean" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "The agent does the first 4 actions of the plan in the environment. Then the\nplanner makes a new plan from the new frame. This loop is model-predictive\ncontrol. It corrects the errors of the model before they add up.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "def run_episodes(starts, goals, policy, num_steps, replan_every=4):\n positions = starts.copy()\n path = [positions]\n for _ in range(num_steps // replan_every):\n actions = policy(positions=positions, goals=goals)\n for t in range(replan_every):\n positions = np.stack(\n [world.step(pos=p, action=a) for p, a in zip(positions, actions[:, t])]\n )\n path.append(positions)\n return np.stack(path, axis=1)\n\n\ndef cem_policy(positions, goals):\n current = encode_positions(positions).to(device)\n goal = encode_positions(goals).to(device)\n return plan(current=current, goal=goal).cpu().numpy()\n\n\ndef random_policy(positions, goals):\n return rng.uniform(-1, 1, size=(len(positions), plan_horizon, 2))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "We evaluate like the LeWM paper. The goal is the frame ``k`` steps later in a\nnew episode, and the agent has ``2 * k`` steps to reach it. An episode is a\nsuccess if the agent ends less than 4 pixels from the goal. A random policy is\nthe baseline.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "results = {}\nfor k in [4, 8, 16]:\n goal_positions, _ = collect(num_episodes=400, num_steps=k, rng=rng)\n distance = np.linalg.norm(goal_positions[:, k] - goal_positions[:, 0], axis=1)\n keep = np.flatnonzero(distance >= 8)[:30]\n starts, goals = goal_positions[keep, 0], goal_positions[keep, k]\n for name, policy in [(\"random actions\", random_policy), (\"CEM\", cem_policy)]:\n paths = run_episodes(starts=starts, goals=goals, policy=policy, num_steps=2 * k)\n miss = np.linalg.norm(paths[:, -1] - goals, axis=1)\n results[k, name] = (paths, goals, miss)\n print(\n f\"goal {k} steps ahead, {name}: success {np.mean(miss < 4):.0%}, \"\n f\"median distance to the goal {np.median(miss):.1f} pixels\"\n )" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "The plot shows four CEM episodes with the goal 8 steps ahead: two that reach\nthe goal and two that miss it. The outline is the goal position of the agent.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "paths, goals, miss = results[8, \"CEM\"]\nepisodes = [*np.flatnonzero(miss < 4)[:2], *np.flatnonzero(miss >= 4)[:2]]\nfig, axes = plt.subplots(1, len(episodes), figsize=(3 * len(episodes), 3.4))\nfor ax, episode in zip(axes, episodes):\n path, goal = paths[episode], goals[episode]\n ax.imshow(world.render(path[0]).transpose(1, 2, 0), extent=(0, 64, 64, 0))\n ax.add_patch(Rectangle(goal, world.agent, world.agent, fill=False, color=\"w\"))\n ax.plot(*(path + world.agent / 2).T, \"c-o\", markersize=3)\n ax.set_title(\"reached\" if miss[episode] < 4 else \"missed\")\n ax.axis(\"off\")\nplt.show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Limits of the planner\n\nThe planner works for near goals. For goals 16 steps ahead, the success rate is\nlower. The plot shows the reason. For pairs of random frames, it compares the\ndistance between the two embeddings with the distance between the two agent\npositions.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "positions = np.stack([world.random_position(rng) for _ in range(2000)])\nembeddings = encode_positions(positions)\ni, j = rng.integers(0, len(positions), size=(2, 20000))\npixel_distance = np.linalg.norm(positions[i] - positions[j], axis=1)\nembedding_distance = (embeddings[i] - embeddings[j]).square().mean(dim=1).numpy()\nbins = np.arange(0, 60, 4)\ncenters = [low + 2 for low in bins]\nmeans = [\n embedding_distance[(pixel_distance >= low) & (pixel_distance < low + 4)].mean()\n for low in bins\n]\nplt.figure(figsize=(6, 4))\nplt.plot(centers, means, marker=\"o\")\nplt.xlabel(\"distance between the agent positions (pixels)\")\nplt.ylabel(\"mean squared distance between the embeddings\")\nplt.show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "The embedding distance increases quickly for the first 10 pixels and then stays\nalmost constant. For a goal that is 30 pixels away, most plans end in the flat\npart of the curve and have almost the same cost. CEM then finds the goal only if\nsome of the sampled plans end near it. The LeWM paper reports\na similar result on its Two-Room environment. The authors explain it with the\nlow diversity and the low intrinsic dimension of the data. In this environment,\nthe position of the agent is the only property that changes.\n\nOther methods change the cost or the embeddings so that a planner can reach far\ngoals:\n\n- [TEMPO](https://arxiv.org/abs/2610.04988) trains a small model on frozen\n LeWM embeddings. In its space, the distance between two states is the number\n of steps between them. The planner uses this distance as the cost.\n- [PLDM](https://arxiv.org/abs/2502.14819) trains with a loss term that keeps\n consecutive embeddings close, so that the embedding distance changes smoothly\n over time.\n- [Search on the Replay Buffer](https://arxiv.org/abs/1906.05253) builds a\n graph of recorded states and plans a path of waypoints to the goal.\n\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Next Steps\n\n- `lewm` describes the LeWM modules and the shapes of their inputs.\n- `lejepa` uses the same SIGReg term for self-supervised learning on\n images.\n- The [official LeWM code](https://github.com/lucas-maes/le-wm) trains LeWM\n on larger environments, for example Push-T and OGBench Cube.\n\n" ] } ], "metadata": { "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.12.9" } }, "nbformat": 4, "nbformat_minor": 0 }