Latest / Embodied AI 101 / Flex-π: Teach Geometry, Choose How Much to Imagine

Transcript

Flex-π: Teach Geometry, Choose How Much to Imagine · Embodied AI 101

• Flex-π: Teach Geometry, Choose How Much to Imagine. Flex-π: A Multi-Stream World-Action Model with Compute Flexibility • Ge Yan, Jinghao Liu, Yuzhi Fan, Lei Cai, Minwen Liao, Jesse Zhang, and Dieter Fox • University of Washington and Allen Institute for AI • arXiv preprint • 2026. A robot needs more than a convincing video. Consider a robot repairing its own gripper. It must retrieve a replacement part, seat it, insert a screw, pick up a powered screwdriver, and align the bit with the screw socket. In Flex-π’s demonstration, the tightest insertion leaves approximately a quarter-millimeter of lateral clearance. This is an unforgiving setting for a policy that understands what a successful repair should look like but cannot reliably organize the corresponding motion. That distinction motivates this episode of Embodied AI 101. Predicting a plausible visual future and producing a useful representation for control are related objectives, but they are not interchangeable. An attractive reconstruction does not tell us which errors matter to the robot. A small displacement around a contact interface may matter more than a much larger appearance error elsewhere. Flex-π addresses this problem by changing what a world-action model learns to predict. The six-billion-parameter model generates actions alongside future RGB latents, three-dimensional pointmap latents, and DINO semantic features. Its second contribution is deployment flexibility: a trained checkpoint can generate those visual futures, only some of them, or no visual futures at all. The intended benefit is richer training supervision without making expensive imagination mandatory at runtime. The important questions are therefore separable. Does structured prediction improve the learned policy? Does generating those predictions during deployment provide additional value? And how much computation does that additional value require?. Keeping those questions separate makes the paper more interesting—and prevents several misleading readings of its results. Changing the prediction target. The motivation should not be interpreted as “RGB models contain no geometry.” That would be too strong. The more defensible argument concerns the training objective: reconstructing appearance does not explicitly prioritize the spatial and semantic distinctions that manipulation depends on. Think about what we would ask a representation to preserve. For inserting a plate into a rack, appearance helps identify the plate and the intended slot. Geometry describes their relative arrangement. Semantic features may help distinguish the target object from distractors. These are overlapping descriptions, but they need not impose identical learning pressures. Flex-π sits at the intersection of several existing approaches. DreamZero, introduced in World Action Models are Zero-shot Policies, builds robot policies on a pretrained video diffusion backbone and jointly models video and action. Its results support the broader proposition that video-generation priors can transfer to manipulation and unfamiliar motions. That establishes an important backdrop: predicting visual futures is not merely an auxiliary visualization exercise. Fast-WAM asks a complementary question: must the policy actually generate those futures at test time? Its controlled variants retain video co-training while removing future generation during deployment. The reported finding is that removing video training hurts substantially more than removing runtime imagination. In other words, world prediction can be valuable because of what it teaches the weights, even when the deployment graph never produces a future video. DINO-WM explores a different axis. Rather than reconstructing images, it predicts pretrained DINO patch features and optimizes action sequences toward observational goals. It is relevant because it treats representation choice as central to world modeling. But its control mechanism—planning through an action-conditioned dynamics model—is different from Flex-π’s action-generation architecture. My reading is that Flex-π combines two useful ideas: make the prediction targets more relevant to manipulation, and disentangle the benefits of training with those targets from the cost of generating them during deployment. That is a more precise contribution than simply adding a third modality to a large transformer. Three descriptions of the same scene. Start with the geometry stream. A pointmap is an image-shaped field of three-dimensional coordinates. Instead of storing red, green, and blue at each pixel, it stores the position of the corresponding visible surface point. Flex-π’s implementation constructs these fields by unprojecting depth, normalizing the coordinates, and assembling the camera views into a composite layout. The representational choice matters. A pointmap retains pixel-to-surface correspondence and regular spatial organization. It is not, by itself, a complete volumetric scene model: occluded surfaces remain a different problem. Nor should an image-shaped coordinate field be confused with a persistent, globally registered map. The surprising observation is that Flex-π can encode pointmaps with the same frozen Wan video VAE used for RGB. The team does not train a new geometry-specific autoencoder. Instead, an existing compression model is reused for a different three-channel signal. DINO features enter separately; they are not passed through this VAE. This suggests a useful engineering hypothesis: before building a specialized tokenizer, test whether an existing one preserves the structure your downstream learner needs. A representation can be useful outside its original training distribution without becoming a universal or exact codec. The distinction is particularly important here. The project reports pointmap reconstruction at 31.1 decibels PSNR in normalized coordinates, alongside a depth-axis reconstruction RMSE of 4.9 centimeters in metric units. That is evidence of useful structural preservation, not proof of sub-millimeter geometric accuracy. It is also an autoencoder reconstruction measurement, not a measurement of future-prediction accuracy. For an adaptation, I would therefore inspect errors near contact surfaces, thin structures, and object boundaries—not merely average reconstruction quality. A visually convincing coordinate field could still erase precisely the local detail a manipulation task needs. The semantic stream comes from a frozen DINOv3 encoder. DINOv3 is designed to produce transferable dense visual features, and Flex-π uses patch-level outputs rather than only a global image embedding. “Object-centric semantics” should be understood in that dense-feature sense, not as a guarantee of explicit object identities, segmentation masks, or persistent object tracks. There is a practical token-budget trick. Flex-π concatenates each two-by-two neighborhood of DINO patches into one wider token. Four 768-dimensional vectors become one 3,072-dimensional vector. This is an invertible rearrangement, not average pooling: spatial feature content is retained while the number of tokens falls by four. Finally, where does depth come from? Depth Anything 3 provides an off-the-shelf route from RGB to geometry, and the released Flex-π workflow includes offline depth annotation for RGB-only datasets. Existing sensed depth can also be used. Thus, additional depth hardware is optional as a system capability; it does not follow that every reported experiment used only RGB observations. What actually communicates inside the model?. The backbone pairs a roughly five-billion-parameter visual transformer with a roughly one-billion-parameter action expert. The released RoboTwin model describes thirty layers, with cross-stream fusion in the middle sixteen and seven stream-specific layers on either side. This is a Mixture-of-Transformers arrangement, not three independently replicated five-billion-parameter visual models. The design separates parameter sharing from information sharing. Visual streams can reuse a large pretrained transformation while retaining modality-specific input and output adapters. The action expert has its own parameters and a narrower hidden representation. This lets the architecture specialize action processing without making every action token pay for another full-sized visual backbone. More important than the parameter count is the direction of attention. Actions read visual representations, including active predicted futures; visual tokens do not read the action stream. Consequently, future generation can influence the evolving action prediction, but this is not unrestricted bidirectional interaction between proposed actions and predicted consequences. That gives us a useful boundary on interpretation. I would not describe Flex-π as performing model-predictive control over alternative action sequences. The reported mechanism does not require sampling candidate controls, rolling each through an action-conditioned simulator, assigning values, and selecting a winner. Instead, the action generator can use a jointly generated, task-conditioned visual future as additional internal context. This is closer to learning an action-producing representation of anticipated task evolution than exposing a general-purpose counterfactual simulator. The contrast with DINO-WM is instructive. In DINO-WM, action-sequence optimization is explicitly part of reaching a visual goal. Flex-π’s runtime mechanism is direct generation. Both approaches use future representations, but “uses a world model” does not specify how decisions are made. This distinction suggests different research extensions. If the goal is fast imitation-based manipulation, Flex-π’s arrangement is attractive. If the goal is evaluating unfamiliar interventions, the next question is whether the visual predictions remain meaningful when conditioned on deliberately different actions. That capability would need its own architecture and evaluation; it should not be inferred from successful joint generation. The training trick: missing inputs are still prediction targets. The generative machinery uses flow matching. In plain language, training constructs intermediate states between noise and a target, then teaches the network which direction moves those states toward data. At inference, a numerical solver follows the learned direction field through several updates. This is the underlying framework, rather than a separate planning algorithm. Flex-π applies supervision to actions and the visual streams. One implementation detail is that the DINO head predicts clean features, with that prediction converted into the velocity representation used by the flow objective. This accommodates the unusually wide folded DINO tokens without changing the overall sampling framework. Compute flexibility comes from two independently sampled decisions. The first determines which current visual streams are available as observations. The second determines which future streams participate in joint generation with the action stream. The training configuration exposes separate probabilities for presence and joint participation, together with cross-modal prediction controls. The critical distinction is that the output selection is an attention mask, not a loss mask. During the flexible training scheme, a future stream can remain supervised even when actions do not read it. Likewise, removing a modality from the current observation does not necessarily remove its future prediction target. Consider a sample with RGB and DINO observations but no observed pointmap. The model must still learn to produce the missing modality’s future representation. In another sample, the missing observation could be semantic features rather than geometry. The authors call this cross-modality forcing; the released helper code randomizes stream presence and joint-generation regimes per sample. The conceptual benefit goes beyond surviving a missing sensor. If geometry prediction were possible only when geometry was already supplied, the model could learn a relatively isolated geometry-processing pathway. Requiring prediction from other observations creates pressure for the streams to share useful information. That is an interpretation of the learning constraint, not proof that the hidden representation has become geometrically consistent. A network might satisfy a prediction loss through several strategies. Nevertheless, it gives us a concrete, testable reason for the masking design. With three visual modalities, there are seven nonempty input subsets and eight output subsets, including no visual output. That yields fifty-six input/output combinations. The repository exposes these through separate runtime flags for observed and generated streams. At deployment, this distinction has real computational consequences. Selecting action-only execution means not generating visual futures—not generating them and discarding the decoded images afterward. The learned weights retain the effects of multimodal training, while the expensive future-generation path can be omitted. What was trained, and what “long horizon” means here. The training pipeline is not a small-data policy learned from scratch. The reported system uses pretrained visual components and approximately five hundred hours from one hundred AGIBOT World-Beta tasks before downstream adaptation. Demonstration efficiency therefore means efficiency in the downstream experiments, not absence of substantial prior training. The real-world policies are fine-tuned per task. Their deployment predicts thirty-two actions, executes them at thirty hertz, and then replans; the synchronous client holds position during the next inference call. Each chunk uses a current observation rather than a long observation history. These details change how we should read “long-horizon performance.” The robot can complete a long sequence through repeated short-horizon predictions. That is not the same claim as forecasting the whole repair procedure in one internally consistent rollout. Likewise, successful self-repair should not automatically be read as autonomous fault diagnosis. The demonstrated competence concerns executing a trained repair task. Identifying an unfamiliar failure, deciding that repair is necessary, and constructing a novel repair procedure would require different evidence. There is also a distinction between data efficiency and data appetite. A method may outperform a baseline at the same demonstration budget while still requiring a substantial collection effort. The authors explicitly acknowledge that the system remains data-hungry. For a practitioner, I would frame the result as a possible reduction in the marginal demonstration cost of a difficult task—not as a guarantee that a handful of demonstrations will suffice. The choice of baselines also matters. ManiFlow is a strong manipulation policy with multimodal conditioning and consistency-based flow training, whereas Fast-WAM studies video co-training without embodied pretraining. Their comparisons are useful system-level evidence, but they are not an experiment in which every prior, training corpus, and optimization choice is identical. The real-robot evidence: progress is not completion. The most persuasive real-world finding is that the fast deployment mode already performs well, while joint generation improves it further. But first, distinguish the metrics. The evaluation reports normalized task completion, which awards partial credit, and binary success, which requires finishing the task. These answer different questions. A policy that reliably grasps a zipper pull but cannot close the bag may accumulate useful partial credit without delivering the requested outcome. Across the five tasks, the reported average completion scores are approximately 83.0% for full-joint Flex-π and 76.4% for action-only Flex-π, compared with 58.0% for ManiFlow and 52.1% for π₀.₅. The corresponding mean binary success rates are 63%, 50%, 27%, and 18%. The advantage therefore appears in both progress and complete execution, rather than only in one favorable metric. The difference between the two Flex-π modes is informative. It supports two separate propositions: richer training can benefit an action-only policy, and generating futures can provide additional value on difficult tasks. Neither proposition requires the other to hold universally. The repair experiment provides a particularly clear endpoint: full-joint Flex-π completes all eight stages in eleven of twenty rollouts, while the strongest evaluated baseline completes one. The result is substantial, but twenty trials are still twenty trials; one additional success changes the measured rate by five percentage points. I would treat that as strong motivation for replication, not a deployment reliability guarantee. It also does not establish that the model estimates all relevant poses to quarter-millimeter accuracy. Mechanical clearance, visual estimation error, policy compensation, and closed-loop task success are different quantities. The unseen-bag condition is another useful test. Full-joint Flex-π retains 63.3% task completion, compared with 17.2% for π₀.₅ and 6.9% for ManiFlow. In a separate plate-placement data-reduction experiment, joint generation with half the demonstrations outperforms the studied baselines trained on the full dataset. These are encouraging but task-specific findings, not universal sample-efficiency multipliers. Simulation: the low-data results matter more than the leaderboard ceiling. On the full RoboTwin evaluation, the flexible checkpoint averages 94.6% success in both action-only and full-joint modes. The model card reports training across fifty tasks using 2,500 clean and 25,000 randomized demonstrations. At this data scale, generating futures barely changes the aggregate result. That does not invalidate imagination. It tells us that its value is conditional. A benchmark on which the cheaper path already succeeds almost everywhere has limited room to reveal the benefit of a more expensive representation. The lower-data comparison is more discriminating: at fifty demonstrations per task, full-joint Flex-π reaches 78.8%, versus 41.9% for Fast-WAM and 31.4% for π₀.₅. LIBERO presents another important distinction. The flexible checkpoint reports 98.4% action-only and 98.5% full-joint success across the four evaluation suites. Those are the relevant numbers for the claim that one checkpoint supports both operating modes. The 99.2% headline belongs to the separately fine-tuned, fixed full-joint variant, marked with an asterisk. That checkpoint was trained without stream dropout. Its model card explicitly warns that missing-modality and action-only deployment are not its intended operating regimes. One should not combine its best score with the flexible checkpoint’s deployment properties. LIBERO-Plus requires protocol care as well. Flex-π reports 88.6% after correcting instruction extraction; π₀.₅ and Fast-WAM are evaluated through the same corrected pipeline, while other comparisons retain published numbers. That limits how cleanly the entire leaderboard can be interpreted. The ablations provide a more direct view of the design. In a five-task, low-data experiment with RGB-only input, changing the generated outputs moves success from 40.2% action-only to 63.8% with all visual futures. The practical interpretation is not that more generation always wins. It is that training-time supervision and inference-time generation are separate experimental factors. For a stronger causal account, I would want matched comparisons that hold observation inputs, pretrained initialization, parameter budget, and training schedule fixed while varying individual prediction losses. Then I would independently vary runtime generation. That would help distinguish improvements from additional perceptual information, auxiliary supervision, cross-modal reconstruction, and extra inference computation. The existing evidence supports the combined recipe. It does not completely allocate credit among all its ingredients. The compute trade-off is real—and heavily engineered. The deployment numbers are approximately sixty milliseconds per action-only call and 193 milliseconds per full-joint call on an RTX 5090, using four denoising steps. These are optimized operating points, not raw eager-PyTorch timings. The released performance ladder makes that clear. Eager execution is about 132 milliseconds action-only and 447 milliseconds full-joint. Compilation brings action-only execution to sixty milliseconds, while full-joint execution remains around 360 milliseconds until the TensorRT optimization stack is applied. The fastest joint path uses separate prefill and decoding engines. This is useful systems work, not an embarrassment to be discounted. But it changes what the result establishes. The paper does not equalize optimization effort across every baseline. A laboratory should therefore compare deployed stacks, not infer that a model with a particular parameter count will automatically reproduce these timings. The architecture, token layout, compilation boundary, solver steps, and hardware all participate in the result. Also keep action-only separate from RGB-only. The former specifies generated outputs. The latter specifies observed inputs. The serving interface controls these independently, and its default fast-path invocation can retain multiple observation streams while disabling all visual futures. Nor is sixty milliseconds equivalent to a demonstrated sixteen-hertz visual feedback loop. A per-call latency must be interpreted alongside the number of actions executed between calls. For a new application, the relevant quantity is how quickly a new observation can change the executed behavior. Memory introduces another trade-off. The documented YAM configurations use roughly 15.8 gigabytes without TensorRT engines and 25.7 gigabytes for the fastest full-joint configuration. The repository’s full-training guidance calls for multiple eighty-gigabyte GPUs. Cheap inference relative to competing large policies does not make training or adaptation computationally trivial. How I would reproduce and adapt it. As of September 28, 2026, the public repository contains training, evaluation, and deployment code, with RoboTwin and LIBERO checkpoints and released datasets. However, the documentation states that the AGIBOT-pretrained checkpoint used in the reported pipeline is not yet released; shipped task configurations begin from base Wan initialization unless a warm-start checkpoint is supplied. That distinction is essential when planning a reproduction. I would first reproduce a published checkpoint before launching a new training run. This separates installation and evaluation failures from training-recipe differences. A successful forward pass is not enough: verify the intended input and output regime, action conventions, camera arrangement, and normalization. The released workflow expects specific dataset organization, depth streams and camera intrinsics when geometry is enabled, and cached text embeddings. Model cards also emphasize preserving the saved configuration and normalization statistics alongside the weights. These artifacts are part of the executable policy, not incidental bookkeeping. For new hardware, my first substantive experiment would be a small factorial study. Train with and without structured future supervision, then evaluate both cheap and expensive deployment modes under the same control schedule. Include geometry-sensitive tasks and appearance-heavy distractor conditions. The goal is to discover where the additional targets help your distribution, rather than reproduce only an aggregate advantage. Next, inspect the representation itself. Check camera-coordinate conventions, depth units, intrinsics after resizing, invalid-depth handling, and temporal alignment. I would prioritize contact-relevant reconstruction errors over pretty decoded videos. This is especially important if geometry targets come from an estimator rather than a calibrated depth source. The deployment sweep should also include wall-clock task duration, not just success and model latency. A slower policy can finish sooner if it avoids retries. Conversely, higher average success may not compensate for long pauses in an application requiring prompt reactions. Neither trade-off is captured by milliseconds per model call alone. An attractive extension is state-dependent compute selection: cheap action generation during easy transport, richer future generation near difficult contacts. The released runtime controls make such an experiment possible, but the mechanism for choosing a mode would itself require validation. Fixed-mode benchmark success does not establish that switching modes online is stable. Finally, I would test whether disagreement among predicted modalities is useful. If appearance suggests successful placement while geometry implies a collision, could that discrepancy trigger a new observation or recovery? This is a research proposal, not a capability established by the paper. Correlated prediction errors may make apparent agreement uninformative. The general lesson is to use the extra streams as testable representations—not automatically as trustworthy explanations of the policy. The takeaway: better supervision, optional imagination. My assessment is that Flex-π makes a meaningful architectural contribution, but its strongest idea is not simply the presence of three-dimensional outputs. It is the combination of structured future supervision and a deployment mechanism that makes those futures optional. That combination turns the question “Should a robot imagine before acting?” into a more useful engineering question: “What should the robot learn to predict, and when is generating those predictions worth the cost?”. The evidence is strongest where the problem is difficult enough to expose differences: reduced demonstration budgets, precise multi-stage manipulation, and selected distribution shifts. The limitations are equally important: task-specific adaptation, substantial prior training, finite real-robot evaluation, and a joint-generation pathway that remains materially more expensive. For a roboticist, the transferable lesson is to reconsider the targets used to shape a policy’s representation. Appearance, geometry, and semantics can provide different supervision—even when they originate from the same images. Then measure, rather than assume, how much of that learned structure needs to be generated again at runtime.