yixiangchen515/flowwam

FlowWAM is the official code release for a research paper about controlling robots using a technique called optical flow, and it is written in Python. The broader idea it builds on is called a World Action Model, which repurposes an AI video generator, a model normally trained to predict what happens next in a video, so that it can also predict robot actions. The problem the paper addresses is that action commands for a robot, such as joint movements, do not naturally look like the pixels a video model was trained on, creating a mismatch between the two. FlowWAM's solution is to represent robot actions using optical flow instead, which is a per pixel map of how things visually move from one video frame to the next. Because optical flow can be extracted directly from ordinary video, even video that was never labeled with robot action data, it fits naturally into the same visual format the underlying video generator already understands. A single shared model, built as what the authors call a dual stream diffusion model, both predicts this flow as the action to take and uses flow as an input when imagining what the world will look like next, and it can be pretrained on large amounts of unlabeled video before being fine tuned for a specific robot task. The repository includes training code, an inference server that a robot control script can talk to over a network connection, and a data generation tool that renders a robot-only view of a scene for computing flow. It is built to plug into an existing robot simulation benchmark called RoboTwin as an add on policy, without needing to edit RoboTwin's own source code, and it ships scripts to evaluate performance across RoboTwin's 50 built in tasks, including support for running multiple evaluations in parallel across several graphics cards. Setting it up requires a GPU, a specific version of PyTorch matched to a particular CUDA version, a separate RoboTwin simulation environment installed alongside it, and downloading multi gigabyte pretrained video model weights along with the project's own released checkpoints and training dataset from Hugging Face.