Walkthrough video (76 s, with narration)
Close Oven · unseen kitchen
From ten demonstrations, a coding agent writes the keypoints and short programs that turn an episodic memory into a policy, and revises them from its own play rollouts, without retraining.
A few demonstrations often contain the motion a robot needs in a new scene; what fails is the policy's notion of which parts of the scene matter. Correcting this in a learned policy means collecting more data and retraining, and even then it is hard to see what changed. We present KeyFlow, in which a coding agent adapts a policy by revising how a small set of demonstrations is represented and reused. The agent describes the scene with a few semantic 3D keypoints, which a vision-language model grounds in each new scene, and writes short programs that define the state, a distance between states, and how retrieved demonstration motion is carried into the current scene. Retrieval from this episodic memory defines a flow field that the end effector follows. The agent runs the policy on simulated play scenes, inspects keypoint-annotated rollouts, and revises the keypoints and programs, without new demonstrations or weight updates. On 13 RoboCasa manipulation tasks covering articulated doors, drawers, sliding racks and appliance controls, KeyFlow outperforms a vision-language-action model fine-tuned on ten times as many demonstrations, and agentic refinement improves held-out success on nearly every task. In IsaacSim, KeyFlow flow fields also guide a base π0.5 policy without retraining it, raising its mean success across six tasks by 37.2 percentage points.
Method
import numpy as np from loop import flowlib as F H = 16 STALL_WINDOW = 24 STALL_EPS = 0.015 JUMP = 24 PUSH_MOVED = 0.03 WIN = 8 def metric(keypoints): return float(np.linalg.norm(keypoints['door_top_edge'] - keypoints['oven_frame_top'])) def state(ep, t): e = ep.eef[t] d = ep.keypoints['door_top_edge'][t] f = ep.keypoints['oven_frame_top'][t] side = np.where(ep.keypoints['oven_frame_top'][0][1] > 0.0, 1.0, -1.0) return np.concatenate([e - d, d - f, e - f, np.array([side])]) def distance(S, s): w = np.array([1.0, 1.0, 1.0, 1.5, 1.5, 1.5, 0.5, 0.5, 0.5, 100.0]) return np.sqrt(np.sum(np.square(S - s) * w, axis=1)) def flow(ep, t, lib, j): jj = j d = ep.keypoints['door_top_edge'][t] d0 = ep.keypoints['door_top_edge'][0] if t > WIN and float(np.linalg.norm(d - d0)) > PUSH_MOVED: A = lib.actions n = len(A) err = np.zeros(n) for i in range(WIN): shifted = np.concatenate([np.full((i, 12), 9.0), A[0:n - i]], axis=0) err = err + np.linalg.norm(shifted - ep.action[t - 1 - i], axis=1) r = int(np.where(err == np.min(err))[0][0]) if float(err[r]) < 1e-4: demo_ids = np.asarray(lib.index)[:, 0] prev_demo = int(demo_ids[r]) e = ep.eef[t] f = ep.keypoints['oven_frame_top'][t] side = np.where(ep.keypoints['oven_frame_top'][0][1] > 0.0, 1.0, -1.0) s = np.concatenate([e - d, d - f, e - f, np.array([side])]) w = np.array([1.0, 1.0, 1.0, 1.5, 1.5, 1.5, 0.5, 0.5, 0.5, 100.0]) dist = np.sqrt(np.sum(np.square(lib.S - s) * w, axis=1)) dist = dist + np.where(demo_ids == prev_demo, 0.0, 1e6) jj = int(np.where(dist == np.min(dist))[0][0]) ei = int(lib.index[jj][0]) fi = int(lib.index[jj][1]) k = jj if t >= STALL_WINDOW: moved = float(np.linalg.norm(ep.eef[t] - ep.eef[t - STALL_WINDOW])) if moved < STALL_EPS: if jj + JUMP < len(lib.index): if int(lib.index[jj + JUMP][0]) == ei and int(lib.index[jj + JUMP][1]) == fi + JUMP: k = jj + JUMP return F.chunk(lib, k, H)
F.chunk replays the next H recorded demonstration actions from the retrieved memory; the agent added the push-phase lock (follow the same demonstration once the door has moved) and the stall jump after reading its play rollouts.The agent names a few scene keypoints with text descriptions, such as “midpoint of the free top edge of the open oven door”. A vision-language model grounds each description in several camera views of the first frame, and the robot's own keypoints come from forward kinematics. The same descriptions label every demonstration, so demonstrations and new scenes share one state space.
door_top_edge: “Midpoint of the free top edge of the open oven door (the edge farthest from the hinge, where the door handle bar is), on the door itself, not the glass window and not the rack.”oven_frame_top: “Midpoint of the top edge of the oven's front opening frame (the fixed oven body just above the cavity, where the closed door's top edge would rest).”robot_eef) come from forward kinematics.At each chunk boundary the policy retrieves the stored moment closest to the current state under the agent's distance, replays the next motion that demonstration made, optionally pulled toward where the demonstration would be in the current scene, and retrieves again. Over all states this defines a flow field whose streamlines are the demonstrations.
A persistent coding-agent session runs each version of the policy on a fixed set of simulated play scenes and reads the outcomes, keypoint-annotated videos and retrieval distances over time. It then edits the keypoints and programs. The best version on play is frozen and evaluated once on unseen kitchens.
Results
Held-out rollouts of the selected KeyFlow policy in kitchens it never saw during refinement (sped up 2×). Each clip is a successful episode; the numbers give success over all 50 held-out scenes.
Composition with a foundation model
The same agent loop produces flow fields that steer a base π0.5 policy, with no retraining of the base model. Across six tasks and 30 held-out scenes per task, the composed policy improves mean success by 37.2 percentage points over base π0.5 and beats it on every task, with the largest gains on Toaster Press (+86.7) and Block Stack (+53.3).
Every clip is the composed policy, base π0.5 guided by KeyFlow, running in IsaacSim, unedited and not sped up. Pink dots mark the grounded scene keypoints where shown. Red labels mark initial conditions or scene layouts that never appear in the demonstrations.