PaperlayerView exact revision on arXiv

Source-linked AI summary

Scaling Proprioceptive-Visual Learning with Heterogeneous Pre-trained Transformers

Lirui Wang, Xinlei Chen, Jialiang Zhao, Kaiming He

arXiv:2409.20537v1cs.ROcs.CVcs.LG

TL;DR

Robotic foundation models must learn across heterogeneous embodiments, tasks, and environments rather than relying on narrow embodiment-specific data. HPT addresses this with embodiment-specific tokenizers and a shared Transformer trunk, and reports improved transfer performance across simulation and real-world settings while scaling across heterogeneous datasets.

  • Problem

    Robotic policy learning is limited by heterogeneous embodiments, tasks, and environments, creating a need for representations that support broader pre-training.

  • Method

    HPT aligns embodiment-specific proprioception and vision inputs with stems, processes fixed token sequences through a shared Transformer trunk, and uses task-specific heads for actions.

  • Results

    HPT transfers learned representations and improves performance in both simulation and the real world while scaling to over 50 heterogeneous datasets.

  • Takeaways & Limitations

    Heterogeneous pre-training can support transferable robotic representations across embodiments and tasks, with performance improvements reported for new settings.

  • Takeaways & Limitations

    Evaluation is restricted to short-horizon manipulation tasks with fixed embodiments, while data quality, objectives, and training scale remain limited.

Abstract

from arXiv · show

One of the roadblocks for training generalist robotic models today is heterogeneity. Previous robot learning methods often collect data to train with one specific embodiment for one task, which is expensive and prone to overfitting. This work studies the problem of learning policy representations through heterogeneous pre-training on robot data across different embodiments and tasks at scale. We propose Heterogeneous Pre-trained Transformers (HPT), which pre-train a large, shareable trunk of a policy neural network to learn a task and embodiment agnostic shared representation. This general architecture aligns the specific proprioception and vision inputs from distinct embodiments to a short sequence of tokens and then processes such tokens to map to control robots for different tasks. Leveraging the recent large-scale multi-embodiment real-world robotic datasets as well as simulation, deployed robots, and human video datasets, we investigate pre-training policies across heterogeneity. We conduct experiments to investigate the scaling behaviors of training objectives, to the extent of 52 datasets. HPTs outperform several baselines and enhance the fine-tuned policy performance by over 20% on unseen tasks in multiple simulator benchmarks and real-world settings. See the project website (https://liruiw.github.io/hpt/) for code and videos.

1 Introduction

HPT addresses robotics’ heterogeneity by aligning embodiment-specific sensory inputs into a shared policy representation, enabling a shared Transformer trunk to learn across datasets and transfer to new embodiments. Experiments examine scaling across diverse data, models, compute, and domains.

  • Robot policies often require embodiment-, task-, and environment-specific data, while learned policies generalize poorly beyond those settings.
  • HPT uses embodiment-specific stems to align proprioception and vision inputs into a shared latent policy space, with a shared trunk and task-specific heads.The trunk is transferred to new embodiments and tasks, while a new embodiment can use a small tokenizer.
  • HPT scales with dataset quantity and diversity, model size, and training compute across more than 50 data sources.The study investigates scaling behavior using more than 50 individual data sources and models exceeding 1 billion parameters.
  • HPT combines heterogeneous data from real robots, simulations, and human videos during pre-training.The framework incorporates datasets spanning these embodiment domains.
  • The modular training procedure can simplify building policies for new embodiments and tasks in terms of data requirements and generalized performance.

2 Related Works

HPT builds on representation alignment and generalist policy learning while extending heterogeneous pre-training to proprioception, vision, and multiple robotic data domains. Its modular stems, shared trunk, and task-specific heads support this broader setting.

  • Prior representation-learning methods transfer features to unseen tasks and can improve robustness to distribution shifts.
  • Foundation-model research uses task-agnostic objectives and small projection layers to align pre-trained feature spaces, but robotics combines limited data with greater heterogeneity.
  • Robotics research has separately explored visual pre-training, multimodal alignment, and generalist policies using real-robot, human-video, or simulation data.
  • HPT handles broader heterogeneity by aligning proprioception and vision while scaling across real robots, human videos, and simulation data.It is evaluated at larger scale in simulation benchmarks than the compared works.
  • HPT’s stem architecture uses attention to map proprioceptive and visual features into fixed-length token sequences for shared processing.This design supports sequential inputs without increasing token size.

3 Heterogenoues Pre-trained Transformers (HPT)

HPT addresses heterogeneous robot learning by modularizing policies into embodiment-specific stems and heads around a shared transformer trunk. The design tokenizes varied proprioception and vision inputs into fixed-length representations, jointly trains across datasets, and transfers the trunk to new embodiments.

  • Architecture: HPT divides the policy into embodiment-specific stems, a shared trunk, and task-specific heads.Stems align heterogeneous sensory inputs, the trunk processes shared representations, and heads map them to task-specific action spaces.
  • Architecture: Proprioceptive and vision tokenizers convert heterogeneous inputs into fixed numbers of tokens with fixed dimensions for shared trunk processing.The proprioceptive tokenizer handles variable-dimensional robot states, while the vision tokenizer processes camera-image sequences using frozen visual features.
  • Architecture: HPT uses short token sequences and pre-trained vision encoders to accommodate arbitrary image sequences while incorporating proprioception before transfer.The default vision tokenizer maps inputs to 16 tokens, contrasting with approaches that use larger token sequences or add proprioception only during transfer.
  • Architecture: The shared transformer trunk has a fixed parameter count independent of the number of embodiments and tasks, and its pooled output feeds dataset-specific policy heads.Each head maps the pooled trunk feature to a normalized action trajectory and is reinitialized for a new embodiment.
  • Training Objective: Heterogeneous pre-training minimizes behavior-cloning loss across datasets using embodiment-specific stems and heads with one shared trunk.During pre-training, the trunk is updated every iteration while stems and heads are updated according to sampled embodiment and task batches; transfer freezes the trunk and reinitializes stem and head dimensions for the new embodiment.

4 Experiments on Pre-training

HPT pre-training is evaluated for scaling across heterogeneous datasets, model sizes, training samples, and additional embodiment domains. Validation performance generally improves with matched increases in data, compute, and model scale, reaching experiments with 200k trajectories and 52 datasets.

  • Scaling Setup: HPT pre-training is evaluated for scaling behavior across heterogeneous data, model size, and training samples.The experiments use validation loss on unseen trajectories and vary dataset counts, trajectory counts, model sizes, batch sizes, and training tokens.
  • Data Scaling: 200k trajectories and 52 datasets define the largest heterogeneous pre-training setting.The scaled setting includes simulation, deployed-robot, and human-video datasets alongside robot teleoperation data.
  • Data Scaling: Increasing trajectory counts improves validation loss when larger models and more compute are scaled alongside the data.Increasing data with fixed model size and iterations can plateau around 1000 trajectories per dataset, whereas matched scaling continues improving.
  • Model Scaling: Scaling models from 1M to 1B parameters with larger compute achieves lower validation losses until performance plateaus.The reported comparison does not find a significant difference between scaling depth and scaling width.
  • Epoch Scaling: Increasing batch sizes and effective training tokens generally improves performance until convergence.The epoch-scaling experiments use HPT-S with 27 datasets and up to 1000 trajectories per dataset.
  • Additional Embodiments: Adding simulation and human-video datasets is feasible despite their embodiment gaps from real robots.These datasets provide complementary embodiment data to pure teleoperation data and are included in the heterogeneous pre-training framework.

5 Experiments on Transfer Learning

Transfer experiments test whether HPT representations transfer to new embodiments, tasks, and environments in simulation and the real world. Pre-trained policies improve simulated task success and outperform no-trunk and from-scratch baselines in real-world tasks, while showing qualitative robustness under changed conditions.

  • Evaluation Scope: Transfer learning evaluates HPT representations on new embodiments, tasks, and environments across simulation and real-world manipulation.Simulation benchmarks include Meta-World, RoboMimic, and Fleet-Tools; real-world experiments cover pet-care and assembly tasks.
  • Simulation Transfer: Pre-trained HPT models from HPT-B to HPT-XL improve task success rates in closed-loop simulation.The simulation pre-training uses real-world embodiments, while transfer is evaluated on downstream simulated tasks.
  • Simulation Transfer: The Simpler benchmark compares HPT with Octo, RT1-X, and RT2-X on Close Drawer, Move Near, and Pick Coke Can.The evaluation uses the Google EDR embodiment and more than 300 episodes across the tasks.
  • Real-World Transfer: Qualitative real-world results show improved generalization and robustness under varying object quantities, poses, camera configurations, and lighting.The tasks involve dynamic and long-horizon contact-rich precision behaviors in pet care and assembly.
  • Ablations: The Sweep Leftover ablation compares fine-tuned HPT with vision-only pre-trained models, from-scratch policies, and models without proprioception during pre-training.The comparison includes R3M, Voltron, and VC-1.

6 Conclusion

HPT is presented as a modular framework for pre-training robotic policy representations across heterogeneous embodiments, tasks, and environments. The paper reports transfer gains in simulation and real-world settings, while identifying dataset quality, convergence, evaluation scope, and reliability as remaining limitations.

  • Limitations: Dataset curation, data filtering, and pre-training objectives remain under-explored, while heterogeneous pre-training can converge slowly.The authors also note that the training scale remains moderate relative to LLM training.
  • Limitations: Evaluation is restricted to short-horizon manipulation tasks with fixed embodiments, and tested policies typically remain below 90% reliability.These constraints may limit the benefits of higher-capacity models and leave failure modes unresolved.
  • Conclusion: HPT addresses robotic generality limits caused by heterogeneous embodiments, tasks, and operating environments.The conclusion frames heterogeneity as a central constraint on robot learning despite increased data availability.
  • Framework and Data: The framework uses a shared trunk with embodiment-specific inputs and supports heterogeneous datasets from real robots, simulation, deployed robots, and human videos.The dataset mixture includes Open-X, simulation, human-video, and in-the-wild deployed-robot data.
  • Supported Outcome: HPT pre-training can transfer learned representations and improve performance in both simulation and the real world.The paper reports extensive simulation evaluation and real-world transfer experiments using heterogeneous pre-training.
  • Modularity: The architecture is modular across stem encoders, transformer trunks, and policy heads.The trunk can be paired with different vision encoders and heads including diffusion policies, MLPs, and transformers.

A.3 Pre-training Experiment Details.

The pre-training details describe weighted dataset sampling, modular policy components, and simulation evaluation procedures. HPT combines scalable transformer trunks with flexible heads and uses training choices designed to balance heterogeneous datasets and action-prediction objectives.

  • Optimization: HPT uses AdamW with weight decay 0.05, a base learning rate of 0.0002, cosine scheduling, warmups, and dropout.The learning rate is proportionally scaled with batch size, and Huber loss uses δ = 0.1 for action prediction.
  • Architecture: The trunk can operate with multiple policy heads, including diffusion, MLP, and transformer heads.The architecture is designed so trunk training and transfer remain independent of stem encoder architecture and head type.
  • Data Loading: Weighted sampling prevents large datasets from dominating training by assigning sampling probabilities based on dataset size.The procedure uses square-root dataset sizes to compute a normalization constant and sample each batch.
  • Simulation Evaluation: Simulation transfer uses HPT-Base after fine-tuning on supervised datasets because the Simpler tasks emphasize language and omit proprioception.The Simpler evaluation focuses on Close Drawer, Move Near, and Pick Coke Can in the Google EDR embodiment.
  • Simulation Evaluation: Simulation training uses 10–100 trajectories per task, 50 test episodes, a one-step observation window, and an MLP policy head.The default setup trains for 20,000 iterations with batch size 512 and learning rate 1e−5.

A.5 Real-World Experiment Details

The real-world experiments evaluate four robotic tool-use tasks across two sensor and action-space configurations, using supervised trials and human oversight. Transfer learning fine-tunes small policy heads with limited active parameters.

  • Task Definition: Four tool-use tasks are evaluated across two Franka Panda setups with different sensor locations and action spaces.The tasks are Sweep Leftover, Fill Water, Scoop Food, and Switch Insertion; approximately 100 demonstrations are collected per task and each task is evaluated for 15 trials.
  • Evaluation: Success is measured by task completion, with episodes terminable for safety concerns, robot faults, or timeout.Some tasks use partial success scoring, such as 0.5 for pushing some, but not all, piles into the plate.
  • Transfer Learning: Fine-tuning uses MLP and diffusion policy heads with no more than 3Mb of active parameters.Training uses an observation history window of 2, learning rate 2e−5, batch size 256, 20000 iterations, and one NVIDIA RTX 2080Ti GPU.
  • Additional Experiments: The section presents additional experiments and ablation studies for the real-world evaluation.

B.1 Additional Simulation Experiments

Additional simulation experiments compare pre-trained HPT representations with baselines and single-task policies across transfer-learning and unseen-benchmark settings. Pre-training improves validation loss, supports sample efficiency, and matches reported task success rates of strong policy architectures.

  • Transfer Learning: Pre-trained HPT representations can achieve lower validation loss during fine-tuning than four baselines across simulator benchmarks.The comparison uses validation-loss curves for transfer learning with and without pre-trained HPT trunks.
  • Unseen Simulation Benchmarks: HPT trunks improve performance in unseen simulation benchmarks and outperform single-task state-of-the-art architectures.The trunks were not pre-trained with diffusion heads or transformer decoder heads, and the experiments also include a sample-efficiency ablation.
  • Benchmark Setup: The Simpler benchmark evaluates Close Drawer, Move Near, and Pick Coke Can in the Google EDR embodiment across more than 300 episodes.Several initializations are tested for each task.
  • Policy Comparisons: HPT matches Diffusion Policy at 78% success on PushT and ACT at 60% on Aloha Transfer Box.The results use diffusion and transformer-decoder heads respectively, demonstrating compatibility with different policy-head architectures.
  • Token Ablation: Removing either visual or proprioceptive tokens hurts downstream-policy performance in simulation transfer learning.The ablation indicates that both input types contribute to the transferred policy representation.

B.2 Ablation Study on the Stem

The stem ablations test proprioception, vision, and vision encoders under a fixed 27-dataset setting. Removing proprioceptive or visual information increases validation loss, while broader encoder and real-world failure analyses identify remaining scope for improvement.

  • Ablation Setup: The stem ablation fixes 27 datasets and limits each dataset to 1000 trajectories while varying stem components.The study examines the pre-training setup for HPT-S.
  • Input Ablations: Removing either proprioception or vision increases validation loss, showing both inputs are critical for policy pre-training at scale.The ablation covers proprioception, vision stems, and vision encoders.
  • Scope and Future Work: Further ablations of image resolution and joint vision-backbone fine-tuning are left for future work.The default implementation focuses on single-view visual information, although the stem can extend to multiple views and other modalities.

B.3 Pre-training Ablation Study

Pre-training ablations examine additional modalities, tokenization choices, observation history, and policy-head designs. Multiple views, language, and action histories improve the default setting in some datasets, while real-world policies still exhibit spatial-precision failures.

  • Additional Modalities: Multiple views, language guidance, and previous action trajectories improve over the default single-view vision-plus-proprioception setting in some datasets.The gains are more pronounced for multiple views on insertion and language on Language Table; prior actions add context and embodiment information.
  • Architecture Ablations: Token size and observation history have little effect on averaged validation loss, suggesting a trade-off between policy information and generalization.The study also compares discretized and continuous action-regression heads and chooses continuous regression for generality.
  • Failure Cases: Real-world policies sometimes overshoot or undershoot, such as pouring before reaching the mug or bowl.The authors attribute these failures to possible spatial-precision and data-quality issues and suggest targeted data recollection or better vision-encoder fine-tuning.

D Discussion and Future Directions

The discussion frames HPT as a scalable approach for heterogeneous robot learning while emphasizing that intrinsic validation metrics have important evaluation gaps. Future work focuses on richer data, embodiment-aware architectures, broader benchmarks, and additional modalities and domains.

  • Validation loss depends on dataset composition, evaluation subsets, and training allocation, and does not directly measure downstream closed-loop task execution.
  • HPT aligns heterogeneous proprioception and vision through modular stems, a shared scalable transformer trunk, and task-specific action heads.
  • Fine-tuning remains required for robotics generalist models, motivating embodiment-specific architectures and training objectives beyond supervised learning.
  • Unified simulation benchmarks are needed for consistent comparisons, while real-world evaluation should extend to longer-horizon fine manipulation with bimanual or mobile setups.
  • Future studies should examine scaling laws and incorporate modalities and domains such as 3D point clouds, tactile data, simulation, and human data.
Loading 2409.20537v1…