DINO-WM: world models on frozen features, then used to plan¶

DINO-WM (code) leaves the visual encoder frozen (DINOv2 patch features) and learns only the dynamics: a ViT predictor with frame-level causal attention, plus small action and proprioception encoders. Planning then runs in the frozen feature space.

Two worlds, both simulated in MuJoCo (pip install torch-world[action]):

1. a Universal Robots UR10e arm with a Robotiq 2F-85 gripper and a cube, after OGBench cube-single;
2. Push-T, pushing a T-shaped block into a goal pose, one of the paper's benchmarks.

In [1]:

import torch
from torch.utils.data import DataLoader

from torch_world import evaluation, metrics, plots
from torch_world.datasets import TrajectoryDataset
from torch_world.envs import PushT, RobotCube
from torch_world.models import DINOWM
from torch_world.nn import PixelPatches
from torch_world.planning import CEM

torch.manual_seed(0)
device = "cuda" if torch.cuda.is_available() else "cpu"

import torch from torch.utils.data import DataLoader from torch_world import evaluation, metrics, plots from torch_world.datasets import TrajectoryDataset from torch_world.envs import PushT, RobotCube from torch_world.models import DINOWM from torch_world.nn import PixelPatches from torch_world.planning import CEM torch.manual_seed(0) device = "cuda" if torch.cuda.is_available() else "cpu"

Frozen encoder¶

With internet access, use DINOv2: encoder = PatchTokens.from_hub("dinov2_vits14"). To stay self-contained and fast, this notebook uses raw 16×16 pixel patches as the frozen features (PixelPatches), so predicted latents can be viewed as images.

In [2]:

encoder = PixelPatches(16)
HISTORY = 3


def train(model, loader, epochs):
    opt = torch.optim.AdamW(model.parameters(), lr=5e-4)
    log = {"visual": [], "proprio": []}
    for epoch in range(epochs):
        for batch in loader:
            out = model(batch["pixels"], batch["actions"], batch["proprio"])
            opt.zero_grad()
            out["loss"].backward()
            opt.step()
            log["visual"].append(out["visual_loss"].item())
            log["proprio"].append(out["proprio_loss"].item())
    return model.eval(), log

encoder = PixelPatches(16) HISTORY = 3 def train(model, loader, epochs): opt = torch.optim.AdamW(model.parameters(), lr=5e-4) log = {"visual": [], "proprio": []} for epoch in range(epochs): for batch in loader: out = model(batch["pixels"], batch["actions"], batch["proprio"]) opt.zero_grad() out["loss"].backward() opt.step() log["visual"].append(out["visual_loss"].item()) log["proprio"].append(out["proprio_loss"].item()) return model.eval(), log

Part 1 · UR10e arm¶

Play data¶

A noisy scripted policy picks the cube up and puts it down at random spots, like OGBench's play datasets. Actions move the gripper by up to 10 cm and open or close it, five times a second; proprioception is the gripper position and opening.

In [3]:

ROBOT = dict(robot="ur10e", size=96, control_dt=0.2, max_step=0.1)
data = {k: v.to(device) for k, v in RobotCube(batch_size=1024, **ROBOT).collect(24, generator=torch.Generator().manual_seed(0)).items()}
data["proprio"] = data["states"][..., :4]
plots.images(data["pixels"][2, :16].cpu(), [f"t={t}" for t in range(16)], size=1.3)

ROBOT = dict(robot="ur10e", size=96, control_dt=0.2, max_step=0.1) data = {k: v.to(device) for k, v in RobotCube(batch_size=1024, **ROBOT).collect(24, generator=torch.Generator().manual_seed(0)).items()} data["proprio"] = data["states"][..., :4] plots.images(data["pixels"][2, :16].cpu(), [f"t={t}" for t in range(16)], size=1.3)

Out[3]:

In [4]:

arm = DINOWM(encoder, num_patches=36, embed_dim=encoder.dim, action_dim=4, proprio_dim=4, history=HISTORY, depth=6, num_heads=8, head_dim=64, mlp_dim=1024).to(device)
print(f"{sum(p.numel() for p in arm.parameters() if p.requires_grad) / 1e6:.2f}M trainable parameters")
arm, log = train(arm, DataLoader(TrajectoryDataset(data, HISTORY + 1), batch_size=256, shuffle=True, drop_last=True), 30)
plots.curves(log, smooth=100)

arm = DINOWM(encoder, num_patches=36, embed_dim=encoder.dim, action_dim=4, proprio_dim=4, history=HISTORY, depth=6, num_heads=8, head_dim=64, mlp_dim=1024).to(device) print(f"{sum(p.numel() for p in arm.parameters() if p.requires_grad) / 1e6:.2f}M trainable parameters") arm, log = train(arm, DataLoader(TrajectoryDataset(data, HISTORY + 1), batch_size=256, shuffle=True, drop_last=True), 30) plots.curves(log, smooth=100)

19.49M trainable parameters

Out[4]:

Imagination¶

From one real frame, the model unrolls the expert's pick-and-place actions in feature space without seeing any more frames. Pixel-patch features decode exactly to images, so imagined frames are scored against the real ones with the image metrics DINO-WM reports for its decoded predictions: SSIM (1 is identical) and PSNR, over 64 held-out episodes.

In [5]:

env = RobotCube(batch_size=64, **ROBOT)
first = env.reset(generator=torch.Generator().manual_seed(5))
target = env.sample_targets(generator=torch.Generator().manual_seed(6))
s0, real, actions = env.state.clone(), [first], []
for t in range(10):
    actions.append(env.expert(target))
    real.append(env.step(actions[-1]))
real = torch.stack(real, 1)
with torch.no_grad():
    z = arm.rollout(arm.encode(first[:, None].to(device), s0[:, None, :4].to(device)), torch.stack(actions, 1).to(device))
imagined = encoder.unpatchify(z.flatten(0, 1), (6, 6)).unflatten(0, z.shape[:2]).clamp(0, 1).cpu()
ssim, psnr = metrics.ssim(imagined, real).mean(0), metrics.psnr(imagined, real).mean(0)
print("step    " + "".join(f"{t:>6d}" for t in range(0, 11, 2)))
print("SSIM    " + "".join(f"{v:6.3f}" for v in ssim[::2].tolist()))
print("PSNR dB " + "".join(f"{v:6.1f}" for v in psnr[::2].tolist()))
keep = list(range(0, 11, 2))
plots.images(torch.cat([real[0, keep], imagined[0, keep]]), ncols=len(keep), size=1.5, rows=["real", "imagined"], titles=[f"t={t}" for t in keep] * 2)

env = RobotCube(batch_size=64, **ROBOT) first = env.reset(generator=torch.Generator().manual_seed(5)) target = env.sample_targets(generator=torch.Generator().manual_seed(6)) s0, real, actions = env.state.clone(), [first], [] for t in range(10): actions.append(env.expert(target)) real.append(env.step(actions[-1])) real = torch.stack(real, 1) with torch.no_grad(): z = arm.rollout(arm.encode(first[:, None].to(device), s0[:, None, :4].to(device)), torch.stack(actions, 1).to(device)) imagined = encoder.unpatchify(z.flatten(0, 1), (6, 6)).unflatten(0, z.shape[:2]).clamp(0, 1).cpu() ssim, psnr = metrics.ssim(imagined, real).mean(0), metrics.psnr(imagined, real).mean(0) print("step " + "".join(f"{t:>6d}" for t in range(0, 11, 2))) print("SSIM " + "".join(f"{v:6.3f}" for v in ssim[::2].tolist())) print("PSNR dB " + "".join(f"{v:6.1f}" for v in psnr[::2].tolist())) keep = list(range(0, 11, 2)) plots.images(torch.cat([real[0, keep], imagined[0, keep]]), ncols=len(keep), size=1.5, rows=["real", "imagined"], titles=[f"t={t}" for t in keep] * 2)

step         0     2     4     6     8    10
SSIM     1.000 0.711 0.698 0.674 0.636 0.625
PSNR dB  120.0  18.5  17.3  16.4  15.5  15.0

Out[5]:

Plan toward goal images¶

Each task gives a start and the image of a goal state. CEM searches 5-step action sequences in feature space and MPC executes the first action of each plan. The "path" objective scores every imagined step, not only the last, so the planner moves now instead of postponing its moves to the end of the horizon.

Two tasks:

• reach: move the open gripper to the pose in the goal image; success within 5 cm.
• pick and place: carry the cube to the spot in the goal image (OGBench cube-single); success when the cube rests within 5 cm of it.

The scripted expert and random actions bracket the planner.

Each policy gets the standard measures of goal reaching: the final success rate with its 95% Wilson interval (16 episodes leave wide intervals), the share of episodes that reached the goal at any step and the mean steps to get there, the median final distance (meters), the mean progress (fraction of the start distance closed), SPL (success weighted by how direct the path was), action smoothness (mean change between consecutive actions; lower is smoother) and the planning time per step.

In [6]:

def plan_with(model, env, goal_state, proprio, action_dim):
    with torch.no_grad():
        goal = model.encode(env.render(goal_state)[:, None].to(device), goal_state[:, None, proprio].to(device))[:, 0]
    planner = CEM(horizon=5, action_dim=action_dim, num_samples=256, num_elites=32, iterations=10, bounds=(-1, 1), objective="path")
    return lambda obs: planner(model, model.encode(obs[:, None].to(device), env.state[:, None, proprio].to(device)), goal)[:, 0].cpu()


arm_runs = {}
for task, steps in [("reach", 12), ("pick and place", 30)]:
    env = RobotCube(batch_size=16, task=task.split()[0], **ROBOT)
    start, goal_state = env.tasks(generator=torch.Generator().manual_seed(1))
    policies = {"DINO-WM + CEM": plan_with(arm, env, goal_state, slice(0, 4), 4), "scripted expert": lambda obs, env=env, goal_state=goal_state: env.scripted(goal_state), "random actions": evaluation.random_policy(env)}
    arm_runs[task] = evaluation.goal_reaching(env, policies, start, goal_state, steps, record=lambda e, goal_state=goal_state: e.render(size=160, goal=goal_state))
    print(f"== {task}\n{arm_runs[task]}")
plots.bars({task: r.intervals() for task, r in arm_runs.items()}, xlabel="success rate", baseline="random actions")

def plan_with(model, env, goal_state, proprio, action_dim): with torch.no_grad(): goal = model.encode(env.render(goal_state)[:, None].to(device), goal_state[:, None, proprio].to(device))[:, 0] planner = CEM(horizon=5, action_dim=action_dim, num_samples=256, num_elites=32, iterations=10, bounds=(-1, 1), objective="path") return lambda obs: planner(model, model.encode(obs[:, None].to(device), env.state[:, None, proprio].to(device)), goal)[:, 0].cpu() arm_runs = {} for task, steps in [("reach", 12), ("pick and place", 30)]: env = RobotCube(batch_size=16, task=task.split()[0], **ROBOT) start, goal_state = env.tasks(generator=torch.Generator().manual_seed(1)) policies = {"DINO-WM + CEM": plan_with(arm, env, goal_state, slice(0, 4), 4), "scripted expert": lambda obs, env=env, goal_state=goal_state: env.scripted(goal_state), "random actions": evaluation.random_policy(env)} arm_runs[task] = evaluation.goal_reaching(env, policies, start, goal_state, steps, record=lambda e, goal_state=goal_state: e.render(size=160, goal=goal_state)) print(f"== {task}\n{arm_runs[task]}") plots.bars({task: r.intervals() for task, r in arm_runs.items()}, xlabel="success rate", baseline="random actions")

== reach
                           success  reached   steps  final distance  progress     SPL  smoothness  ms / step
DINO-WM + CEM        19% [7%, 43%]      81%    3.92          0.0709      0.61   0.112       0.567   2.69e+05
scripted expert   100% [81%, 100%]     100%    2.44           0.015     0.927   0.943       0.158      0.186
random actions        0% [0%, 19%]       6%       3           0.313    -0.374       0        1.56       0.17

== pick and place
                           success  reached   steps  final distance  progress     SPL  smoothness  ms / step
DINO-WM + CEM         0% [0%, 19%]       0%     nan           0.189         0       0       0.523   2.21e+05
scripted expert   100% [81%, 100%]     100%    10.1         0.00254     0.984   0.955       0.442       2.68
random actions        0% [0%, 19%]       0%     nan           0.189    0.0255       0        1.55      0.241

Out[6]:

Reaching episodes of the planner. The translucent orange gripper marks the goal (videos only; the model sees the plain scene):

In [7]:

r = arm_runs["reach"]
plots.episodes(r.frames["DINO-WM + CEM"][:8], r.success["DINO-WM + CEM"][:8], [f"{100 * d:.0f} cm → {100 * e:.0f} cm" for d, e in r.distances["DINO-WM + CEM"][:8, [0, -1]].tolist()])

r = arm_runs["reach"] plots.episodes(r.frames["DINO-WM + CEM"][:8], r.success["DINO-WM + CEM"][:8], [f"{100 * d:.0f} cm → {100 * e:.0f} cm" for d, e in r.distances["DINO-WM + CEM"][:8, [0, -1]].tolist()])

Out[7]:

Pick and place is where greedy latent planning breaks down. The planner reaches the arm pose shown in the goal image, but carrying the cube needs a detour (descend, grasp, lift, carry) that a short horizon never finds attractive, and the cube covers only a few patches. The paper reaches such tasks with pre-trained DINOv2 features, more data and longer horizons; here the expert and planner videos show the gap:

In [8]:

r = arm_runs["pick and place"]
plots.episodes(torch.cat([r.frames["scripted expert"][:2], r.frames["DINO-WM + CEM"][:6]]), torch.cat([r.success["scripted expert"][:2], r.success["DINO-WM + CEM"][:6]]), ["expert"] * 2 + ["DINO-WM + CEM"] * 6)

r = arm_runs["pick and place"] plots.episodes(torch.cat([r.frames["scripted expert"][:2], r.frames["DINO-WM + CEM"][:6]]), torch.cat([r.success["scripted expert"][:2], r.success["DINO-WM + CEM"][:6]]), ["expert"] * 2 + ["DINO-WM + CEM"] * 6)

Out[8]:

Part 2 · Push-T¶

The agent is a round pusher that moves on a table; actions shift its target by up to a few centimeters, and proprioception is its position. Play data comes from a noisy policy that walks up to the T and pushes it from random points of its edge.

In [9]:

pusht = {k: v.to(device) for k, v in PushT(batch_size=2048, size=96).collect(32, generator=torch.Generator().manual_seed(0)).items()}
pusht["proprio"] = pusht["states"][..., :2]
plots.images(pusht["pixels"][0, ::2].cpu(), [f"t={t}" for t in range(0, 32, 2)], size=1.3)
pt = DINOWM(encoder, num_patches=36, embed_dim=encoder.dim, action_dim=2, proprio_dim=2, history=HISTORY, depth=6, num_heads=8, head_dim=64, mlp_dim=1024).to(device)
pt, log = train(pt, DataLoader(TrajectoryDataset(pusht, HISTORY + 1), batch_size=256, shuffle=True, drop_last=True), 20)
plots.curves(log, smooth=100)

pusht = {k: v.to(device) for k, v in PushT(batch_size=2048, size=96).collect(32, generator=torch.Generator().manual_seed(0)).items()} pusht["proprio"] = pusht["states"][..., :2] plots.images(pusht["pixels"][0, ::2].cpu(), [f"t={t}" for t in range(0, 32, 2)], size=1.3) pt = DINOWM(encoder, num_patches=36, embed_dim=encoder.dim, action_dim=2, proprio_dim=2, history=HISTORY, depth=6, num_heads=8, head_dim=64, mlp_dim=1024).to(device) pt, log = train(pt, DataLoader(TrajectoryDataset(pusht, HISTORY + 1), batch_size=256, shuffle=True, drop_last=True), 20) plots.curves(log, smooth=100)

Out[9]:

Goals are where the pushing policy left the T after 20 steps. Success means the T lies within 2.5 cm and about 15° of its goal pose. The table adds coverage, the score of Diffusion Policy's Push-T: the best fraction of the goal T's area the block covered during the episode.

In [10]:

env = PushT(batch_size=16, size=96)
start, goal_state = env.tasks(steps=20, generator=torch.Generator().manual_seed(1))
policies = {"DINO-WM + CEM": plan_with(pt, env, goal_state, slice(0, 2), 2), "random actions": evaluation.random_policy(env)}
pt_run = evaluation.goal_reaching(env, policies, start, goal_state, steps=30, record=lambda e: e.render(size=160, goal=goal_state))
print(pt_run)
plots.distances(pt_run.distances, radius=env.radius)

env = PushT(batch_size=16, size=96) start, goal_state = env.tasks(steps=20, generator=torch.Generator().manual_seed(1)) policies = {"DINO-WM + CEM": plan_with(pt, env, goal_state, slice(0, 2), 2), "random actions": evaluation.random_policy(env)} pt_run = evaluation.goal_reaching(env, policies, start, goal_state, steps=30, record=lambda e: e.render(size=160, goal=goal_state)) print(pt_run) plots.distances(pt_run.distances, radius=env.radius)

                          success  reached   steps  final distance  progress  smoothness  ms / step  coverage
DINO-WM + CEM        0% [0%, 19%]       0%     nan          0.0734     0.116       0.499   1.82e+05       21%
random actions       0% [0%, 19%]       0%     nan          0.0778    0.0235        1.06      0.463       14%

Out[10]:

In [11]:

plots.episodes(pt_run.frames["DINO-WM + CEM"][:8], pt_run.success["DINO-WM + CEM"][:8])

plots.episodes(pt_run.frames["DINO-WM + CEM"][:8], pt_run.success["DINO-WM + CEM"][:8])

Out[11]:

Where the T went (the path darkens over time), from above:

In [12]:

paths = pt_run.paths["DINO-WM + CEM"][:8]
plots.trajectories(env.position(paths), env.position(goal_state[:8]), radius=env.radius / env.top_size, scale=env.top_size, background=env.render(goal_state[:8], camera="top", size=192), success=pt_run.success["DINO-WM + CEM"][:8])

paths = pt_run.paths["DINO-WM + CEM"][:8] plots.trajectories(env.position(paths), env.position(goal_state[:8]), radius=env.radius / env.top_size, scale=env.top_size, background=env.render(goal_state[:8], camera="top", size=192), success=pt_run.success["DINO-WM + CEM"][:8])

Out[12]:

Summary¶

In [13]:

runs = {"UR10e reach": arm_runs["reach"], "UR10e pick and place": arm_runs["pick and place"], "Push-T": pt_run}
plots.bars({name: {k: r.intervals()[k] for k in ("DINO-WM + CEM", "random actions")} for name, r in runs.items()}, xlabel="success rate", baseline="random actions")

runs = {"UR10e reach": arm_runs["reach"], "UR10e pick and place": arm_runs["pick and place"], "Push-T": pt_run} plots.bars({name: {k: r.intervals()[k] for k in ("DINO-WM + CEM", "random actions")} for name, r in runs.items()}, xlabel="success rate", baseline="random actions")

Out[13]: