"""FACT — Failure-Aware Causal Training for World-Action Models. Interactive counterfactual imagination: FACT plans a 3.2 s action chunk for a real RoboTwin observation, then imagines the future twice — once under its own plan and once under an action you edited — and scores both with its task-progress / failure head. """ import json import os from pathlib import Path os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1") import spaces # noqa: E402 (must precede torch/CUDA imports) import gradio as gr # noqa: E402 import torch # noqa: E402 from diffusers.models import AutoencoderKLWan # noqa: E402 from diffusers.schedulers import UniPCMultistepScheduler # noqa: E402 from huggingface_hub import snapshot_download # noqa: E402 from transformers import AutoTokenizer, UMT5EncoderModel # noqa: E402 import fact_demo # noqa: E402 from fact_demo import COUNTERFACTUALS # noqa: E402 from world_action_model.models.transformer_wa_casual import CasualWorldActionTransformer # noqa: E402 from world_action_model.pipeline import WAPipeline # noqa: E402 from world_action_model.pipeline.utils import ( # noqa: E402 extract_normalization_tensors, infer_robot_adapter_rank, load_stats, ) FACT_REPO = "Bariona/fact-wam" WAN_REPO = "Wan-AI/Wan2.2-TI2V-5B-Diffusers" DTYPE = torch.bfloat16 HERE = Path(__file__).parent # --------------------------------------------------------------------------- # Weights. Only the Wan VAE / UMT5 text encoder are needed from the base repo — # FACT replaces the 20 GB Wan transformer with its own 5B world-action transformer. # --------------------------------------------------------------------------- fact_dir = Path(snapshot_download(FACT_REPO)) wan_dir = Path( snapshot_download( WAN_REPO, allow_patterns=["model_index.json", "scheduler/*", "vae/*", "tokenizer/*", "text_encoder/*"], ) ) STATS = load_stats(str(fact_dir / "norm_stats_delta.json")) _transformer_dir = fact_dir / "transformer" _config = dict(CasualWorldActionTransformer.load_config(_transformer_dir)) _state_dict = torch.load(_transformer_dir / "diffusion_pytorch_model.bin", map_location="cpu", weights_only=True) _config["robot_adapter_rank"] = max(int(_config.get("robot_adapter_rank", 0) or 0), infer_robot_adapter_rank(_state_dict)) # Build straight into bf16 (the reference loads in fp32 then casts; same result, half the RAM). _default_dtype = torch.get_default_dtype() try: torch.set_default_dtype(DTYPE) transformer = CasualWorldActionTransformer.from_config(_config) finally: torch.set_default_dtype(_default_dtype) _load = transformer.load_state_dict(_state_dict, strict=False) if _load.missing_keys: print(f"[fact] missing transformer keys ({len(_load.missing_keys)}): {_load.missing_keys[:8]}", flush=True) if _load.unexpected_keys: print(f"[fact] unexpected transformer keys ({len(_load.unexpected_keys)}): {_load.unexpected_keys[:8]}", flush=True) del _state_dict transformer = transformer.to(DTYPE).eval() vae = AutoencoderKLWan.from_pretrained(wan_dir, subfolder="vae", torch_dtype=DTYPE) scheduler = UniPCMultistepScheduler.from_pretrained(wan_dir, subfolder="scheduler") _expand_timesteps = bool(json.loads((wan_dir / "model_index.json").read_text()).get("expand_timesteps", True)) pipe = WAPipeline( tokenizer=None, text_encoder=None, # text is encoded separately, see fact_demo.encode_instruction vae=vae, scheduler=scheduler, transformer=transformer, expand_timesteps=_expand_timesteps, ) for _name in ("scheduler", "action_scheduler", "future_state_scheduler", "value_scheduler"): _sched = getattr(pipe, _name, None) _updates = {k: fact_demo.FLOW_SHIFT for k in ("flow_shift", "shift") if k in _sched.config} if _updates: _sched.register_to_config(**_updates) pipe.to("cuda") tokenizer = AutoTokenizer.from_pretrained(wan_dir / "tokenizer") text_encoder = UMT5EncoderModel.from_pretrained(wan_dir / "text_encoder", torch_dtype=DTYPE) text_encoder = text_encoder.eval().requires_grad_(False).to("cuda") # --------------------------------------------------------------------------- # Bundled RoboTwin observations (real frames from Bariona/robotwin-v2, MIT licensed) # --------------------------------------------------------------------------- SCENARIOS = json.loads((HERE / "examples" / "scenarios.json").read_text()) def _views(scenario): base = HERE / "examples" / scenario["slug"] return [str(base / "cam_high.png"), str(base / "cam_left_wrist.png"), str(base / "cam_right_wrist.png")] def _state_text(scenario): return ", ".join(f"{v:.4f}" for v in scenario["state"]) DEFAULT = SCENARIOS[0] # Per-scenario counterfactual picked by measuring this checkpoint's value head — each of # these visibly moves FACT's predicted time-to-go on its scenario. SCENARIO_MODES = { "place_dual_shoes": "Freeze the arms", "stack_blocks_three": "Random jitter", "beat_block_hammer": "Swap left and right arm", "handover_block": "Open the grippers (let go)", "place_bread_basket": "Squeeze the grippers shut", } DEFAULT_MODE = SCENARIO_MODES.get(DEFAULT["slug"], list(COUNTERFACTUALS)[0]) def _mode_for(scenario): return SCENARIO_MODES.get(scenario["slug"], DEFAULT_MODE) EXAMPLES = [_views(s) + [s["instruction"], _state_text(s), _mode_for(s)] for s in SCENARIOS] EXAMPLE_LABELS = [f"{s['label']} → {_mode_for(s).lower()}" for s in SCENARIOS] @spaces.GPU(duration=25) def imagine(cam_high, cam_left, cam_right, instruction=DEFAULT["instruction"], state_text=_state_text(DEFAULT), mode=DEFAULT_MODE, strength=1.0, steps=fact_demo.NUM_STEPS, seed=0, progress=gr.Progress(track_tqdm=True)): models = fact_demo.Models( pipe=pipe, tokenizer=tokenizer, text_encoder=text_encoder, norm=extract_normalization_tensors( STATS, device=torch.device("cuda"), state_dim=fact_demo.STATE_DIM, action_dim=fact_demo.ACTION_DIM, ), device=torch.device("cuda"), dtype=DTYPE, ) try: return fact_demo.run(models, cam_high, cam_left, cam_right, instruction, state_text, mode, float(strength), int(steps), int(seed)) except ValueError as err: raise gr.Error(str(err)) CSS = """ #col-container { max-width: 1200px; margin: 0 auto; } .dark .gradio-container { color: var(--body-text-color); } """ with gr.Blocks(title="FACT world-action model") as demo: with gr.Column(elem_id="col-container"): gr.Markdown( """ # FACT · what does a robot world model think happens if you do the wrong thing? [**FACT**](https://huggingface.co/papers/2608.10232) is a *world-action model*: one causal diffusion transformer that jointly denoises **the next 48 actions**, **the future video** and **a task-progress value** for a bimanual robot. Failure-aware causal training is what makes the video branch actually *react* to the action it is conditioned on — so you can hand it a deliberately bad action and watch it imagine the consequences. Pick a RoboTwin situation below, choose a way to **sabotage the plan**, and FACT imagines both futures from the same observation and the same noise — the only difference is the action. [paper](https://huggingface.co/papers/2608.10232) · [project page](https://fact-wam.github.io) · [code](https://github.com/Bariona/FACT) · [checkpoint](https://huggingface.co/Bariona/fact-wam) · [RoboTwin 2.0 data](https://huggingface.co/datasets/Bariona/robotwin-v2) """ ) with gr.Row(): with gr.Column(scale=5): cam_high = gr.Image(label="Overhead camera", type="numpy", height=240, value=_views(DEFAULT)[0]) with gr.Row(): cam_left = gr.Image(label="Left wrist", type="numpy", height=150, value=_views(DEFAULT)[1]) cam_right = gr.Image(label="Right wrist", type="numpy", height=150, value=_views(DEFAULT)[2]) instruction = gr.Textbox(label="Language instruction", lines=2, value=DEFAULT["instruction"]) mode = gr.Dropdown(label="Counterfactual action — how to break the plan", choices=list(COUNTERFACTUALS), value=DEFAULT_MODE) strength = gr.Slider(0.0, 2.0, value=1.0, step=0.05, label="Edit strength (0 = FACT's own plan, 1 = full edit)") run = gr.Button("Imagine both futures", variant="primary") with gr.Accordion("Robot state & sampling", open=False): state_text = gr.Textbox( label="Joint state — 14 numbers: left arm j1..j6 + gripper, then right arm", lines=3, value=_state_text(DEFAULT), ) steps = gr.Slider(8, 40, value=fact_demo.NUM_STEPS, step=1, label="Denoising steps (20 = released eval setting)") seed = gr.Slider(0, 2**31 - 2, value=0, step=1, label="Seed") with gr.Column(scale=7): with gr.Row(): video_policy = gr.Video(label="FACT's own plan", autoplay=True, loop=True, height=260) video_cf = gr.Video(label="Counterfactual action", autoplay=True, loop=True, height=260) verdict = gr.Markdown() filmstrip = gr.Image(label="Predicted keyframes (now, +0.8 s, +1.6 s, +2.4 s, +3.2 s)", type="filepath", height=260) with gr.Accordion("Action chunk & predicted future joint state", open=False): action_plot = gr.Image(label="48-step action chunk", type="filepath") inputs = [cam_high, cam_left, cam_right, instruction, state_text, mode] outputs = [video_policy, video_cf, verdict, filmstrip, action_plot] run.click(imagine, inputs=inputs + [strength, steps, seed], outputs=outputs) gr.Examples( examples=EXAMPLES, example_labels=EXAMPLE_LABELS, inputs=inputs, outputs=outputs, fn=imagine, cache_examples=True, cache_mode="lazy", label="Real RoboTwin 2.0 observations (mid-episode frames from the authors' dataset)", ) gr.Markdown( """ ### Reading the output * **Videos / keyframes** — FACT predicts 4 future keyframes covering the next 3.2 s (48 control steps at 15 Hz); each is held over its segment, so the clip is a slideshow, not an interpolation. * **Value** = the fraction of the episode still remaining 3.2 s from now (0 = task finished). It is trained with a **+1 penalty on failure states**, so a value above 1 is FACT saying *"this rollout broke the task"*. The released policy uses exactly this head to rank candidate action chunks. On all five bundled observations FACT's *own-plan* value lands within `0.02` of the ground-truth time-to-go computed from the RoboTwin episode itself (e.g. `+0.477` predicted vs `0.484` actual) — a useful sanity check that the head is calibrated, not just plausible. * **Edit strength** interpolates between FACT's own action chunk and the sabotaged one, in de-normalised delta-action space, so you can watch the imagined future degrade continuously. Both rollouts share the same observation, instruction, noise and seed — only the action conditioning differs. Inference follows the released recipe: 384×192 three-view canvas, 48-action chunk, 5 frames, 20 UniPC steps, flow shift 3.0, no CFG, bf16. Model & code: [Bariona/FACT](https://github.com/Bariona/FACT), Apache-2.0. Example observations are frames from [Bariona/robotwin-v2](https://huggingface.co/datasets/Bariona/robotwin-v2) (MIT), built on [RoboTwin 2.0](https://robotwin-platform.github.io/). """ ) if __name__ == "__main__": # Gradio 6 takes theme/css on launch(), not on the Blocks constructor. demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)