Latent Policy Steering with Embodiment-Agnostic Pretrained World Models

01 动机Motivation

收集特定机器人形态的大规模演示数据成本高昂，而现有跨形态数据集因动作空间不一致、形态差异大，难以被直接利用。 如何从已有多形态数据中提炼可迁移的"技能先验"，以极少量目标形态演示即可大幅提升策略性能，是本文的核心出发点。Collecting large-scale demonstration data for one specific robot embodiment is expensive, while existing cross-embodiment datasets are hard to use directly because their action spaces are inconsistent and their embodiments differ greatly. How to distill a transferable “skill prior” from existing multi-embodiment data, so that a very small number of target-embodiment demonstrations already improves policy performance substantially, is the core starting point of this paper.

"Skills performed across different embodiments produce visual similarities in motions that can be captured using off-the-shelf action representations such as optical flow."

02 方法Method

LPS 分三步：① 以光流为动作表征在多形态数据上预训练视觉世界模型； ② 用少量目标形态演示微调世界模型并训练 Robust Value Function； ③ 推理时对 diffusion policy 采样的多条动作规划进行隐空间评估，选取价值最高的计划执行。LPS proceeds in three steps: ① pretrain a visual world model on multi-embodiment data with optical flow as the action representation; ② fine-tune the world model with a few target-embodiment demonstrations and train a Robust Value Function; ③ at inference time, evaluate in latent space the multiple action plans sampled from the diffusion policy and execute the plan with the highest value.

Flow-as-Action：形态无关的世界模型预训练Flow-as-Action: embodiment-agnostic world model pretraining

传统世界模型以具体机器人动作（如关节角度或末端位姿）为条件，无法直接跨形态迁移。 LPS 将光流替换为动作输入：卷积编码器将光流场压缩为 n 维向量（n 等于目标形态的动作空间维度）， 迫使网络捕获与形态无关的显著运动特征，同时抑制噪声和形态差异。 该设计使预训练数据可来自任意机器人甚至人类演示。Conventional world models are conditioned on concrete robot actions (such as joint angles or end-effector poses) and therefore cannot transfer across embodiments directly. LPS replaces the action input with optical flow: a convolutional encoder compresses the flow field into an n-dimensional vector (n equals the action-space dimension of the target embodiment), forcing the network to capture salient embodiment-agnostic motion features while suppressing noise and embodiment differences. This design lets the pretraining data come from any robot, or even from human demonstrations.

目标形态微调与 Robust Value FunctionTarget-embodiment fine-tuning and the Robust Value Function

在获得少量目标形态（如 Franka）演示后，将光流编码器替换为归一化的机器人真实动作， 同时联合训练世界模型与价值函数。 价值函数需处理 distribution shift 问题——推理时策略访问的状态可能偏离训练分布。 为此，作者设计了Robust Value Function，同时在专家演示状态和策略访问状态上训练， 并引入 cosine similarity 奖励惩罚偏离专家轨迹的行为：Once a small number of target-embodiment (e.g. Franka) demonstrations are available, the optical flow encoder is replaced by normalized real robot actions, and the world model and the value function are trained jointly. The value function must cope with the distribution shift problem — at inference time the states visited by the policy may deviate from the training distribution. To this end the authors design a Robust Value Function that is trained on both expert demonstration states and policy-visited states, and introduce a cosine similarity reward that penalizes behavior deviating from expert trajectories:

"r_t:t+h′ = r_t:t+h + (sim(s_t:t+h, s_t:t+h′) − 1) / 2"

其中 s_t:t+h 为专家隐状态序列，s_t:t+h′ 为策略在世界模型中展开的状态。 该奖励促使价值函数在 out-of-distribution 状态下仍能给出保守估计，避免高估偏差。Here s_t:t+h is the expert latent state sequence and s_t:t+h′ is the sequence rolled out by the policy inside the world model. This reward makes the value function stay conservative in out-of-distribution states, avoiding overestimation bias.

推理时隐空间策略引导（Latent Policy Steering）Inference-time steering in latent space (Latent Policy Steering)

给定当前观测，从 diffusion policy 采样 B 条候选动作规划，规划长度为 h。 世界模型将每条规划在隐空间中前向展开，得到未来隐状态序列， 再由价值函数以加权平均（未来时刻权重更大）计算规划级价值。 执行价值最高的规划前 1 步后重新规划，实现滚动决策。Given the current observation, B candidate action plans of length h are sampled from the diffusion policy. The world model rolls each plan forward in latent space to obtain the future latent state sequence, and the value function then computes a plan-level value as a weighted average (later timesteps carry larger weights). The first 1 step of the highest-value plan is executed before replanning, yielding receding-horizon decision making.

03 实验Experiments

实验在两个场景下评估：① Robomimic 仿真（Lift / Can / Square / Transport）， 预训练使用 IIWA、UR5e、Kinova3 三种形态，目标形态为 Franka； ② 真实 Franka 机器人上的 4 个操作任务。 基线为行为克隆（BC），对比方法为 LPS-mix（使用混合多形态预训练数据）。Experiments are evaluated in two settings: ① Robomimic simulation (Lift / Can / Square / Transport), pretrained on the three embodiments IIWA, UR5e and Kinova3 with Franka as the target embodiment; ② 4 manipulation tasks on a real Franka robot. The baseline is behavior cloning (BC) and the compared method is LPS-mix (using the mixed multi-embodiment pretraining data).

真实世界实验（Table I）Real-world experiments (Table I)

在 30–50 条演示下，LPS-mix 相对 BC 实现 70% 相对提升（21.2% → 36.2%）； 60–100 条演示下实现 44% 相对提升（51.2% → 73.8%）。With 30–50 demonstrations, LPS-mix attains a 70% relative improvement over BC (21.2% → 36.2%); with 60–100 demonstrations it attains a 44% relative improvement (51.2% → 73.8%).

Robomimic 仿真实验（Table II，50 条 Franka 演示）Robomimic simulation experiments (Table II, 50 Franka demonstrations)

四个任务平均实现 10.6% 相对提升，预训练数据大部分来自真实世界场景而非仿真。Averaged over the four tasks it attains a 10.6% relative improvement, even though most of the pretraining data comes from real-world scenes rather than simulation.

消融实验Ablations

光流 vs. EEF 动作（Table III）： 以光流预训练的世界模型（LPS-sim flow，均值 62.4%）优于以末端位姿（EEF）预训练版本（59.1%）， 在 Square 任务上差距最显著（52.4% vs. 45.3%），验证了光流作为形态无关表征的优越性。Optical flow vs. EEF actions (Table III): the world model pretrained with optical flow (LPS-sim flow, mean 62.4%) outperforms the version pretrained with end-effector poses (EEF) (59.1%), with the largest gap on the Square task (52.4% vs. 45.3%), confirming the superiority of optical flow as an embodiment-agnostic representation.

价值函数设计（Table IV，100 条演示无预训练）： 完整 LPS（均值 68.7%）优于 LPS-vanilla（65.2%）和 LPS-bootstrap（64.3%）， 而 BC 基线为 62.9%。分布偏移惩罚项对最终性能贡献最大。Value function design (Table IV, 100 demonstrations, no pretraining): the full LPS (mean 68.7%) outperforms LPS-vanilla (65.2%) and LPS-bootstrap (64.3%), while the BC baseline reaches 62.9%. The distribution shift penalty term contributes the most to the final performance.

04 局限性Limitations