"""Core inference + rendering helpers for the FACT counterfactual-imagination demo. Kept separate from ``app.py`` so the whole pipeline call / post-processing path can be exercised with stub models (no 10 GB checkpoint) while developing. Everything here mirrors the reference implementation shipped with the paper (``scripts/inference_server.py`` in https://github.com/Bariona/FACT): same 384x192 three-view reference canvas, 48-step action chunk, 5 predicted keyframes, 20 UniPC steps, flow shift 3.0, guidance 0.0, z-score normalisation and delta-action decoding. """ from __future__ import annotations import json import tempfile from dataclasses import dataclass from typing import Dict, List, Optional, Sequence, Tuple import numpy as np import torch from PIL import Image, ImageDraw from world_action_model.image_layouts import ( ROBOTWIN_VIEW_KEYS, build_robotwin_ref_tensor, infer_robotwin_main_view_size, ) from world_action_model.pipeline.utils import ( NormalizationTensors, add_state_to_action, denormalize_action, denormalize_state, denormalize_value, normalize_state, pad_t5_embedding, ) # --------------------------------------------------------------------------- # Constants — all taken from evaluation/robotwin/launch_config.yml (released setup) # --------------------------------------------------------------------------- DST_W, DST_H = 384, 192 # reference canvas (main view + stacked wrist views) ACTION_CHUNK = 48 # predicted actions per plan NUM_FRAMES = 5 # 1 reference + 4 predicted future keyframes NUM_STEPS = 20 # UniPC steps GUIDANCE = 0.0 # no CFG FLOW_SHIFT = 3.0 TEXT_LEN = 512 # tokenizer max length T5_LEN = 64 # T5 tokens kept as conditioning ACTION_DIM = STATE_DIM = 14 # 2 arms x (6 joints + gripper) NORM_MODE = "zscore" FPS = 15 # RoboTwin control/record rate DELTA_MASK: Tuple[int, ...] = (1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 0) GRIPPER_DIMS = (6, 13) ARM_DIMS = tuple(d for d in range(ACTION_DIM) if d not in GRIPPER_DIMS) KEYFRAME_OFFSETS = [round(ACTION_CHUNK * i / (NUM_FRAMES - 1)) for i in range(NUM_FRAMES)] # 0,12,24,36,48 JOINT_LABELS = [f"joint {i + 1}" for i in range(6)] STATE_FIELDS = ( [f"left {n}" for n in JOINT_LABELS] + ["left gripper"] + [f"right {n}" for n in JOINT_LABELS] + ["right gripper"] ) COUNTERFACTUALS: Tuple[str, ...] = ( "Freeze the arms", "Reverse the motion", "Overshoot (move twice as far)", "Open the grippers (let go)", "Squeeze the grippers shut", "Swap left and right arm", "Random jitter", ) @dataclass class Models: pipe: object tokenizer: object text_encoder: object norm: NormalizationTensors device: torch.device dtype: torch.dtype # --------------------------------------------------------------------------- # Inputs # --------------------------------------------------------------------------- def parse_state(text: str) -> List[float]: """Parse the 14-dim joint state from a textbox (JSON list or comma/space separated).""" raw = (text or "").strip() if not raw: raise ValueError("The robot joint state is empty.") try: values = json.loads(raw if raw.startswith("[") else f"[{raw}]") except json.JSONDecodeError: values = [p for p in raw.replace(",", " ").split() if p] try: out = [float(v) for v in values] except (TypeError, ValueError): raise ValueError("The robot joint state must be 14 numbers.") if len(out) != STATE_DIM: raise ValueError(f"The robot joint state must have exactly {STATE_DIM} numbers, got {len(out)}.") return out def _to_chw_uint8(image) -> torch.Tensor: arr = np.asarray(image) if arr.ndim == 2: arr = np.stack([arr] * 3, axis=-1) if arr.shape[-1] == 4: arr = arr[..., :3] if arr.dtype != np.uint8: arr = np.clip(arr, 0, 255).astype(np.uint8) return torch.from_numpy(np.ascontiguousarray(arr)).permute(2, 0, 1) def build_reference(cam_high, cam_left, cam_right) -> torch.Tensor: """Compose the three camera views into the 3x192x384 float[0,1] reference canvas.""" views = { ROBOTWIN_VIEW_KEYS[0]: _to_chw_uint8(cam_high), ROBOTWIN_VIEW_KEYS[1]: _to_chw_uint8(cam_left), ROBOTWIN_VIEW_KEYS[2]: _to_chw_uint8(cam_right), } return build_robotwin_ref_tensor(views, main_dst_size=infer_robotwin_main_view_size((DST_W, DST_H))) @torch.no_grad() def encode_instruction(models: Models, instruction: str) -> torch.Tensor: """UMT5 encode -> [T5_LEN, 4096] float32, exactly as the reference server does.""" text = (instruction or "").strip() if not text: raise ValueError("The language instruction is empty.") batch = models.tokenizer( [text], padding="max_length", max_length=TEXT_LEN, truncation=True, add_special_tokens=True, return_attention_mask=True, return_tensors="pt", ) input_ids = batch.input_ids.to(models.device) attention_mask = batch.attention_mask.to(models.device) hidden = models.text_encoder(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state seq_len = int(attention_mask.gt(0).sum(dim=1)[0].item()) return pad_t5_embedding(hidden[0, :seq_len].detach().float().cpu(), target_len=T5_LEN) # --------------------------------------------------------------------------- # Counterfactual action edits (in de-normalised delta space) # --------------------------------------------------------------------------- def counterfactual_target(delta: torch.Tensor, mode: str, generator: torch.Generator, action_std: torch.Tensor) -> torch.Tensor: """The 'extreme' version of the edit; `strength` later interpolates plan -> target.""" target = delta.clone() arm = list(ARM_DIMS) grip = list(GRIPPER_DIMS) if mode == "Freeze the arms": target[:, arm] = 0.0 elif mode == "Reverse the motion": target[:, arm] = -delta[:, arm] elif mode == "Overshoot (move twice as far)": target[:, arm] = 2.0 * delta[:, arm] elif mode == "Open the grippers (let go)": target[:, grip] = 1.0 elif mode == "Squeeze the grippers shut": target[:, grip] = 0.0 elif mode == "Swap left and right arm": target = torch.cat([delta[:, 7:], delta[:, :7]], dim=1) elif mode == "Random jitter": noise = torch.randn(delta.shape, generator=generator, device=delta.device, dtype=delta.dtype) target[:, arm] = delta[:, arm] + 0.5 * noise[:, arm] * action_std[arm] else: raise ValueError(f"unknown counterfactual: {mode}") return target def _clamp(x: torch.Tensor, lo: torch.Tensor, hi: torch.Tensor) -> torch.Tensor: return torch.maximum(torch.minimum(x, hi), lo) def _normalize_action(delta: torch.Tensor, norm: NormalizationTensors) -> torch.Tensor: return (delta - norm.action_mean) / norm.action_std.clamp_min(1e-8) def absolute_joints(delta: torch.Tensor, state: torch.Tensor, norm: NormalizationTensors) -> torch.Tensor: """delta-action -> absolute joint targets, mirroring the reference post-processing.""" mask = torch.tensor(DELTA_MASK, device=delta.device, dtype=torch.bool) abs_a = add_state_to_action(delta, state, action_chunk=delta.shape[0], mask=mask) fallback = state.unsqueeze(0).repeat(delta.shape[0], 1) abs_a = torch.where(torch.isfinite(abs_a), abs_a, fallback) return _clamp(abs_a, norm.state_min, norm.state_max) # --------------------------------------------------------------------------- # Rendering # --------------------------------------------------------------------------- def _keyframes_uint8(imgs: torch.Tensor) -> np.ndarray: """[C, F, H, W] in [-1, 1] -> [F, H, W, C] uint8.""" x = ((imgs.float().clamp(-1, 1) + 1.0) * 127.5).to(torch.uint8) return x.permute(1, 2, 3, 0).cpu().numpy() def write_rollout_video(keyframes: np.ndarray, path: str, upscale: int = 2) -> str: """Hold each predicted keyframe over its 12-frame segment (as in the reference dump).""" import imageio.v3 as iio frames: List[np.ndarray] = [] lead_in = 6 for _ in range(lead_in): # short look at the current observation frames.append(keyframes[0]) for i in range(1, keyframes.shape[0]): hold = KEYFRAME_OFFSETS[i] - KEYFRAME_OFFSETS[i - 1] for _ in range(hold): frames.append(keyframes[i]) stack = np.stack(frames) if upscale > 1: stack = stack.repeat(upscale, axis=1).repeat(upscale, axis=2) iio.imwrite(path, stack, fps=FPS, codec="libx264") return path def _font(size: int): from PIL import ImageFont try: return ImageFont.load_default(size=size) except TypeError: # Pillow < 10.1 return ImageFont.load_default() def build_filmstrip(policy: np.ndarray, counterfactual: np.ndarray, cf_label: str) -> str: """Two rows of keyframes (plan vs counterfactual) with time captions.""" tile_h, tile_w = policy.shape[1], policy.shape[2] gap, head, side = 4, 22, 104 cols = policy.shape[0] width = side + cols * (tile_w + gap) height = head + 2 * (tile_h + gap) + 18 canvas = Image.new("RGB", (width, height), (17, 17, 19)) draw = ImageDraw.Draw(canvas) small, tiny = _font(13), _font(12) for c in range(cols): x = side + c * (tile_w + gap) secs = KEYFRAME_OFFSETS[c] / FPS caption = "now (observation)" if c == 0 else f"+{secs:.1f}s" draw.text((x + 2, 4), caption, fill=(215, 215, 220), font=small) rows = [("FACT's own plan", policy, (120, 220, 150)), (cf_label, counterfactual, (250, 170, 120))] for r, (label, strip, colour) in enumerate(rows): y = head + r * (tile_h + gap) words, line, lines = label.split(), "", [] for w in words: trial = f"{line} {w}".strip() if len(trial) > 13: lines.append(line) line = w else: line = trial lines.append(line) for k, ln in enumerate(lines[:4]): draw.text((6, y + 4 + k * 15), ln, fill=colour, font=tiny) for c in range(cols): x = side + c * (tile_w + gap) canvas.paste(Image.fromarray(strip[c]), (x, y)) if c == 0: draw.rectangle([x, y, x + tile_w - 1, y + tile_h - 1], outline=(90, 90, 96)) draw.text((6, height - 16), "left column = the observation the model was given; the rest is imagined by FACT", fill=(140, 140, 148), font=tiny) path = tempfile.NamedTemporaryFile(suffix=".png", delete=False).name canvas.save(path) return path def build_action_plot(policy_abs: np.ndarray, cf_abs: np.ndarray, state: np.ndarray, policy_future: np.ndarray, cf_future: np.ndarray, cf_label: str) -> str: import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt t = np.arange(1, ACTION_CHUNK + 1) / FPS fig, axes = plt.subplots(3, 1, figsize=(7.2, 7.6), sharex=True) cmap = plt.get_cmap("tab10") groups = [(range(6), "Left arm joints (rad)"), (range(7, 13), "Right arm joints (rad)")] for ax, (dims, title) in zip(axes[:2], groups): for k, d in enumerate(dims): colour = cmap(k % 10) ax.plot(t, policy_abs[:, d], color=colour, lw=1.7, label=f"j{k + 1}") ax.plot(t, cf_abs[:, d], color=colour, lw=1.4, ls="--", alpha=0.85) ax.scatter([0.0], [state[d]], color=colour, s=12, zorder=5) ax.scatter([t[-1]], [policy_future[d]], color=colour, s=26, marker="*", zorder=5) ax.set_title(title, fontsize=10) ax.grid(alpha=0.25) ax.legend(fontsize=6, ncols=6, loc="upper right", framealpha=0.6) ax = axes[2] for k, (d, name) in enumerate([(6, "left gripper"), (13, "right gripper")]): colour = cmap(k) ax.plot(t, policy_abs[:, d], color=colour, lw=1.8, label=name) ax.plot(t, cf_abs[:, d], color=colour, lw=1.4, ls="--", alpha=0.85) ax.set_title("Grippers (1 = open, 0 = closed)", fontsize=10) ax.set_xlabel("time from now (s)") ax.grid(alpha=0.25) ax.legend(fontsize=7, loc="upper right", framealpha=0.6) fig.suptitle(f"solid = FACT's own plan dashed = {cf_label}\n" "dots = current joints, stars = FACT's predicted future state", fontsize=9) fig.tight_layout(rect=(0, 0, 1, 0.94)) path = tempfile.NamedTemporaryFile(suffix=".png", delete=False).name fig.savefig(path, dpi=110) plt.close(fig) return path def _bar(fraction: float, width: int = 22) -> str: filled = int(round(max(0.0, min(1.0, fraction)) * width)) return "█" * filled + "░" * (width - filled) def verdict_markdown(v_policy: float, v_cf: float, cf_label: str) -> str: def row(name: str, v: float) -> str: togo = max(0.0, min(1.0, v)) done = 1.0 - togo flag = "🔴 **failure predicted**" if v > 1.0 else ("🟡 borderline" if v > 0.95 else "🟢 on track") return f"| {name} | `{_bar(done)}` {done * 100:4.0f}% | `{v:+.3f}` | {flag} |" delta = v_cf - v_policy if v_cf > 1.0 >= v_policy: headline = (f"**FACT flags the counterfactual as a failure.** Following its own plan it expects the " f"episode to keep progressing, but under *{cf_label.lower()}* the predicted cost crosses " f"the failure threshold (value > 1).") elif delta > 0.05: headline = (f"**FACT expects *{cf_label.lower()}* to set the task back** — predicted time-to-go grows " f"by {delta:+.3f} (that is {delta * 100:.0f}% of the episode).") elif delta < -0.05: headline = (f"**FACT thinks *{cf_label.lower()}* would actually help here** — predicted time-to-go " f"shrinks by {-delta:.3f}. Not every edit is a mistake.") else: headline = (f"**FACT is nearly indifferent to *{cf_label.lower()}*** (Δ time-to-go {delta:+.3f}); " f"at this moment the edit barely changes how far the task is from done.") return ( "### Task-progress / failure head\n\n" "| rollout | progress after 3.2 s | value (time-to-go) | verdict |\n|---|---|---|---|\n" + row("FACT's own plan", v_policy) + "\n" + row(cf_label, v_cf) + "\n\n" + headline + "\n\n" "The value head regresses the *fraction of the episode still remaining* 3.2 s from now " "(0 = done, lower is better) and is trained with a +1 penalty on failure states, so anything " "above 1 means FACT believes the rollout has broken the task." ) # --------------------------------------------------------------------------- # Main entry point # --------------------------------------------------------------------------- @torch.no_grad() def run(models: Models, cam_high, cam_left, cam_right, instruction: str, state_text: str, mode: str, strength: float, steps: int, seed: int): for name, img in (("overhead", cam_high), ("left wrist", cam_left), ("right wrist", cam_right)): if img is None: raise ValueError(f"The {name} camera image is missing — all three views are required.") if mode not in COUNTERFACTUALS: raise ValueError(f"unknown counterfactual: {mode}") device, dtype, norm = models.device, models.dtype, models.norm steps = int(max(4, min(40, int(steps)))) seed = int(seed) % (2**31 - 1) state_list = parse_state(state_text) reference = build_reference(cam_high, cam_left, cam_right) prompt_embeds = encode_instruction(models, instruction).to(device) state = torch.tensor(state_list, dtype=torch.float32, device=device).unsqueeze(0) norm_state = normalize_state(state, norm, mode=NORM_MODE).to(device=device, dtype=dtype) common = dict( image=reference, action_chunk=ACTION_CHUNK, action_dim=ACTION_DIM, height=DST_H, width=DST_W, num_frames=NUM_FRAMES, num_inference_steps=steps, guidance_scale=GUIDANCE, return_dict=False, ) # Stage 1 only: what does FACT itself want to do here? _, action_raw, _, _ = models.pipe( state=norm_state, prompt_embeds=prompt_embeds.unsqueeze(0), generator=torch.Generator(device=device).manual_seed(seed), action_only=True, skip_future_state_value=True, enable_prefix_cache=True, **common, ) plan = denormalize_action(action_raw[0].float(), norm, mode=NORM_MODE) plan = torch.nan_to_num(plan, nan=0.0, posinf=0.0, neginf=0.0) plan = _clamp(plan, norm.action_min, norm.action_max) gen = torch.Generator(device=device).manual_seed(seed + 1) target = counterfactual_target(plan, mode, gen, norm.action_std) edited = _clamp(plan + float(strength) * (target - plan), norm.action_min, norm.action_max) # One batched call: identical noise for both rows (per-row generators with the same # seed), so the ONLY difference between the two imagined futures is the action. gt_action = torch.stack([_normalize_action(plan, norm), _normalize_action(edited, norm)]).to(dtype) imgs, _, future_state_raw, value_raw = models.pipe( state=norm_state.repeat(2, 1), prompt_embeds=prompt_embeds.unsqueeze(0).repeat(2, 1, 1), generator=[torch.Generator(device=device).manual_seed(seed) for _ in range(2)], action_only=False, gt_action_condition=gt_action, **common, ) values = [] futures = [] for i in range(2): v = denormalize_value(value_raw[i].float(), norm) v = float(torch.nan_to_num(v, nan=0.0).view(-1)[0].item()) # the head is trained on [value_min, value_max]; keep the display inside it values.append(min(max(v, float(norm.value_min)), float(norm.value_max))) fs = denormalize_state(future_state_raw[i].float(), norm, mode=NORM_MODE) fs = torch.nan_to_num(fs, nan=0.0, posinf=0.0, neginf=0.0) futures.append(_clamp(fs, norm.state_min, norm.state_max).view(-1).cpu().numpy()) policy_frames = _keyframes_uint8(imgs[0]) cf_frames = _keyframes_uint8(imgs[1]) video_policy = write_rollout_video(policy_frames, tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name) video_cf = write_rollout_video(cf_frames, tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name) filmstrip = build_filmstrip(policy_frames, cf_frames, mode) plot = build_action_plot( absolute_joints(plan, state[0], norm).cpu().numpy(), absolute_joints(edited, state[0], norm).cpu().numpy(), state[0].cpu().numpy(), futures[0], futures[1], mode, ) return video_policy, video_cf, verdict_markdown(values[0], values[1], mode), filmstrip, plot