PaperlayerView exact revision on arXiv

Source-linked AI summary

Any-point Trajectory Modeling for Policy Learning

Chuan Wen, Xingyu Lin, John So, Kai Chen, Qi Dou, Yang Gao, Pieter Abbeel

arXiv:2401.00025v3cs.ROcs.CV

TL;DR

Robot policy learning needs action-labeled demonstrations, but collecting them is costly, while unlabeled videos lack direct action supervision. ATM pre-trains a model to predict future trajectories of arbitrary image points and uses those trajectories to guide policy learning; across more than 130 tasks, it reports stronger performance than video-pre-training baselines and transfer from human and other-robot videos.

  • Problem

    Action-labeled demonstrations are expensive to collect, while unlabeled videos are difficult to use for policy learning because they lack action labels.

  • Method

    ATM pre-trains a track transformer on action-free videos to predict future point trajectories, then trains a track-guided policy with limited action-labeled demonstrations.

  • Results

    Across over 130 language-conditioned simulation and real-world tasks, ATM outperforms video-pre-training baselines and learns from human and robot videos across embodiments.

  • Takeaways & Limitations

    Particle trajectories provide an interpretable, structured representation that incorporates physical inductive biases such as object permanence for video pre-training.

  • Takeaways & Limitations

    ATM still requires action-labeled demonstrations, and its video dataset contains only small domain gaps.

Abstract

from arXiv · show

Learning from demonstration is a powerful method for teaching robots new skills, and having more demonstration data often improves policy learning. However, the high cost of collecting demonstration data is a significant bottleneck. Videos, as a rich data source, contain knowledge of behaviors, physics, and semantics, but extracting control-specific information from them is challenging due to the lack of action labels. In this work, we introduce a novel framework, Any-point Trajectory Modeling (ATM), that utilizes video demonstrations by pre-training a trajectory model to predict future trajectories of arbitrary points within a video frame. Once trained, these trajectories provide detailed control guidance, enabling the learning of robust visuomotor policies with minimal action-labeled data. Across over 130 language-conditioned tasks we evaluated in both simulation and the real world, ATM outperforms strong video pre-training baselines by 80% on average. Furthermore, we show effective transfer learning of manipulation skills from human videos and videos from a different robot morphology. Visualizations and code are available at: \url{https://xingyu-lin.github.io/atm}.

I. INTRODUCTION

Robot learning benefits from more demonstrations, but action-labeled data are expensive to collect. ATM addresses the difficulty of using unlabeled videos by modeling structured future point trajectories that guide policy learning.

  • Motivation: Collecting action-labeled human demonstrations is time-consuming and labor-intensive, creating a major bottleneck for robot learning.One cited effort required 17 months to collect 130K trajectories.
  • Motivation: Videos encode behaviors, physics, and semantics, but their missing action labels make them difficult to use directly for policy learning.Feature-based video pre-training also largely neglects transition dynamics and future states.
  • Contribution: ATM represents video states as points and predicts their future 2D camera-frame trajectories to bridge video pre-training and policy learning.The representation is intended to reduce assumptions about camera calibration and transfer across domains and tasks.
  • Contribution: Predicted trajectories provide subgoals for learning robust control policies from a minimal amount of action-labeled data.The framework uses action-free videos for pre-training and limited demonstrations for trajectory-guided policy learning.
  • Contribution: ATM is presented as a structured particle-trajectory framework that outperforms video-pre-training baselines and supports learning from human and robot videos across embodiments.The paper reports extensive simulated and real-world experiments and cross-embodiment transfer.
  • Preliminary setup: The approach uses tracked 2D point coordinates and visibility values as camera-frame trajectory data for policy learning.Tracking is performed without external supervision using an off-the-shelf video tracker.

IV. METHOD

ATM is a two-stage framework: it first learns future point trajectories from large action-free video datasets, then uses those predictions to guide policy learning with few action-labeled demonstrations.

  • IV. METHOD: ATM first pre-trains a model to predict future point trajectories from action-free videos.The trajectory model is learned from an action-free video dataset before policy training.
  • IV. METHOD: ATM then trains a track-guided policy using predicted trajectories and a limited action-labeled trajectory dataset.The predicted tracks serve as guidance for mapping visual observations to control actions.

A. Trajectory Modeling from Video Datasets

ATM pre-trains a track transformer to forecast future positions of arbitrary image points, using tracker-generated trajectories from action-free videos as supervision.

  • Multimodal Track Modeling: The trajectory model predicts future coordinates for queried 2D image points over a horizon H in the camera frame.The model conditions on the current image observation, query-point positions, and language instruction.
  • Self-supervised Track Annotation: A vision tracker generates point trajectories from sampled frames in action-free videos for self-supervised pre-training.Points are sampled at a random timestep and tracked through the video.
  • Framework overview: The first framework stage combines tracker-generated trajectories with a transformer that predicts masked future point positions.The overview depicts trajectory extraction followed by trajectory prediction from current observations.
  • Self-supervised Track Annotation: Static points are filtered by tracking an n × n grid and thresholding the variance of point positions over time.Points are then resampled around locations that exhibited movement.
  • Multimodal Track Modeling: The framework masks future point positions and predicts them from current positions, image observations, and language instructions.This formulates future forecasting as a multimodal masked prediction problem.
  • Multimodal Track Modeling: Track-transformer pre-training learns a motion prior for particles within video frames.The resulting representation models temporal motion rather than only current visual state.

B. Track-guided Policy Learning

After pre-training, ATM feeds predicted tracks alongside current observations into a policy trained on limited demonstrations, using trajectories as dense subgoal guidance for action prediction.

  • Arbitrary Points Tracking: ATM uses a fixed grid of 32 points during policy inference instead of computing movement-based track filtering.The fixed-point sampling avoids the complexity of learning or identifying task-specific key points.
  • Track-guided Policy Learning: The policy takes the current observation and predicted future tracks as inputs and outputs control actions.Image observations remain in the policy so contextual information is not discarded.
  • Track-guided Policy Learning: Track tokens are fused with image tokens both before and after the policy transformer.These are called early fusion and late fusion, respectively.
  • Track-guided Policy Learning: The predicted tracks can serve as fine-grained subgoals, reducing policy learning to a subgoal-following problem.The method finds that language instruction is no longer needed during this policy stage for task specification.
  • Experiments: The experiments compare ATM with video-pre-training and behavioral-cloning baselines on more than 130 language-conditioned manipulation tasks.Evaluations span simulation and the real world, with additional experiments on human videos and ablations.

A. Video Pre-training for Imitation Learning

ATM bridges video pre-training and policy learning by using structured point trajectories as future guidance for language-conditioned manipulation. It outperforms video pre-training baselines across LIBERO tasks and also improves diffusion policies and real-robot performance.

  • Benchmark setup: LIBERO evaluates language-conditioned manipulation across suites targeting spatial reasoning, object reasoning, task understanding, and long-horizon behavior, including LIBERO-90's 90 tasks.Each standard suite contains 10 tasks, while LIBERO-90 contains 90 tasks.
  • Simulation results: ATM outperforms video pre-training baselines on all LIBERO tasks and excels on LIBERO-Goal and LIBERO-Long.The comparison includes baselines with fast computation on LIBERO-90.
  • Method: ATM bridges video data and policy learning with a structured representation of point trajectories.The main results evaluate this approach against video pre-training baselines.
  • Real-world results: ATM shows significant average improvements over state-of-the-art video pre-training baselines in real-robot experiments involving five dining-table tasks.The policy execution closely follows the predicted particle trajectories.
  • Mechanism: ATM's closed-loop future trajectory provides changing guidance about what the policy should do next, unlike static language instructions.The trajectories are used as future subgoals based on the current observation.
  • Policy integration: ATM consistently improves Diffusion Policy across all benchmark suites when predicted future tracks are added as conditioning.The ATM Diffusion Policy keeps the remaining setup the same as the base diffusion policy.

B. Human-to-robot and Robot-to-robot Transfer

ATM uses low-level any-point trajectories to learn manipulation skills from human videos and videos of robots with different embodiments. Additional cross-embodiment videos improve trajectory prediction robustness and policy learning when action-labeled demonstrations are limited.

  • Transfer rationale: ATM enables learning from cross-embodiment videos of humans or different robots by modeling low-level any-point trajectories.This is presented as a way to use more scalable data sources.
  • Human-to-robot transfer: Human-video experiments collect 100 human videos and 10 teleoperation demonstrations across three tasks.Comparisons include ATM trained without human videos and ATM trained with them.
  • Transfer results: Training the trajectory model with additional cross-embodiment videos makes trajectory prediction more robust and accurate, significantly improving policy learning.The experiments compare behavioral cloning, ATM trained on limited robot data, and ATM trained on human and robot data.
  • Human-to-robot transfer: ATM trained with human videos significantly outperforms behavioral cloning and ATM trained with only 10 robot videos in average success rate.The result demonstrates ATM's cross-embodiment capability.

C. Ablation Analysis

The ablations examine data efficiency, trajectory horizon, image masking, and track fusion. ATM benefits from limited action-labeled data and structured track inputs, with performance depending on horizon and architectural choices.

  • Effect of the number of action-labeled trajectories: 4% action-labeled demonstrations perform comparably to BC with 20% on LIBERO-Spatial, Object, and GOAL, and better on LIBERO-Spatial.
  • Effect of trajectory prediction horizon: A trajectory length of 16 yields optimal average performance, while length 0 reduces ATM to the behavior cloning baseline.
  • Effect of trajectory prediction horizon: Longer trajectories generally improve success rates, but gains tend to plateau after 16 steps.
  • Effect of image masking: Omitting image masking causes a slight policy-performance decline except on LIBERO-Spatial.The authors hypothesize that masking encourages joint reasoning across image regions, which may matter less when language specifies the spatial target.
  • Effect of early and late fusion: Removing late fusion produces the largest performance drop, supporting predicted tracks as useful subgoals in multi-task learning.The predicted tracks are supplied both before and after the policy transformer through early and late fusion.

VI. LIMITATIONS

The method still requires action-labeled demonstrations to map trajectories to actions, and the video data cover only small domain gaps. These constraints limit policy generalization and leave in-the-wild video as future work.

  • ATM still relies on action-labeled demonstration trajectories for mapping to actions, limiting generalization of the learned policies.
  • The video dataset contains only small domain gaps, while in-the-wild videos introduce multimodal distributions, diverse camera motions, and sub-optimal motions.

APPENDIX A ADDITIONAL EXPERIMENTAL RESULTS

Additional experiments compare ATM with UniPi variants, assess computation, and integrate ATM with Diffusion Policy. The results favor structured trajectory subgoals and show that ATM can improve another policy architecture.

  • Attention map visualization: ATM’s attention maps focus on task-relevant regions, while BC often attends to irrelevant backgrounds.Examples include attention to the cream cheese, bowl, and wine bottle specified by textual instructions.
  • Discussions on the UniPi Baselines: ATM uses structured point trajectories as subgoals, whereas UniPi predicts future image frames with a language-conditioned video diffusion model.
  • Computation: ATM is the most computationally efficient method and is feasible for closed-loop control, unlike the computationally prohibitive UniPi-Replan.
  • Discussions on the UniPi Baselines: UniPi-Replan performs worse than UniPi, suggesting that structured subgoals are more effective than image subgoals.Predicting an image goal at a fixed future time can produce inconsistent subgoals.
  • Applying ATM to Diffusion Policy: The ATM Diffusion Policy consistently outperforms the standard Diffusion Policy on LIBERO with limited training data.

D. Human-to-robot Transfer Details

Human-to-robot transfer experiments evaluate deformable-object, long-horizon, and tool-use tasks using action-free human videos alongside limited robot demonstrations. Attention visualizations compare how human-video pretraining affects spatial focus.

  • Human-to-robot transfer settings: The transfer study covers deformable-object manipulation, long-horizon tasks, and tool use with 10 robot trajectories and 100 human videos.
  • Comparative results: The broader LIBERO results report that ATM outperforms baselines across suites, while Diffusion Policy is further improved by ATM.
  • Experimental protocol: The study evaluates behavioral cloning, ATM with robot demonstrations, and ATM combining human videos with robot data across three random seeds and ten real-world trials.
  • Attention visualization: Track Transformers trained with human videos focus primarily on the object and robot arm, whereas models without human videos show more dispersed attention.

APPENDIX B IMPLEMENTATION DETAILS

The policy architecture processes temporally stacked multi-view images and proprioception through spatial encoding, temporal decoding, and an action head that incorporates reconstructed tracks.

  • Policy Architecture: ViT-T takes temporally stacked multi-view images and proprioception as policy inputs.These inputs are processed across multiple views and timesteps.
  • Implementation Details: The appendix provides separate hyperparameter tables for track-transformer and policy training.These are listed as Tables VII and VIII.
  • Spatial Encoding: Spatial encoding combines image patches, tracks, and modality tokens across views using self-attention with a spatial CLS token.The spatial representation is formed separately at each timestep.
  • Temporal Decoding: Temporal decoding interleaves proprioception, spatial CLS tokens, and learned action CLS tokens across timesteps with causal self-attention.This produces temporally integrated representations for action prediction.
  • Action Head: The action head independently predicts each timestep’s actions by fusing its action CLS token with reconstructed tracks.An MLP parameterizes the actions after this fusion.

B. Efficient Training with Point Filtering

The implementation filters static background points, generates tracks for large-motion points, and examines attention and rollout behavior with and without human-video data.

  • Efficient Training with Point Filtering: Static background points are filtered before an off-the-shelf tracker generates tracks for large-motion points.Points are sampled before and after filtering, as shown in Figure 13.
  • Efficient Training with Point Filtering: Training uses four A100 GPUs with DeepSpeed, while LIBERO policies use teleoperated expert demonstrations and last-epoch checkpoints for rollout evaluation.The demonstrations were collected by human experts using a 3Dconnexion Spacemouse.
  • Efficient Training with Point Filtering: The track transformer is trained on CoTracker-generated ground-truth tracks and selected using the lowest validation loss.Frame stacking is omitted to avoid causal confusion.
  • Cross-Domain Video Data: Including large-scale human videos produces clearer attention maps focused on objects and robot arms than training without human videos.Without human videos, attention can focus on incorrect regions such as background walls.
  • Cross-Domain Video Data: ATM policies trained with human data generate more precise tracks and achieve better policy performance than policies trained without human data.The comparison uses human demonstrations and rollout videos.
Loading 2401.00025v3…