# plan_cem.py """ Minimal WM evaluation: CEM planning with latent-L2-to-goal as the cost. Validation protocol (single task): - Sample an episode from the expert partition uniformly at random. - Use the last observation in that episode as the goal. - Sample the initial state uniformly from the first 20% of the episode. - CEM in action space over horizon H = (L - 1 - start_idx). Terminal cost: ||z_{start + H} - z_goal||_2 - Baselines: expert (GT) action slice, random actions, shuffled actions. Expert demos are near-optimal (slight-noise expert policies), so their cost under this metric is the reference quality bar. - Save per-plan visualization: real-GT frames / WM rollout under GT actions / WM rollout under CEM plan. Typical usage: python plan_cem.py --task og-point-maze \ --n_episodes 4 --n_starts_per_ep 3 \ --cem_samples 256 --cem_iters 4 \ --viz_dir ./logs/plan_cem_viz \ --output ./logs/plan_cem_og-point-maze.csv """ import argparse import glob import json import math import os import sys import time from dataclasses import dataclass from typing import List, Optional, Tuple import numpy as np import torch from PIL import Image sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) from interactive import ( load_tokenizer_from_ckpt, load_dynamics_from_ckpt, make_tau_schedule, pack_bottleneck_to_spatial, unpack_spatial_to_bottleneck, sample_one_timestep_packed, ) from model import temporal_patchify, temporal_unpatchify, symexp torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True # -------------------- data loading (single episode) -------------------- # @dataclass class Episode: ep_id: int frames: torch.Tensor # (L, 3, H, W) float in [0,1] actions: torch.Tensor # (L, 16) float; actions[0] is nan by convention rewards: torch.Tensor # (L,) float; rewards[0] is nan by convention def load_episode( task: str, data_dir: str, frames_dir: str, ep_id: int, *, shard_size: int = 4096, ) -> Episode: td = torch.load(os.path.join(data_dir, f"{task}.pt"), map_location="cpu", weights_only=False) ep = td["episode"].to(torch.int64) idxs = (ep == ep_id).nonzero(as_tuple=False).flatten().tolist() if not idxs: raise ValueError(f"episode {ep_id} not found in {task}") shard_paths = sorted(glob.glob(os.path.join(frames_dir, task, "*shard*.pt"))) if not shard_paths: raise FileNotFoundError(f"no shards under {frames_dir}/{task}") frames_out: List[Optional[torch.Tensor]] = [None] * len(idxs) by_shard: dict[int, list[tuple[int, int]]] = {} for out_i, raw_i in enumerate(idxs): s, off = raw_i // shard_size, raw_i % shard_size by_shard.setdefault(s, []).append((out_i, off)) for s_idx, picks in by_shard.items(): sd = torch.load(shard_paths[s_idx], map_location="cpu", weights_only=False) fr = sd["frames"] if fr.ndim == 4 and fr.shape[-1] == 3 and fr.shape[1] != 3: fr = fr.permute(0, 3, 1, 2).contiguous() for out_i, off in picks: frames_out[out_i] = fr[off] frames = torch.stack(frames_out).to(torch.float32) / 255.0 actions = td["action"][idxs].to(torch.float32) rewards = td["reward"][idxs].to(torch.float32) return Episode(ep_id=int(ep_id), frames=frames, actions=actions, rewards=rewards) def list_episode_ids(task: str, data_dir: str) -> List[int]: td = torch.load(os.path.join(data_dir, f"{task}.pt"), map_location="cpu", weights_only=False) ep = td["episode"].to(torch.int64) return ep.unique().tolist() # -------------------- encoding / decoding -------------------- # @torch.inference_mode() def encode_frames_to_packed( encoder, frames: torch.Tensor, *, patch: int, n_spatial: int, packing_factor: int, use_amp: bool, ) -> torch.Tensor: """frames: (N, 3, H, W) float [0,1] on device -> (N, n_spatial, d_spatial).""" N, C, H, W = frames.shape patches = temporal_patchify(frames.view(N, 1, C, H, W), patch) device = frames.device with torch.autocast(device_type=device.type, enabled=(use_amp and device.type == "cuda")): z_btLd, _ = encoder(patches) z_packed = pack_bottleneck_to_spatial(z_btLd, n_spatial=n_spatial, k=packing_factor) return z_packed[:, 0].to(torch.float32) @torch.inference_mode() def decode_packed_sequence( decoder, z_packed_seq: torch.Tensor, *, H: int, W: int, C: int, patch: int, packing_factor: int, d_bottleneck: int, use_amp: bool, ) -> torch.Tensor: """z_packed_seq: (T, n_spatial, d_spatial) -> (T, C, H, W) float [0,1].""" T = z_packed_seq.shape[0] z_bt = z_packed_seq.unsqueeze(0) # (1, T, n_spatial, d_spatial) z_btLd = unpack_spatial_to_bottleneck(z_bt, k=packing_factor, d_bottleneck=d_bottleneck) device = z_packed_seq.device with torch.autocast(device_type=device.type, enabled=(use_amp and device.type == "cuda")): patches = decoder(z_btLd) frames = temporal_unpatchify(patches, H, W, C, patch) # (1, T, C, H, W) return frames[0].clamp(0, 1).to(torch.float32) # -------------------- batched rollout -------------------- # @torch.inference_mode() def imagine_batch( dyn, *, z_start: torch.Tensor, action_plan: torch.Tensor, horizon: int, sched: dict, k_max: int, act_mask_1d: torch.Tensor, lang_emb: Optional[torch.Tensor], use_amp: bool, return_h: bool = False, use_kv_cache: bool = True, ): """ Inputs: z_start: (B, n_spatial, d_spatial) action_plan: (B, H, 16); action_plan[:,t] produces z_{t+1}. Returns: return_h=False (default): z_seq: (B, 1+H, n_spatial, d_spatial) with z_seq[:,0] = z_start. return_h=True: (z_seq, h_seq) where h_seq is (B, H, n_agent, D) or (B, H, D). h_seq[:,t] is h at rollout position t+1 (the one that produced z_{t+1}), which aligns with the reward head's MTP l=0 convention: at position t+1 with context (z_0..z_{t+1}, a_0..a_t), l=0 predicts r_t = the reward received from applying action_plan[t]. Summing these H predictions gives the total predicted return over the plan. use_kv_cache=True (default): at each rollout step, prefill the transformer's K,V for the past context tokens once and reuse them across the K=sched["K"] denoising iterations — so each denoising iter only runs the single new noisy token through attention instead of the full context+new sequence. Changes per-rollout compute from O(K·H³) to O(H³ + K·H²) with no change to the output. Set to False to disable caching. """ B = z_start.shape[0] device = z_start.device # Prepend the zero action at position 0 (backward convention: action[k] produced z[k]). zero_a = torch.zeros((B, 1, 16), device=device, dtype=action_plan.dtype) full_actions = torch.cat([zero_a, action_plan], dim=1) full_mask = act_mask_1d.view(1, 1, -1).expand(B, 1 + horizon, -1).contiguous() past = z_start.unsqueeze(1).contiguous() z_outs = [z_start] h_outs: List[torch.Tensor] = [] for t in range(horizon): out = sample_one_timestep_packed( dyn, past_packed=past, k_max=k_max, sched=sched, actions=full_actions[:, : past.shape[1] + 1], act_mask=full_mask[:, : past.shape[1] + 1], use_amp=use_amp, return_h=return_h, tau_ctx=0.0, lang_emb=lang_emb.expand(B, -1) if lang_emb is not None else None, use_kv_cache=use_kv_cache, ) if return_h: z_next, h_new, _inst = out h_outs.append(h_new[:, 0]) # (B, n_agent, D) or (B, D) else: z_next, _inst = out z_outs.append(z_next) past = torch.cat([past, z_next.unsqueeze(1)], dim=1) z_seq = torch.stack(z_outs, dim=1).to(torch.float32) if not return_h: return z_seq h_seq = torch.stack(h_outs, dim=1) return z_seq, h_seq LATENT_COST_KINDS = ("terminal",) REWARD_COST_KINDS = ("reward",) def latent_cost( z_seq: torch.Tensor, z_goal: torch.Tensor, *, kind: str = "terminal", ) -> torch.Tensor: """z_seq: (B, 1+H, n_spatial, d_spatial); scores z_seq[:, 1:] vs z_goal.""" z_goal_b = z_goal.view(1, 1, *z_goal.shape).to(z_seq.dtype) diff = (z_seq[:, 1:] - z_goal_b).flatten(2) per_step_l2 = diff.pow(2).sum(dim=-1).sqrt() if kind == "terminal": return per_step_l2[:, -1] raise ValueError(f"unknown latent cost kind {kind}") @torch.inference_mode() def reward_cost( rew_head, h_seq: torch.Tensor, *, use_amp: bool, discount: float = 1.0, ) -> torch.Tensor: """Cost = -sum_t γ^t r_hat_t, where r_hat_t is the predicted reward at rollout step t (t = 0..H-1 within the plan) read from the reward head's MTP l=0 position. Sign convention matches the latent-cost modes (CEM minimizes), so CEM picking the lowest `reward_cost` is equivalent to picking the highest predicted (discounted) return. discount=1.0 (default) gives the undiscounted cumulative reward used by open-loop planning. discount<1.0 applies per-step γ weighting to the predicted return. h_seq: (B, H, n_agent, D) or (B, H, D), one hidden state per rollout step. """ device = h_seq.device with torch.autocast(device_type=device.type, enabled=(use_amp and device.type == "cuda")): logits_btlk, centers_log = rew_head(h_seq) # (B, H, L, K), (K,) logits_l0 = logits_btlk[..., 0, :].float() # (B, H, K); head l=0 -> r at that step probs = logits_l0.softmax(dim=-1) expected_symlog = (probs * centers_log.float().view(1, 1, -1)).sum(dim=-1) # (B, H) r_hat = symexp(expected_symlog) # (B, H), predicted reward per step if float(discount) < 1.0: H = r_hat.shape[-1] gammas = torch.pow( torch.tensor(float(discount), device=r_hat.device, dtype=r_hat.dtype), torch.arange(H, device=r_hat.device, dtype=r_hat.dtype), ) # (H,) return -(r_hat * gammas.view(1, H)).sum(dim=-1) # (B,) return -r_hat.sum(dim=-1) # (B,) # -------------------- CEM -------------------- # def colored_noise( n_samples: int, horizon: int, act_dim: int, beta: float, device: torch.device, generator: Optional[torch.Generator] = None, ) -> torch.Tensor: """Power-law-spectrum Gaussian noise, shape (n_samples, horizon, act_dim). Power spectral density scales as 1 / f^beta along the time axis (iCEM convention, Pinneri et al. 2020). beta=0 reproduces white noise; beta=1 is pink; beta=2 is Brownian (random-walk actions, high temporal correlation). The DC component is zeroed so each sampled plan has zero mean, and the output is rescaled to unit per-dim std across (samples, time) to keep `init_std` interpretable across beta settings. """ w = torch.randn((n_samples, horizon, act_dim), device=device, generator=generator) if beta == 0.0: return w W = torch.fft.rfft(w, dim=1) # (n_samples, horizon//2 + 1, act_dim) freqs = torch.arange(W.shape[1], device=device, dtype=torch.float32) scale = torch.zeros_like(freqs) scale[1:] = freqs[1:].pow(-beta / 2.0) W = W * scale.view(1, -1, 1).to(W.dtype) colored = torch.fft.irfft(W, n=horizon, dim=1).to(torch.float32) std = colored.flatten(0, 1).std(dim=0, keepdim=True).clamp_min(1e-8) # (1, act_dim) return colored / std.view(1, 1, -1) @torch.inference_mode() def cem_plan( dyn, *, z_start: torch.Tensor, z_goal: torch.Tensor, horizon: int, act_dim: int, act_mask_1d: torch.Tensor, sched: dict, k_max: int, lang_emb: Optional[torch.Tensor], cost_kind: str, n_samples: int, n_elites: int, n_iters: int, init_std: float, min_std: float, rollout_batch: int, noise_beta: float, device: torch.device, use_amp: bool, generator: Optional[torch.Generator] = None, init_mu: Optional[torch.Tensor] = None, n_rollouts_elite: int = 1, rew_head=None, discount: float = 1.0, ) -> dict: """Diagonal-Gaussian CEM over (horizon, act_dim) plans, clipped to [-1,1]. `noise_beta` controls the time-axis noise color (0 = white, 2 = Brownian). `init_mu` optionally seeds the mean (e.g. from a BC warm-start rollout). `n_rollouts_elite` scores each sample by averaging over k denoising rollouts before elite selection (guards against single-rollout overfit).""" if init_mu is not None: assert init_mu.shape == (horizon, act_dim), \ f"init_mu shape {tuple(init_mu.shape)} != (horizon={horizon}, act_dim={act_dim})" mu = init_mu.detach().to(device=device, dtype=torch.float32).clamp(-1.0, 1.0).clone() else: mu = torch.zeros((horizon, act_dim), device=device, dtype=torch.float32) std = torch.full((horizon, act_dim), float(init_std), device=device, dtype=torch.float32) best_cost = float("inf") best_plan = mu.clone() history = [] z_start_b = z_start.unsqueeze(0) k_sel = max(1, int(n_rollouts_elite)) for it in range(n_iters): noise = colored_noise( n_samples, horizon, act_dim, float(noise_beta), device=device, generator=generator, ) samples = (mu + std * noise).clamp(-1.0, 1.0) full = torch.zeros((n_samples, horizon, 16), device=device, dtype=torch.float32) full[:, :, :act_dim] = samples full = full * act_mask_1d.view(1, 1, 16) # Replicate each plan k_sel times; each replica gets an independent # denoising seed inside the sampler, so averaging their costs guards # against single-rollout luck. total = n_samples * k_sel full_rep = full.unsqueeze(1).expand(-1, k_sel, -1, -1).reshape(total, horizon, 16).contiguous() use_reward_cost = cost_kind in REWARD_COST_KINDS if use_reward_cost and rew_head is None: raise ValueError(f"cost_kind={cost_kind!r} requires a reward head; got rew_head=None") costs_flat = torch.empty(total, device=device, dtype=torch.float32) for s in range(0, total, rollout_batch): e = min(s + rollout_batch, total) z0_b = z_start_b.expand(e - s, -1, -1).contiguous() if use_reward_cost: _z_seq, h_seq = imagine_batch( dyn, z_start=z0_b, action_plan=full_rep[s:e], horizon=horizon, sched=sched, k_max=k_max, act_mask_1d=act_mask_1d, lang_emb=lang_emb, use_amp=use_amp, return_h=True, ) costs_flat[s:e] = reward_cost(rew_head, h_seq, use_amp=use_amp, discount=discount) else: z_seq = imagine_batch( dyn, z_start=z0_b, action_plan=full_rep[s:e], horizon=horizon, sched=sched, k_max=k_max, act_mask_1d=act_mask_1d, lang_emb=lang_emb, use_amp=use_amp, ) costs_flat[s:e] = latent_cost(z_seq, z_goal, kind=cost_kind) costs = costs_flat.view(n_samples, k_sel).mean(dim=1) elite_idx = torch.topk(costs, k=n_elites, largest=False).indices elites = samples[elite_idx] new_mu = elites.mean(dim=0) new_std = elites.std(dim=0).clamp_min(float(min_std)) iter_best_cost = float(costs[elite_idx[0]].item()) if iter_best_cost < best_cost: best_cost = iter_best_cost best_plan = samples[elite_idx[0]].clone() history.append({ "iter": it, "mean_cost": float(costs.mean().item()), "elite_mean_cost": float(costs[elite_idx].mean().item()), "best_cost": float(iter_best_cost), "running_best_cost": float(best_cost), "n_rollouts_elite": int(k_sel), }) mu, std = new_mu, new_std return {"best_plan": best_plan, "best_cost": float(best_cost), "history": history} # -------------------- reference-plan scoring (with returned latents) -------------------- # @torch.inference_mode() def score_plan( dyn, *, z_start: torch.Tensor, z_goal: torch.Tensor, action_plan: torch.Tensor, # (H, act_dim) in [-1,1] act_dim: int, act_mask_1d: torch.Tensor, horizon: int, sched: dict, k_max: int, lang_emb: Optional[torch.Tensor], cost_kind: str, use_amp: bool, device: torch.device, n_rollouts: int = 1, return_first_latents: bool = False, rew_head=None, discount: float = 1.0, ) -> Tuple[float, Optional[torch.Tensor]]: """Averages cost over n_rollouts stochastic denoising seeds. If return_first_latents, also returns the first rollout's (1+H, ...) latents for visualization.""" use_reward_cost = cost_kind in REWARD_COST_KINDS if use_reward_cost and rew_head is None: raise ValueError(f"cost_kind={cost_kind!r} requires a reward head; got rew_head=None") full = torch.zeros((n_rollouts, horizon, 16), device=device, dtype=torch.float32) full[:, :, :act_dim] = action_plan.unsqueeze(0).expand(n_rollouts, -1, -1) full = full * act_mask_1d.view(1, 1, 16) z_start_b = z_start.unsqueeze(0).expand(n_rollouts, -1, -1).contiguous() if use_reward_cost: z_seq, h_seq = imagine_batch( dyn, z_start=z_start_b, action_plan=full, horizon=horizon, sched=sched, k_max=k_max, act_mask_1d=act_mask_1d, lang_emb=lang_emb, use_amp=use_amp, return_h=True, ) costs = reward_cost(rew_head, h_seq, use_amp=use_amp, discount=discount) else: z_seq = imagine_batch( dyn, z_start=z_start_b, action_plan=full, horizon=horizon, sched=sched, k_max=k_max, act_mask_1d=act_mask_1d, lang_emb=lang_emb, use_amp=use_amp, ) costs = latent_cost(z_seq, z_goal, kind=cost_kind) z0 = z_seq[0].detach() if return_first_latents else None return float(costs.mean().item()), z0 # -------------------- visualization -------------------- # def _pick_viz_indices(T: int, max_frames: int) -> List[int]: """Linearly-spaced indices into a T-length sequence, capped at max_frames. Always includes 0 and T-1 (start and terminal).""" if T <= max_frames: return list(range(T)) return sorted(set(np.linspace(0, T - 1, num=max_frames, dtype=int).tolist())) def _tile_row(frames_tchw: torch.Tensor, gap_px: int = 4) -> torch.Tensor: """(T, C, H, W) -> (C, H, T*W + (T-1)*gap_px) with black gaps between frames.""" T, C, Hp, Wp = frames_tchw.shape if gap_px <= 0 or T == 1: return frames_tchw.permute(1, 2, 0, 3).contiguous().view(C, Hp, T * Wp) panels = [] for i in range(T): panels.append(frames_tchw[i]) if i < T - 1: panels.append(torch.zeros((C, Hp, gap_px), dtype=frames_tchw.dtype, device=frames_tchw.device)) return torch.cat([p if p.dim() == 3 else p for p in panels], dim=-1) def save_viz( *, path: str, real_gt: torch.Tensor, # (T_viz, C, H, W) wm_gt: torch.Tensor, # (T_viz, C, H, W) wm_cem: torch.Tensor, # (T_viz, C, H, W) row_gap_px: int = 8, ) -> None: assert real_gt.shape == wm_gt.shape == wm_cem.shape, "row shapes must match" C, Hp, Wp = real_gt.shape[1:] r1 = _tile_row(real_gt) r2 = _tile_row(wm_gt) r3 = _tile_row(wm_cem) gap = torch.zeros((C, row_gap_px, r1.shape[-1]), dtype=r1.dtype, device=r1.device) big = torch.cat([r1, gap, r2, gap, r3], dim=1) # (C, 3H + 2*gap, W*T + gaps) big_u8 = (big.clamp(0, 1) * 255.0).to(torch.uint8).cpu().numpy() big_hwc = np.transpose(big_u8, (1, 2, 0)) Image.fromarray(big_hwc, mode="RGB").save(path) # -------------------- model / task setup -------------------- # @torch.inference_mode() def bc_rollout_integrated( policy_head, dyn, *, z_start: torch.Tensor, horizon: int, act_dim: int, act_mask_1d: torch.Tensor, sched: dict, k_max: int, lang_emb: Optional[torch.Tensor], use_amp: bool, ) -> Tuple[torch.Tensor, torch.Tensor]: """Closed-loop BC rollout using the *integrated* policy head. The integrated policy is `policy_head(h_t)` where h_t is the dynamics transformer's task-conditioned agent-token state at rollout slot t. Task conditioning is handled inside the dynamics via `task_proj(lang_emb)`, so this rollout is zero-shot across all 200 tasks. Convention recap: training shifts BC targets so at slot t (context z_0..z_t, a_0..a_t in led-to convention, a_t led to z_t), head l=0 predicts a_{t+1} (the action to apply at z_t producing z_{t+1}). To predict the first action a_1 at z_0 we need h_0 — obtained here via a one-shot prefill dynamics pass on past=[z_0] before the rollout loop starts. Returns (z_seq, bc_plan): z_seq: (1 + horizon, n_spatial, d_spatial) with z_seq[0] = z_start bc_plan: (horizon, act_dim) of the post-clip, post-mask actions taken """ device = z_start.device emax = int(round(math.log2(int(k_max)))) n_spatial, d_spatial = z_start.shape lang_emb_1 = lang_emb.expand(1, -1) if lang_emb is not None else None past = z_start.unsqueeze(0).unsqueeze(0).contiguous() # (1, 1, n_spatial, d_spatial) full_actions = torch.zeros((1, horizon + 1, 16), device=device, dtype=torch.float32) full_mask = act_mask_1d.view(1, 1, -1).expand(1, horizon + 1, -1).contiguous() # Prefill: one clean dynamics pass on past=[z_0] to extract h_0 (the agent- # token state at slot 0). step_idx=emax + signal_idx=k_max mark the slot as # fully denoised (matches how context slots are flagged in the denoising # sampler, so behavior is consistent with the learned training distribution). with torch.autocast(device_type=device.type, enabled=(use_amp and device.type == "cuda")): step_idx_pre = torch.full((1, 1), emax, device=device, dtype=torch.long) sig_idx_pre = torch.full((1, 1), int(k_max), device=device, dtype=torch.long) _, h_full = dyn( full_actions[:, :1], step_idx_pre, sig_idx_pre, past, act_mask=full_mask[:, :1], agent_tokens=None, lang_emb=lang_emb_1, ) if h_full is None: raise RuntimeError("bc_rollout_integrated: dynamics returned h_t=None; need n_agent>0 ckpt.") h_cur = h_full[:, -1:] # (1, 1, n_agent, D) — h at slot 0 bc_plan = torch.zeros((horizon, act_dim), device=device, dtype=torch.float32) for t in range(horizon): with torch.autocast(device_type=device.type, enabled=(use_amp and device.type == "cuda")): a_all = policy_head(h_cur) # (1, 1, L, A) -- tanh-squashed in [-1,1] a_pred = a_all[:, 0, 0, :act_dim].float().clamp(-1.0, 1.0).squeeze(0) full_actions[0, t + 1, :act_dim] = a_pred * act_mask_1d[:act_dim] bc_plan[t] = a_pred z_next, h_new, _inst = sample_one_timestep_packed( dyn, past_packed=past, k_max=k_max, sched=sched, actions=full_actions[:, : past.shape[1] + 1], act_mask=full_mask[:, : past.shape[1] + 1], use_amp=use_amp, return_h=True, tau_ctx=0.0, lang_emb=lang_emb_1, use_kv_cache=True, ) past = torch.cat([past, z_next.unsqueeze(1)], dim=1) h_cur = h_new # (1, 1, n_agent, D) — h at newly predicted slot t+1 return past[0], bc_plan def build_models(args, device): tok, tok_info = load_tokenizer_from_ckpt(args.tokenizer_ckpt, device) dyn, rew_head, policy_head, dyn_info = load_dynamics_from_ckpt( args.dynamics_ckpt, device=device, d_bottleneck=int(tok_info["d_bottleneck"]), n_latents=int(tok_info["n_latents"]), packing_factor=int(args.packing_factor), ) sched = make_tau_schedule( k_max=int(dyn_info["k_max"]), schedule=args.schedule, d=(args.eval_d if args.schedule == "shortcut" else None), ) info = {**tok_info, **dyn_info, "packing_factor": int(args.packing_factor)} return tok.encoder, tok.decoder, dyn, rew_head, policy_head, sched, info def load_lang_emb(tasks_json: str, task: str, lang_dim: int, device) -> Optional[torch.Tensor]: if not tasks_json or not os.path.exists(tasks_json): return None with open(tasks_json, "r") as f: meta = json.load(f) te = meta.get(task, {}).get("text_embedding") if te is None: return None emb = torch.tensor(te, dtype=torch.float32, device=device) if emb.numel() != lang_dim: return None return emb.view(1, -1) def load_act_meta(tasks_json: str, task: str, device=None): with open(tasks_json, "r") as f: meta = json.load(f) act_dim = int(meta.get(task, {}).get("action_dim", 16)) act_dim = max(0, min(16, act_dim)) mask = torch.zeros(16, dtype=torch.float32) if act_dim > 0: mask[:act_dim] = 1.0 if device is not None: mask = mask.to(device) return act_dim, mask # -------------------- per-plan evaluation -------------------- # def plan_one( *, dyn, encoder, decoder, info, sched, args, device, use_amp, episode: Episode, start_idx: int, act_dim: int, act_mask_1d: torch.Tensor, lang_emb: Optional[torch.Tensor], rng: torch.Generator, rew_head=None, policy_head=None, ) -> Optional[dict]: H_img, W_img, C, patch = int(info["H"]), int(info["W"]), int(info["C"]), int(info["patch"]) n_spatial, packing_factor = int(info["n_spatial"]), int(info["packing_factor"]) d_bottleneck, k_max = int(info["d_bottleneck"]), int(info["k_max"]) frames = episode.frames.to(device) actions_stored = episode.actions.to(device) L = frames.shape[0] max_h_configured = (args.max_horizon if args.max_horizon and args.max_horizon > 0 else L - 1) horizon = min(L - 1 - start_idx, max_h_configured) if horizon < 2: return None # Encode start, goal, and real-GT viz frames. viz_idx_real = _pick_viz_indices(horizon + 1, args.n_viz_frames) viz_abs_idx = [start_idx + i for i in viz_idx_real] # Include both endpoints; also ensure we encode them. to_encode_idx = sorted(set([start_idx, L - 1] + viz_abs_idx)) frames_to_encode = frames[to_encode_idx] z_encoded = encode_frames_to_packed( encoder, frames_to_encode, patch=patch, n_spatial=n_spatial, packing_factor=packing_factor, use_amp=use_amp, ) idx_to_z = {i: z for i, z in zip(to_encode_idx, z_encoded)} z_start = idx_to_z[start_idx] z_goal = idx_to_z[L - 1] # GT action slice: actions[start_idx+1 : start_idx+1+horizon] under the # stored convention (action[k] produced obs[k]). These produce z_{start+1..start+horizon}. gt_plan_full = actions_stored[start_idx + 1: start_idx + 1 + horizon] gt_plan = gt_plan_full[:, :act_dim].clamp(-1, 1) if torch.isnan(gt_plan).any(): return None # Baselines. rand_plan = torch.rand((horizon, act_dim), device=device, generator=rng) * 2.0 - 1.0 ep_acts = actions_stored[1:, :act_dim] perm = torch.randperm(ep_acts.shape[0], generator=rng, device=device) shuf_plan = ep_acts[perm[:horizon]].clamp(-1, 1) def _sync(): if device.type == "cuda": torch.cuda.synchronize() # Warm-start: BC rollout -> action sequence used to seed CEM's mu. bc_init_plan = None cost_bc_init = float("nan") warm_start = str(getattr(args, "warm_start", "none")) cem_init_std_eff = float(args.cem_init_std) if warm_start == "bc": assert policy_head is not None, "warm_start=bc requires a loaded policy head" _sync(); t_bcr = time.time() _z_bc, bc_init_plan = bc_rollout_integrated( policy_head, dyn, z_start=z_start, horizon=horizon, act_dim=act_dim, act_mask_1d=act_mask_1d, sched=sched, k_max=k_max, lang_emb=lang_emb, use_amp=use_amp, ) bc_init_plan = bc_init_plan.clamp(-1.0, 1.0) _sync(); dt_bcr = time.time() - t_bcr t_bcs = time.time() cost_bc_init, _ = score_plan( dyn, z_start=z_start, z_goal=z_goal, action_plan=bc_init_plan, act_dim=act_dim, act_mask_1d=act_mask_1d, horizon=horizon, sched=sched, k_max=k_max, lang_emb=lang_emb, cost_kind=args.cost, use_amp=use_amp, device=device, n_rollouts=args.baseline_rollouts, rew_head=rew_head, ) _sync(); dt_bcs = time.time() - t_bcs print(f" [bc-warm] rollout={dt_bcr:.2f}s rescore={dt_bcs:.2f}s " f"cost_bc_init={cost_bc_init:.3f}", flush=True) cem_init_std_eff = float(args.warm_start_init_std) # CEM. _sync(); t_cem = time.time() out = cem_plan( dyn, z_start=z_start, z_goal=z_goal, horizon=horizon, act_dim=act_dim, act_mask_1d=act_mask_1d, sched=sched, k_max=k_max, lang_emb=lang_emb, cost_kind=args.cost, n_samples=args.cem_samples, n_elites=args.cem_elites, n_iters=args.cem_iters, init_std=cem_init_std_eff, min_std=args.cem_min_std, rollout_batch=args.rollout_batch, noise_beta=args.noise_beta, device=device, use_amp=use_amp, generator=rng, init_mu=bc_init_plan, n_rollouts_elite=int(args.cem_n_rollouts), rew_head=rew_head, ) _sync(); dt_cem = time.time() - t_cem print(f" [cem] {dt_cem:.2f}s n_samples={args.cem_samples} iters={args.cem_iters} " f"n_rollouts_elite={args.cem_n_rollouts} rollout_batch={args.rollout_batch} " f"best={out['best_cost']:.3f}", flush=True) # Score all plans with matched rollout budget; keep latents for viz. _sync(); t_sc = time.time() cost_gt, z_seq_gt = score_plan( dyn, z_start=z_start, z_goal=z_goal, action_plan=gt_plan, act_dim=act_dim, act_mask_1d=act_mask_1d, horizon=horizon, sched=sched, k_max=k_max, lang_emb=lang_emb, cost_kind=args.cost, use_amp=use_amp, device=device, n_rollouts=args.baseline_rollouts, return_first_latents=True, rew_head=rew_head, ) cost_rand, _ = score_plan( dyn, z_start=z_start, z_goal=z_goal, action_plan=rand_plan, act_dim=act_dim, act_mask_1d=act_mask_1d, horizon=horizon, sched=sched, k_max=k_max, lang_emb=lang_emb, cost_kind=args.cost, use_amp=use_amp, device=device, n_rollouts=args.baseline_rollouts, rew_head=rew_head, ) cost_shuf, _ = score_plan( dyn, z_start=z_start, z_goal=z_goal, action_plan=shuf_plan, act_dim=act_dim, act_mask_1d=act_mask_1d, horizon=horizon, sched=sched, k_max=k_max, lang_emb=lang_emb, cost_kind=args.cost, use_amp=use_amp, device=device, n_rollouts=args.baseline_rollouts, rew_head=rew_head, ) cost_cem, z_seq_cem = score_plan( dyn, z_start=z_start, z_goal=z_goal, action_plan=out["best_plan"], act_dim=act_dim, act_mask_1d=act_mask_1d, horizon=horizon, sched=sched, k_max=k_max, lang_emb=lang_emb, cost_kind=args.cost, use_amp=use_amp, device=device, n_rollouts=args.baseline_rollouts, return_first_latents=True, rew_head=rew_head, ) _sync(); dt_sc = time.time() - t_sc print(f" [score] {dt_sc:.2f}s gt={cost_gt:.3f} rand={cost_rand:.3f} " f"shuf={cost_shuf:.3f} cem={cost_cem:.3f}", flush=True) row = { "task": str(args.task), "ep_id": int(episode.ep_id), "start": int(start_idx), "ep_len": int(L), "horizon": int(horizon), "cost_kind": args.cost, "cost_gt": cost_gt, "cost_random": cost_rand, "cost_shuffled": cost_shuf, "cost_bc_init": cost_bc_init, "cost_cem_best": cost_cem, "cem_seconds": dt_cem, "cem_iters": int(args.cem_iters), "cem_samples": int(args.cem_samples), "noise_beta": float(args.noise_beta), "warm_start": warm_start, "cem_init_std_eff": float(cem_init_std_eff), "cem_n_rollouts": int(args.cem_n_rollouts), "iters_running_best": ",".join(f"{h['running_best_cost']:.3f}" for h in out["history"]), } # Visualization + per-plan JSON (incremental, crash-safe). Both live in # viz_dir with a shared basename so they can be inspected together. if args.viz_dir: os.makedirs(args.viz_dir, exist_ok=True) base = f"{args.task}_ep{episode.ep_id:03d}_s{start_idx:03d}" z_gt_viz = z_seq_gt[viz_idx_real] z_cem_viz = z_seq_cem[viz_idx_real] wm_gt_frames = decode_packed_sequence( decoder, z_gt_viz, H=H_img, W=W_img, C=C, patch=patch, packing_factor=packing_factor, d_bottleneck=d_bottleneck, use_amp=use_amp, ) wm_cem_frames = decode_packed_sequence( decoder, z_cem_viz, H=H_img, W=W_img, C=C, patch=patch, packing_factor=packing_factor, d_bottleneck=d_bottleneck, use_amp=use_amp, ) real_gt_frames = frames[viz_abs_idx] save_viz(path=os.path.join(args.viz_dir, base + ".png"), real_gt=real_gt_frames, wm_gt=wm_gt_frames, wm_cem=wm_cem_frames) detail = dict(row) detail["task"] = str(args.task) detail["seed"] = int(args.seed) detail["cem_history"] = out["history"] detail["plans"] = { "gt": gt_plan.detach().cpu().tolist(), "random": rand_plan.detach().cpu().tolist(), "shuffled": shuf_plan.detach().cpu().tolist(), "cem_best": out["best_plan"].detach().cpu().tolist(), "bc_init": (bc_init_plan.detach().cpu().tolist() if bc_init_plan is not None else None), } json_path = os.path.join(args.viz_dir, base + ".json") tmp_path = json_path + ".tmp" with open(tmp_path, "w") as f: json.dump(detail, f, indent=2) os.replace(tmp_path, json_path) return row # -------------------- main -------------------- # def evaluate(args): rank = int(args.rank) world_size = int(args.world_size) assert world_size >= 1 and 0 <= rank < world_size, \ f"invalid rank/world_size: rank={rank} world_size={world_size}" tag = f"[rank {rank}/{world_size}] " if world_size > 1 else "" # Auto-suffix shard outputs so concurrent ranks don't collide. if world_size > 1: if args.output: base, ext = os.path.splitext(args.output) args.output = f"{base}.rank{rank}{ext}" if args.viz_dir: args.viz_dir = f"{args.viz_dir}.rank{rank}" device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") torch.manual_seed(args.seed) np.random.seed(args.seed) py_rng = np.random.default_rng(args.seed) torch_rng = torch.Generator(device=device).manual_seed(args.seed) encoder, decoder, dyn, rew_head, policy_head, sched, info = build_models(args, device) use_amp = bool(args.amp) and device.type == "cuda" if str(args.cost) in REWARD_COST_KINDS and rew_head is None: raise SystemExit( f"--cost={args.cost!r} requires a reward head in the dynamics ckpt. " f"Loaded ckpt {args.dynamics_ckpt!r} has none; use a ckpt " f"trained with --reward_weight > 0 and --n_agent > 0, or switch " f"--cost back to {'|'.join(LATENT_COST_KINDS)}." ) if str(args.warm_start) == "bc" and policy_head is None: raise SystemExit( f"--warm_start=bc requires a policy head in the dynamics ckpt. " f"Loaded ckpt {args.dynamics_ckpt!r} has none; use a ckpt " f"trained with --bc_weight > 0 and --n_agent > 0, or switch " f"--warm_start to 'none'." ) act_dim, act_mask_1d = load_act_meta(args.tasks_json, args.task, device=device) lang_emb = load_lang_emb(args.tasks_json, args.task, int(args.lang_dim), device) if act_dim == 0: raise SystemExit(f"task {args.task} has act_dim=0; nothing to plan") if str(args.warm_start) == "bc": print(f"{tag}[plan_cem] BC warm-start via integrated policy head; " f"init_std={args.warm_start_init_std}") all_ep_ids = list_episode_ids(args.task, args.data_dir) n_eps = min(args.n_episodes, len(all_ep_ids)) chosen_eps = py_rng.choice(all_ep_ids, size=n_eps, replace=False).tolist() print(f"{tag}[plan_cem] task={args.task} act_dim={act_dim} " f"n_spatial={info['n_spatial']} k_max={info['k_max']} sched.K={sched['K']}") print(f"{tag}[plan_cem] sampled episodes: {chosen_eps}") all_rows = [] t0 = time.time() plan_idx = 0 # global counter across all ranks; partition via plan_idx % world_size for ep_id in chosen_eps: epi = load_episode(args.task, args.data_dir, args.frames_dir, ep_id, shard_size=int(args.shard_size)) L = epi.frames.shape[0] upper = max(1, min(int(args.max_start_frame), L - 1)) # Draw starts identically on every rank (same seed -> same py_rng state), # so the plan_idx-based partition below is consistent across the cluster. starts = py_rng.integers(0, upper, size=int(args.n_starts_per_ep)).tolist() starts = sorted(set(int(s) for s in starts)) for s in starts: if plan_idx % world_size != rank: plan_idx += 1 continue row = plan_one( dyn=dyn, encoder=encoder, decoder=decoder, info=info, sched=sched, args=args, device=device, use_amp=use_amp, episode=epi, start_idx=int(s), act_dim=act_dim, act_mask_1d=act_mask_1d, lang_emb=lang_emb, rng=torch_rng, rew_head=rew_head, policy_head=policy_head, ) plan_idx += 1 if row is None: print(f"{tag}[plan_cem] ep {ep_id} s={s}: skipped (short horizon or nan actions)") continue all_rows.append(row) bc_str = f" bc_init={row['cost_bc_init']:.3f}" if row.get("warm_start") == "bc" else "" print(f"{tag}[ep {ep_id:3d} s={s:3d} H={row['horizon']:3d}] " f"gt={row['cost_gt']:.3f} rand={row['cost_random']:.3f} " f"shuf={row['cost_shuffled']:.3f}{bc_str} cem={row['cost_cem_best']:.3f} " f"(iters=[{row['iters_running_best']}], {row['cem_seconds']:.1f}s)") if not all_rows: print(f"{tag}[plan_cem] no rows collected; exiting") return def arr(k): return np.array([r[k] for r in all_rows], dtype=np.float64) gt, rnd, shf, cem = arr("cost_gt"), arr("cost_random"), arr("cost_shuffled"), arr("cost_cem_best") horizons = arr("horizon") scope = "this rank only" if world_size > 1 else "all plans" print(f"\n{tag}========== aggregate over {len(all_rows)} plans " f"({scope}; H range {int(horizons.min())}..{int(horizons.max())}, " f"mean {horizons.mean():.1f}) ==========") print(f" cost_gt mean={gt.mean():.3f} median={np.median(gt):.3f}") print(f" cost_random mean={rnd.mean():.3f} median={np.median(rnd):.3f}") print(f" cost_shuffled mean={shf.mean():.3f} median={np.median(shf):.3f}") if any(r.get("warm_start") == "bc" for r in all_rows): bci = arr("cost_bc_init") print(f" cost_bc_init mean={bci.mean():.3f} median={np.median(bci):.3f} " f"(CEM warm-started from BC plan)") print(f" cost_cem_best mean={cem.mean():.3f} median={np.median(cem):.3f}") # Validation framing depends on cost semantics. Latent distances are # non-negative: GT is near-zero under good WM, so gt*1.1 is a meaningful # "match-expert" tolerance and (cem-gt)/gt is a signed relative error. # Reward-cost is negated predicted return and typically spans both signs; # gt*1.1 flips meaning and relative error blows up near zero. Report # reward-appropriate statistics instead. if str(args.cost) in REWARD_COST_KINDS: r_gt, r_rnd, r_cem = -gt, -rnd, -cem print(f" reward_gt mean={r_gt.mean():+.3f} median={np.median(r_gt):+.3f}") print(f" reward_random mean={r_rnd.mean():+.3f} median={np.median(r_rnd):+.3f}") print(f" reward_cem mean={r_cem.mean():+.3f} median={np.median(r_cem):+.3f}") print(f" (reward_gt > reward_random) rate : {(r_gt > r_rnd).mean():.1%} <- validation #1 (predicted reward is informative)") print(f" (reward_cem >= reward_gt) rate : {(r_cem >= r_gt).mean():.1%} <- validation #2 (CEM >= GT predicted return)") print(f" reward_cem - reward_gt : mean={(r_cem - r_gt).mean():+.3f} median={np.median(r_cem - r_gt):+.3f} (positive = CEM improves on GT)") else: print(f" (gt < random) rate : {(gt < rnd).mean():.1%} <- validation #1 (cost is informative)") print(f" (cem <= gt * 1.1) rate : {(cem <= gt * 1.1).mean():.1%} <- validation #2 (CEM recovers expert-quality plans)") print(f" (cem < random) rate : {(cem < rnd).mean():.1%}") rel = (cem - gt) / np.maximum(gt, 1e-6) print(f" (cem - gt)/gt : mean={rel.mean():+.1%} median={np.median(rel):+.1%}") print(f" total wallclock: {time.time() - t0:.1f}s") if args.output: os.makedirs(os.path.dirname(args.output) or ".", exist_ok=True) import csv with open(args.output, "w", newline="") as f: w = csv.DictWriter(f, fieldnames=list(all_rows[0].keys())) w.writeheader() w.writerows(all_rows) print(f"[plan_cem] wrote {args.output}") if args.viz_dir: print(f"[plan_cem] wrote visualizations to {args.viz_dir}") def main(): p = argparse.ArgumentParser() # task / data p.add_argument("--task", type=str, default="og-point-maze") p.add_argument("--data_dir", type=str, default="./data/expert") p.add_argument("--frames_dir", type=str, default="./data/expert-shards") p.add_argument("--tasks_json", type=str, default="../tasks.json") p.add_argument("--shard_size", type=int, default=4096) p.add_argument("--lang_dim", type=int, default=512) # checkpoints p.add_argument("--tokenizer_ckpt", type=str, default="./logs/tokenizer_ckpts/latest.pt") p.add_argument("--dynamics_ckpt", type=str, default="./logs/dynamics_ckpts/latest.pt") # rollout / denoising schedule p.add_argument("--packing_factor", type=int, default=2) p.add_argument("--schedule", type=str, default="shortcut", choices=["finest", "shortcut"]) p.add_argument("--eval_d", type=float, default=0.125) p.add_argument("--amp", action="store_true", default=True) # planning problem p.add_argument( "--cost", type=str, default="terminal", choices=list(LATENT_COST_KINDS) + list(REWARD_COST_KINDS), help=( "CEM cost. Latent modes ({LATENT}) score z_seq vs z_goal " "(reward-head-free). 'reward' uses the dynamics ckpt's reward head: " "cost = -sum_t r_hat_t (l=0 MTP slot at each rollout step). Requires " "a ckpt trained with n_agent>0 and reward_weight>0." ).format(LATENT="|".join(LATENT_COST_KINDS)), ) p.add_argument("--n_episodes", type=int, default=4) p.add_argument("--n_starts_per_ep", type=int, default=3) p.add_argument("--max_start_frame", type=int, default=5, help="Sample each plan's start index uniformly from [0, max_start_frame). " "Default 5 keeps starts close to the episode beginning so horizons are " "comparable across tasks with different episode lengths.") p.add_argument("--max_horizon", type=int, default=0, help="0 = use full remaining length (L - 1 - start). Cap for compute if needed.") # CEM p.add_argument("--cem_samples", type=int, default=256) p.add_argument("--cem_elites", type=int, default=32) p.add_argument("--cem_iters", type=int, default=4) p.add_argument("--cem_init_std", type=float, default=1.0) p.add_argument("--cem_min_std", type=float, default=0.05) p.add_argument("--rollout_batch", type=int, default=128) p.add_argument("--cem_n_rollouts", type=int, default=1, help="Rollouts per sample during CEM elite selection. " ">1 averages denoising-seed noise so the selector can't " "overfit to a lucky single rollout.") p.add_argument("--noise_beta", type=float, default=0.0, help="Action-noise power spectrum 1/f^beta. " "0 = white (default); " "1 = pink; 2 = Brownian (high temporal correlation, " "recommended for navigation / smooth-velocity tasks).") # baselines p.add_argument("--baseline_rollouts", type=int, default=4, help="stochastic rollouts to average for each reference plan's cost") # BC warm-start p.add_argument("--warm_start", type=str, default="none", choices=["none", "bc"], help=( "If 'bc', seed CEM's mu with a closed-loop rollout of the 200-task " "language-conditioned policy head baked into the dynamics " "ckpt (reads dynamics h_t agent tokens, zero-shot across tasks), " "and use --warm_start_init_std instead of --cem_init_std." )) p.add_argument("--warm_start_init_std", type=float, default=0.3, help="CEM init_std when warm-starting from BC (should be < cem_init_std).") # outputs p.add_argument("--output", type=str, default="./logs/plan_cem_og-point-maze.csv") p.add_argument("--viz_dir", type=str, default="./logs/plan_cem_viz") p.add_argument("--n_viz_frames", type=int, default=12, help="number of timesteps to display in each visualization row") # misc p.add_argument("--seed", type=int, default=0) # multi-GPU sharding (plans are embarrassingly parallel; no torch.distributed needed). # Launch N processes with CUDA_VISIBLE_DEVICES=i and --rank i --world_size N; # each rank owns plans whose global index satisfies `i % world_size == rank`. # --output and --viz_dir are auto-suffixed with `.rank{i}` when world_size > 1. # Concatenate the per-rank shards to obtain the global aggregate. p.add_argument("--rank", type=int, default=0) p.add_argument("--world_size", type=int, default=1) args = p.parse_args() evaluate(args) if __name__ == "__main__": main()