.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "tutorials/package/tutorial_lewm.py" .. LINE NUMBERS ARE GIVEN BELOW. .. only:: html .. note:: :class: sphx-glr-download-link-note :ref:`Go to the end ` to download the full example code. .. rst-class:: sphx-glr-example-title .. _sphx_glr_tutorials_package_tutorial_lewm.py: .. _lightly-lewm-tutorial-9: Train a World Model and Plan with It ================================================ A world model predicts what an agent will see after it does an action. In this tutorial, we train LeWM from `LeWorldModel `_ and use it to move an agent to a goal. LeWM learns from frames and actions only. It does not use rewards or labels. An encoder turns each frame into an embedding. A predictor predicts the embedding of the next frame from the current embedding and the action. The SIGReg term of the loss prevents the collapse of all embeddings to one point. The environment is a small room with a wall. We write it in this tutorial, so there is no dataset to download. Training takes about 5 minutes on one GPU. In this tutorial you will learn: - How to collect episodes from an environment - How to train LeWM with lightly - How to make sure that the model uses the actions - How to read the predictions with a probe and a decoder - How to plan with the model to reach a goal .. GENERATED FROM PYTHON SOURCE LINES 34-42 Imports ------- This tutorial needs timm for the encoder and matplotlib for the plots. .. code-block:: console pip install "lightly[timm]" matplotlib .. GENERATED FROM PYTHON SOURCE LINES 42-60 .. code-block:: Python import copy import math import matplotlib.pyplot as plt import numpy as np import timm import torch from matplotlib.patches import Rectangle from torch import nn from torch.utils.data import DataLoader, TensorDataset from lightly.loss import LeWMLoss from lightly.models.modules import ( ActionEncoder, LatentDynamicsPredictor, LeWMProjectionHead, ) .. GENERATED FROM PYTHON SOURCE LINES 61-67 Configuration ------------- The default configuration trains for 10 epochs on 1000 episodes. Training needs about 2 GB of GPU memory. On a CPU, use fewer episodes and epochs to make the run shorter. .. GENERATED FROM PYTHON SOURCE LINES 67-85 .. code-block:: Python seed = 0 num_episodes = 1000 episode_length = 32 clip_length = 4 batch_size = 128 epochs = 10 embed_dim = 192 if torch.cuda.is_available(): device = "cuda" elif torch.backends.mps.is_available(): device = "mps" else: device = "cpu" torch.manual_seed(seed) rng = np.random.default_rng(seed) .. GENERATED FROM PYTHON SOURCE LINES 86-94 The environment --------------- The environment is a room of 64x64 pixels. A vertical wall divides the room in two halves, and the agent can only go through the door in the wall. The state is the position of the top-left corner of the agent. An action is a vector ``(dx, dy)`` in ``[-1, 1]``. It moves the agent by up to 6 pixels on each axis. The agent stops when it touches the wall or the border. .. GENERATED FROM PYTHON SOURCE LINES 94-150 .. code-block:: Python class WallWorld: size = 64 # Width and height of a frame in pixels. agent = 8 # Side length of the square agent. wall = (30.0, 34.0) # Columns of the wall. door = (24.0, 40.0) # Rows of the door in the wall. max_step = 6.0 # Largest movement on each axis in one step. def blocked(self, pos): x, y = pos in_wall_columns = x < self.wall[1] and x + self.agent > self.wall[0] in_door_rows = y >= self.door[0] and y + self.agent <= self.door[1] return in_wall_columns and not in_door_rows def random_position(self, rng): while True: pos = rng.uniform(0, self.size - self.agent, size=2) if not self.blocked(pos): return pos def step(self, pos, action): # Move along x, then along y. Stop at the wall on each axis. delta = np.clip(action, -1, 1) * self.max_step pos = np.array(pos, dtype=np.float64) pos[0] = np.clip(pos[0] + delta[0], 0, self.size - self.agent) if self.blocked(pos): pos[0] = self.wall[0] - self.agent if delta[0] > 0 else self.wall[1] pos[1] = np.clip(pos[1] + delta[1], 0, self.size - self.agent) if self.blocked(pos): pos[1] = self.door[1] - self.agent if delta[1] > 0 else self.door[0] return pos def render(self, positions): # Draw frames with shape (..., 3, 64, 64) for positions (..., 2). centers = np.arange(self.size) + 0.5 x = positions[..., 0, None, None] y = positions[..., 1, None, None] agent = (centers[None, :] >= x) & (centers[None, :] < x + self.agent) agent = agent & (centers[:, None] >= y) & (centers[:, None] < y + self.agent) wall = (centers >= self.wall[0]) & (centers < self.wall[1]) door = (centers >= self.door[0]) & (centers < self.door[1]) wall = wall[None, :] & ~door[:, None] background = (0.08, 0.08, 0.12) wall_color = (0.55, 0.55, 0.6) agent_color = (1.0, 0.6, 0.1) frames = np.empty(positions.shape[:-1] + (3, self.size, self.size), np.float32) for c in range(3): frames[..., c, :, :] = np.where( agent, agent_color[c], np.where(wall, wall_color[c], background[c]) ) return frames world = WallWorld() .. GENERATED FROM PYTHON SOURCE LINES 151-158 Collect episodes ---------------- LeWM learns from recorded episodes. It does not need a good policy, so we use a random one. The policy keeps its last action with a probability of 0.8. As a result, the agent moves in straight lines for some steps and often goes through the door. .. GENERATED FROM PYTHON SOURCE LINES 158-183 .. code-block:: Python def collect(num_episodes, num_steps, rng, keep_prob=0.8): positions = np.empty((num_episodes, num_steps + 1, 2), np.float32) actions = np.empty((num_episodes, num_steps, 2), np.float32) for episode in range(num_episodes): pos = world.random_position(rng) action = rng.uniform(-1, 1, size=2) positions[episode, 0] = pos for t in range(num_steps): if rng.random() > keep_prob: action = rng.uniform(-1, 1, size=2) pos = world.step(pos=pos, action=action) actions[episode, t] = action positions[episode, t + 1] = pos return positions, actions train_positions, train_actions = collect( num_episodes=num_episodes, num_steps=episode_length, rng=rng ) test_positions, test_actions = collect( num_episodes=64, num_steps=episode_length, rng=rng ) .. GENERATED FROM PYTHON SOURCE LINES 184-186 The plot shows one of the episodes. The agent slides along the wall and then goes through the door. .. GENERATED FROM PYTHON SOURCE LINES 186-197 .. code-block:: Python wall_center = sum(world.wall) / 2 left = train_positions[..., 0] + world.agent / 2 < wall_center crossing = np.argmax(left.any(axis=1) & ~left.all(axis=1)) # 0 if none crosses frames = world.render(train_positions[crossing, ::4]) fig, axes = plt.subplots(1, len(frames), figsize=(2 * len(frames), 2.4)) for t, (ax, frame) in enumerate(zip(axes, frames)): ax.imshow(frame.transpose(1, 2, 0)) ax.set_title(f"step {4 * t}") ax.axis("off") plt.show() .. image-sg:: /tutorials/package/images/sphx_glr_tutorial_lewm_001.png :alt: step 0, step 4, step 8, step 12, step 16, step 20, step 24, step 28, step 32 :srcset: /tutorials/package/images/sphx_glr_tutorial_lewm_001.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 198-201 The training data are clips of ``clip_length`` consecutive frames and the actions between them. Each episode gives 30 clips that overlap. We store only the positions and draw the frames when we need them. .. GENERATED FROM PYTHON SOURCE LINES 201-213 .. code-block:: Python clip_starts = range(episode_length + 2 - clip_length) clip_positions = np.concatenate( [train_positions[:, s : s + clip_length] for s in clip_starts] ) clip_actions = np.concatenate( [train_actions[:, s : s + clip_length - 1] for s in clip_starts] ) dataset = TensorDataset( torch.from_numpy(clip_positions), torch.from_numpy(clip_actions) ) dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, drop_last=True) .. GENERATED FROM PYTHON SOURCE LINES 214-229 The model --------- LeWM has four parts: - The encoder is a ViT-tiny from timm. It turns a frame into a vector of 192 numbers. - The projection head maps this vector to the embedding. - The action encoder maps an action to a vector with the width of the embedding. - The predictor is a causal transformer. It predicts the embedding of the next frame from the embeddings and the actions so far. The predictor sees at most 3 frames, which is the context of the training clips. .. GENERATED FROM PYTHON SOURCE LINES 229-267 .. code-block:: Python class LeWM(nn.Module): def __init__(self, embed_dim, history): super().__init__() self.backbone = timm.create_model( "vit_tiny_patch16_224", pretrained=False, img_size=64, num_classes=0, dynamic_img_size=True, ) self.projection_head = LeWMProjectionHead( input_dim=self.backbone.num_features, output_dim=embed_dim ) self.action_encoder = ActionEncoder(action_dim=2, output_dim=embed_dim) self.predictor = LatentDynamicsPredictor( num_frames=history, input_dim=embed_dim, hidden_dim=384, depth=4, num_heads=6, ) def encode(self, frames): # Encode frames (B, T, 3, H, W) to embeddings (B, T, D). batch_size, num_frames = frames.shape[:2] features = self.backbone(frames.flatten(0, 1)) embeddings = self.projection_head(features) return embeddings.unflatten(0, (batch_size, num_frames)) def predict(self, embeddings, actions): # Predict the next embedding for every frame (B, T, D). return self.predictor(embeddings, action_emb=self.action_encoder(actions)) model = LeWM(embed_dim=embed_dim, history=clip_length - 1).to(device) .. GENERATED FROM PYTHON SOURCE LINES 268-277 Train the model --------------- For each clip, the predictor gets the embeddings of frames 1 to 3 and the three actions. It predicts the embeddings of frames 2 to 4. :class:`lightly.loss.LeWMLoss` adds two terms. The prediction term is the mean squared error between the predicted and the real embeddings. The SIGReg term pushes the distribution of the embeddings to a standard Gaussian, so that the encoder cannot map all frames to one point. .. GENERATED FROM PYTHON SOURCE LINES 277-323 .. code-block:: Python criterion = LeWMLoss(lambda_param=0.1).to(device) # Moves the SIGReg buffers. optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.05) total_steps = epochs * len(dataloader) warmup_steps = max(1, total_steps // 20) scheduler = torch.optim.lr_scheduler.LambdaLR( optimizer, lambda step: min(1.0, (step + 1) / warmup_steps) * 0.5 * (1 + math.cos(math.pi * step / total_steps)), ) history = [] for epoch in range(epochs): model.train() term_sums = torch.zeros(2, device=device) for positions, actions in dataloader: frames = torch.from_numpy(world.render(positions.numpy())).to(device) actions = actions.to(device) embeddings = model.encode(frames) predicted = model.predict(embeddings[:, :-1], actions) loss = criterion( predicted=predicted, target=embeddings[:, 1:], embeddings=embeddings ) optimizer.zero_grad() loss.backward() optimizer.step() scheduler.step() # Log both terms for the plot. SIGReg expects the batch in dimension -2. with torch.no_grad(): embeddings = embeddings.detach() prediction = (predicted - embeddings[:, 1:]).square().mean() sigreg = criterion.sigreg(embeddings.transpose(0, 1)) term_sums += torch.stack([prediction, sigreg]) history.append((term_sums / len(dataloader)).tolist()) print( f"epoch {epoch}: prediction {history[-1][0]:.4f}, SIGReg {history[-1][1]:.3f}" ) fig, axes = plt.subplots(1, 2, figsize=(10, 3.5)) for ax, values, title in zip(axes, zip(*history), ["prediction term", "SIGReg term"]): ax.plot(range(1, epochs + 1), values, marker="o") ax.set_xlabel("epoch") ax.set_title(title) fig.tight_layout() plt.show() .. image-sg:: /tutorials/package/images/sphx_glr_tutorial_lewm_002.png :alt: prediction term, SIGReg term :srcset: /tutorials/package/images/sphx_glr_tutorial_lewm_002.png :class: sphx-glr-single-img .. rst-class:: sphx-glr-script-out .. code-block:: none epoch 0: prediction 0.3918, SIGReg 6.490 epoch 1: prediction 0.1246, SIGReg 2.615 epoch 2: prediction 0.1061, SIGReg 2.131 epoch 3: prediction 0.1008, SIGReg 1.885 epoch 4: prediction 0.0941, SIGReg 1.730 epoch 5: prediction 0.0872, SIGReg 1.631 epoch 6: prediction 0.0784, SIGReg 1.558 epoch 7: prediction 0.0708, SIGReg 1.522 epoch 8: prediction 0.0653, SIGReg 1.497 epoch 9: prediction 0.0626, SIGReg 1.491 .. GENERATED FROM PYTHON SOURCE LINES 324-332 Check the batch norm -------------------- The projection head has a batch norm layer. In training mode, it uses the statistics of the batch. In eval mode, it uses running statistics. We plan in eval mode, so the two modes must give almost the same embeddings. We compare them on a batch of test frames. A copy of the head runs in training mode, so that the running statistics do not change. .. GENERATED FROM PYTHON SOURCE LINES 332-342 .. code-block:: Python model.eval() with torch.no_grad(): frames = torch.from_numpy(world.render(test_positions[:, 0])).to(device) features = model.backbone(frames) eval_embeddings = model.projection_head(features) train_embeddings = copy.deepcopy(model.projection_head).train()(features) gap = (train_embeddings - eval_embeddings).square().mean() gap = gap / eval_embeddings.var(dim=0).mean() print(f"difference between train and eval mode: {gap:.3f} of the variance") .. rst-class:: sphx-glr-script-out .. code-block:: none difference between train and eval mode: 0.052 of the variance .. GENERATED FROM PYTHON SOURCE LINES 343-361 Does the model use the actions? ------------------------------- A low prediction loss does not prove that the model uses the actions. In many environments, the motion in the past frames predicts most of the next frame. To make sure that the model uses the actions, we do rollouts on test episodes. A rollout starts from the embedding of one frame. The predictor predicts the next embedding and then uses this prediction as its next input. We compare three rollouts of 10 steps: - with the true actions of the episode - with the actions of a different episode - without the model: the embedding of the first frame, repeated The plot shows the mean squared error to the real embeddings, divided by the variance of the embeddings. A model that uses the actions has a much lower error with the true actions. .. GENERATED FROM PYTHON SOURCE LINES 361-405 .. code-block:: Python @torch.no_grad() def encode_positions(positions, chunk_size=512): # Draw and encode frames for positions (..., 2). Returns (..., D) on the CPU. flat = positions.reshape(-1, 2) embeddings = [] for i in range(0, len(flat), chunk_size): frames = torch.from_numpy(world.render(flat[i : i + chunk_size])).to(device) embeddings.append(model.encode(frames[None])[0].cpu()) return torch.cat(embeddings).reshape(*positions.shape[:-1], -1) @torch.no_grad() def rollout(context, actions): # Predict one embedding per action. context: (B, 1, D), actions: (B, H, 2). action_emb = model.action_encoder(actions) return model.predictor.rollout( context, action_emb=action_emb, steps=actions.shape[1] ) horizon = 10 test_embeddings = encode_positions(test_positions).to(device) test_actions_t = torch.from_numpy(test_actions).to(device) context = test_embeddings[:, :1] target = test_embeddings[:, 1 : horizon + 1] variance = target.var(dim=(0, 1)).mean() rollouts = { "true actions": rollout(context=context, actions=test_actions_t[:, :horizon]), "actions of a different episode": rollout( context=context, actions=test_actions_t[:, :horizon].roll(1, dims=0) ), "first embedding, repeated": context.expand_as(target), } plt.figure(figsize=(6, 4)) for label, predicted in rollouts.items(): error = (predicted - target).square().mean(dim=(0, 2)) / variance plt.plot(range(1, horizon + 1), error.cpu(), marker="o", label=label) plt.xlabel("rollout step") plt.ylabel("normalized error") plt.legend() plt.show() .. image-sg:: /tutorials/package/images/sphx_glr_tutorial_lewm_003.png :alt: tutorial lewm :srcset: /tutorials/package/images/sphx_glr_tutorial_lewm_003.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 406-420 Read the predictions with a probe --------------------------------- An embedding is a vector of 192 numbers, and we do not know what each number means. A probe is a small model that reads one property from the embedding, here the position of the agent. We know the position for every frame because we made the frames. Thus, we can train the probe with this position as the target. The probe is not part of the world model. We train it after the world model, and the encoder does not change. The LeWM paper uses linear probes and MLP probes. Here, we use an MLP with one hidden layer, because it reads the position more exactly. If the probe gives the position of new frames with a small error, the embedding contains the position. .. GENERATED FROM PYTHON SOURCE LINES 420-441 .. code-block:: Python probe_positions = train_positions[:200].reshape(-1, 2) probe_embeddings = encode_positions(probe_positions) probe_targets = torch.from_numpy(probe_positions) / world.size probe = nn.Sequential(nn.Linear(embed_dim, 256), nn.ReLU(), nn.Linear(256, 2)) probe_optimizer = torch.optim.Adam(probe.parameters(), lr=1e-3) for step in range(2000): loss = (probe(probe_embeddings) - probe_targets).square().mean() probe_optimizer.zero_grad() loss.backward() probe_optimizer.step() @torch.no_grad() def read_position(embeddings): # Apply the probe to embeddings (..., D). Returns positions (..., 2) in pixels. return probe(embeddings.cpu()) * world.size error = read_position(test_embeddings) - torch.from_numpy(test_positions) print(f"probe error on test frames: {error.norm(dim=-1).mean():.1f} pixels") .. rst-class:: sphx-glr-script-out .. code-block:: none probe error on test frames: 0.6 pixels .. GENERATED FROM PYTHON SOURCE LINES 442-445 Now we apply the same probe to the embeddings of a rollout. These embeddings come from the predictor, not from a frame. Thus, the probe shows where the model expects the agent to be. .. GENERATED FROM PYTHON SOURCE LINES 445-449 .. code-block:: Python imagined = read_position(rollouts["true actions"]) error = imagined - torch.from_numpy(test_positions[:, 1 : horizon + 1]) print(f"probe error on predicted embeddings: {error.norm(dim=-1).mean():.1f} pixels") .. rst-class:: sphx-glr-script-out .. code-block:: none probe error on predicted embeddings: 2.5 pixels .. GENERATED FROM PYTHON SOURCE LINES 450-454 The wall is a good test for the model. In the plot, the agent moves right for 10 steps from two start positions. On the left, the wall blocks the agent. On the right, the agent goes through the door. The dashed line is the rollout, read with the probe. .. GENERATED FROM PYTHON SOURCE LINES 454-484 .. code-block:: Python def push_right(starts, num_steps=10): actions = np.tile(np.float32([1.0, 0.0]), (len(starts), num_steps, 1)) real = [starts] for t in range(num_steps): real.append( np.stack( [world.step(pos=p, action=a) for p, a in zip(real[-1], actions[:, t])] ) ) context = encode_positions(starts[:, None]).to(device) predicted = rollout(context=context, actions=torch.from_numpy(actions).to(device)) real = np.stack(real, axis=1) imagined = read_position(predicted).numpy() imagined = np.concatenate([starts[:, None], imagined], axis=1) return real, imagined real, imagined = push_right(np.array([[10.0, 8.0], [10.0, 28.0]])) fig, axes = plt.subplots(1, 2, figsize=(8, 4)) for ax, real_path, imagined_path in zip(axes, real, imagined): ax.imshow(world.render(real_path[0]).transpose(1, 2, 0), extent=(0, 64, 64, 0)) center = world.agent / 2 ax.plot(*(real_path + center).T, "w-o", markersize=3, label="real") ax.plot(*(imagined_path + center).T, "c--o", markersize=3, label="rollout") ax.axis("off") axes[0].legend(loc="lower left") plt.show() .. image-sg:: /tutorials/package/images/sphx_glr_tutorial_lewm_004.png :alt: tutorial lewm :srcset: /tutorials/package/images/sphx_glr_tutorial_lewm_004.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 485-488 The same test with 200 random start positions in the left half of the room shows how often the rollout agrees with the environment on one question: does the agent get through the wall? .. GENERATED FROM PYTHON SOURCE LINES 488-496 .. code-block:: Python starts = np.stack([rng.uniform(0, 22, 200), rng.uniform(0, 56, 200)], axis=1) real, imagined = push_right(starts) through_real = real[:, -1, 0] + world.agent / 2 > wall_center through_imagined = imagined[:, -1, 0] + world.agent / 2 > wall_center print( f"agreement with the environment: {np.mean(through_real == through_imagined):.0%}" ) .. rst-class:: sphx-glr-script-out .. code-block:: none agreement with the environment: 96% .. GENERATED FROM PYTHON SOURCE LINES 497-503 Look at the predictions with a decoder -------------------------------------- A decoder turns an embedding back into a frame. LeWM does not need a decoder, and the decoder does not change the world model. We train it after the world model, on frozen embeddings, only to look at the predictions. .. GENERATED FROM PYTHON SOURCE LINES 503-538 .. code-block:: Python class FrameDecoder(nn.Module): def __init__(self, embed_dim): super().__init__() self.project = nn.Linear(embed_dim, 512 * 4 * 4) layers = [] for in_channels, out_channels in [(512, 256), (256, 128), (128, 64), (64, 64)]: layers += [ nn.Upsample(scale_factor=2), nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1), nn.GroupNorm(32, out_channels), nn.SiLU(), ] layers += [nn.Conv2d(64, 3, kernel_size=3, padding=1), nn.Sigmoid()] self.layers = nn.Sequential(*layers) def forward(self, embeddings): return self.layers(self.project(embeddings).view(-1, 512, 4, 4)) decoder = FrameDecoder(embed_dim).to(device) decoder_optimizer = torch.optim.AdamW(decoder.parameters(), lr=1e-3) decoder_positions = train_positions.reshape(-1, 2) decoder_embeddings = encode_positions(decoder_positions) for step in range(1000): index = rng.integers(0, len(decoder_positions), 256) frames = torch.from_numpy(world.render(decoder_positions[index])).to(device) reconstructed = decoder(decoder_embeddings[index].to(device)) loss = nn.functional.mse_loss(reconstructed, frames) decoder_optimizer.zero_grad() loss.backward() decoder_optimizer.step() print(f"decoder error: {loss.item():.5f}") .. rst-class:: sphx-glr-script-out .. code-block:: none decoder error: 0.00005 .. GENERATED FROM PYTHON SOURCE LINES 539-542 The first row shows a test episode. The second row decodes a rollout with the true actions. The third row decodes a rollout with the opposite actions. All rows start from the same frame. .. GENERATED FROM PYTHON SOURCE LINES 542-558 .. code-block:: Python opposite = rollout(context=context[:1], actions=-test_actions_t[:1, :horizon]) with torch.no_grad(): first = torch.from_numpy(world.render(test_positions[0, :1])).to(device) rows = { "real": torch.from_numpy(world.render(test_positions[0, : horizon + 1])), "true actions": torch.cat([first, decoder(rollouts["true actions"][0])]), "opposite actions": torch.cat([first, decoder(opposite[0])]), } fig, axes = plt.subplots(3, horizon + 1, figsize=(horizon + 1, 3.6)) for row, (label, frames) in enumerate(rows.items()): for t in range(horizon + 1): axes[row, t].imshow(frames[t].cpu().numpy().transpose(1, 2, 0).clip(0, 1)) axes[row, t].axis("off") axes[row, 0].set_title(label, fontsize=8, loc="left") plt.show() .. image-sg:: /tutorials/package/images/sphx_glr_tutorial_lewm_005.png :alt: real, true actions, opposite actions :srcset: /tutorials/package/images/sphx_glr_tutorial_lewm_005.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 559-587 Plan with the model ------------------- Planning uses the model to select actions. The planner gets the current frame and a goal frame. It tries many action sequences in the model and keeps the sequence whose predicted embeddings are nearest to the embedding of the goal frame. The model does not change during planning. We use the cross-entropy method (CEM). A plan is a sequence of 8 actions. CEM keeps a Gaussian distribution over plans: ``mean`` is the best plan so far, and ``std`` is the range around it that CEM still tries. Each iteration does these steps: 1. Sample 300 plans from the distribution. 2. Do a rollout for each plan from the current embedding. 3. Calculate the cost of each plan: the mean squared distance between the predicted embeddings and the goal embedding, averaged over the 8 steps. 4. Set ``mean`` and ``std`` to the mean and standard deviation of the 30 plans with the lowest cost. After 10 iterations, ``mean`` is the plan. ``plan`` makes plans for many episodes at the same time, so every tensor has a batch dimension ``B``. .. note:: The LeWM paper uses the distance at the last step of the plan only. The average over all steps makes the agent go to the goal early and stay there. On this environment, it gives a much higher success rate. .. GENERATED FROM PYTHON SOURCE LINES 587-615 .. code-block:: Python plan_horizon = 8 @torch.no_grad() def plan( current, goal, horizon=plan_horizon, num_samples=300, num_iters=10, num_elites=30 ): # current, goal: embeddings (B, D). Returns plans (B, horizon, 2). batch, dim = current.shape mean = torch.zeros(batch, horizon, 2, device=device) std = torch.ones(batch, horizon, 2, device=device) context = current.repeat_interleave(num_samples, dim=0).unsqueeze(1) for _ in range(num_iters): noise = torch.randn(batch, num_samples, horizon, 2, device=device) candidates = (mean[:, None] + std[:, None] * noise).clamp(-1, 1) predicted = rollout(context=context, actions=candidates.flatten(0, 1)) predicted = predicted.view(batch, num_samples, horizon, dim) costs = (predicted - goal[:, None, None]).square().mean(dim=(2, 3)) # Refit the distribution to the plans with the lowest cost. elite_idx = torch.topk(costs, num_elites, dim=1, largest=False).indices elites = torch.take_along_dim(candidates, elite_idx[:, :, None, None], dim=1) mean = elites.mean(dim=1) std = elites.std(dim=1) return mean .. GENERATED FROM PYTHON SOURCE LINES 616-619 The agent does the first 4 actions of the plan in the environment. Then the planner makes a new plan from the new frame. This loop is model-predictive control. It corrects the errors of the model before they add up. .. GENERATED FROM PYTHON SOURCE LINES 619-644 .. code-block:: Python def run_episodes(starts, goals, policy, num_steps, replan_every=4): positions = starts.copy() path = [positions] for _ in range(num_steps // replan_every): actions = policy(positions=positions, goals=goals) for t in range(replan_every): positions = np.stack( [world.step(pos=p, action=a) for p, a in zip(positions, actions[:, t])] ) path.append(positions) return np.stack(path, axis=1) def cem_policy(positions, goals): current = encode_positions(positions).to(device) goal = encode_positions(goals).to(device) return plan(current=current, goal=goal).cpu().numpy() def random_policy(positions, goals): return rng.uniform(-1, 1, size=(len(positions), plan_horizon, 2)) .. GENERATED FROM PYTHON SOURCE LINES 645-649 We evaluate like the LeWM paper. The goal is the frame ``k`` steps later in a new episode, and the agent has ``2 * k`` steps to reach it. An episode is a success if the agent ends less than 4 pixels from the goal. A random policy is the baseline. .. GENERATED FROM PYTHON SOURCE LINES 649-664 .. code-block:: Python results = {} for k in [4, 8, 16]: goal_positions, _ = collect(num_episodes=400, num_steps=k, rng=rng) distance = np.linalg.norm(goal_positions[:, k] - goal_positions[:, 0], axis=1) keep = np.flatnonzero(distance >= 8)[:30] starts, goals = goal_positions[keep, 0], goal_positions[keep, k] for name, policy in [("random actions", random_policy), ("CEM", cem_policy)]: paths = run_episodes(starts=starts, goals=goals, policy=policy, num_steps=2 * k) miss = np.linalg.norm(paths[:, -1] - goals, axis=1) results[k, name] = (paths, goals, miss) print( f"goal {k} steps ahead, {name}: success {np.mean(miss < 4):.0%}, " f"median distance to the goal {np.median(miss):.1f} pixels" ) .. rst-class:: sphx-glr-script-out .. code-block:: none goal 4 steps ahead, random actions: success 3%, median distance to the goal 18.0 pixels goal 4 steps ahead, CEM: success 97%, median distance to the goal 0.9 pixels goal 8 steps ahead, random actions: success 10%, median distance to the goal 20.9 pixels goal 8 steps ahead, CEM: success 70%, median distance to the goal 0.6 pixels goal 16 steps ahead, random actions: success 7%, median distance to the goal 25.3 pixels goal 16 steps ahead, CEM: success 53%, median distance to the goal 1.2 pixels .. GENERATED FROM PYTHON SOURCE LINES 665-667 The plot shows four CEM episodes with the goal 8 steps ahead: two that reach the goal and two that miss it. The outline is the goal position of the agent. .. GENERATED FROM PYTHON SOURCE LINES 667-679 .. code-block:: Python paths, goals, miss = results[8, "CEM"] episodes = [*np.flatnonzero(miss < 4)[:2], *np.flatnonzero(miss >= 4)[:2]] fig, axes = plt.subplots(1, len(episodes), figsize=(3 * len(episodes), 3.4)) for ax, episode in zip(axes, episodes): path, goal = paths[episode], goals[episode] ax.imshow(world.render(path[0]).transpose(1, 2, 0), extent=(0, 64, 64, 0)) ax.add_patch(Rectangle(goal, world.agent, world.agent, fill=False, color="w")) ax.plot(*(path + world.agent / 2).T, "c-o", markersize=3) ax.set_title("reached" if miss[episode] < 4 else "missed") ax.axis("off") plt.show() .. image-sg:: /tutorials/package/images/sphx_glr_tutorial_lewm_006.png :alt: reached, reached, missed, missed :srcset: /tutorials/package/images/sphx_glr_tutorial_lewm_006.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 680-687 Limits of the planner --------------------- The planner works for near goals. For goals 16 steps ahead, the success rate is lower. The plot shows the reason. For pairs of random frames, it compares the distance between the two embeddings with the distance between the two agent positions. .. GENERATED FROM PYTHON SOURCE LINES 687-704 .. code-block:: Python positions = np.stack([world.random_position(rng) for _ in range(2000)]) embeddings = encode_positions(positions) i, j = rng.integers(0, len(positions), size=(2, 20000)) pixel_distance = np.linalg.norm(positions[i] - positions[j], axis=1) embedding_distance = (embeddings[i] - embeddings[j]).square().mean(dim=1).numpy() bins = np.arange(0, 60, 4) centers = [low + 2 for low in bins] means = [ embedding_distance[(pixel_distance >= low) & (pixel_distance < low + 4)].mean() for low in bins ] plt.figure(figsize=(6, 4)) plt.plot(centers, means, marker="o") plt.xlabel("distance between the agent positions (pixels)") plt.ylabel("mean squared distance between the embeddings") plt.show() .. image-sg:: /tutorials/package/images/sphx_glr_tutorial_lewm_007.png :alt: tutorial lewm :srcset: /tutorials/package/images/sphx_glr_tutorial_lewm_007.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 705-724 The embedding distance increases quickly for the first 10 pixels and then stays almost constant. For a goal that is 30 pixels away, most plans end in the flat part of the curve and have almost the same cost. CEM then finds the goal only if some of the sampled plans end near it. The LeWM paper reports a similar result on its Two-Room environment. The authors explain it with the low diversity and the low intrinsic dimension of the data. In this environment, the position of the agent is the only property that changes. Other methods change the cost or the embeddings so that a planner can reach far goals: - `TEMPO `_ trains a small model on frozen LeWM embeddings. In its space, the distance between two states is the number of steps between them. The planner uses this distance as the cost. - `PLDM `_ trains with a loss term that keeps consecutive embeddings close, so that the embedding distance changes smoothly over time. - `Search on the Replay Buffer `_ builds a graph of recorded states and plans a path of waypoints to the goal. .. GENERATED FROM PYTHON SOURCE LINES 726-734 Next Steps ---------- - :ref:`lewm` describes the LeWM modules and the shapes of their inputs. - :ref:`lejepa` uses the same SIGReg term for self-supervised learning on images. - The `official LeWM code `_ trains LeWM on larger environments, for example Push-T and OGBench Cube. .. rst-class:: sphx-glr-timing **Total running time of the script:** (3 minutes 48.971 seconds) .. _sphx_glr_download_tutorials_package_tutorial_lewm.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: tutorial_lewm.ipynb ` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: tutorial_lewm.py ` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: tutorial_lewm.zip ` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_