Self-Improving Embodied Foundation Models

Abstract: Foundation models trained on web-scale data have revolutionized robotics, but their application to low-level control remains largely limited to behavioral cloning. Drawing inspiration from the success of the reinforcement learning stage in fine-tuning LLMs, we propose a two-stage post-training approach for robotics. The first stage, Supervised Fine-Tuning (SFT), fine-tunes pretrained foundation models using both: a) behavioral cloning, and b) steps-to-go prediction objectives. In the second stage, Self-Improvement, steps-to-go prediction enables the extraction of a well-shaped reward function and a robust success detector, enabling a fleet of robots to autonomously practice downstream tasks with minimal human supervision. Through extensive experiments on real-world and simulated robot embodiments, our novel post-training recipe unveils significant results on Embodied Foundation Models. First, we demonstrate that the combination of SFT and Self-Improvement is significantly more sample-efficient than scaling imitation data collection for supervised learning, and that it leads to policies with significantly higher success rates. Further ablations highlight that the combination of web-scale pretraining and Self-Improvement is the key to this sample-efficiency. Next, we demonstrate that our proposed combination uniquely unlocks a capability that current methods cannot achieve: autonomously practicing and acquiring novel skills that generalize far beyond the behaviors observed in the imitation learning datasets used during training. These findings highlight the transformative potential of combining pretrained foundation models with online Self-Improvement to enable autonomous skill acquisition in robotics. Our project website can be found at https://self-improving-efms.github.io .

• The paper introduces a two-stage framework combining imitation learning with steps-to-go prediction to extract dense rewards and enhance policy success.
• The methodology boosts sample efficiency and real-world performance, increasing success rates significantly with minimal additional robot time.
• The framework leverages web-scale multimodal pretraining to generalize beyond imitation datasets, enabling autonomous acquisition of novel robotic skills.

Self-Improving Embodied Foundation Models: A Two-Stage Post-Training Framework for Autonomous Robotic Skill Acquisition

Overview and Motivation

This paper introduces a two-stage post-training framework for Embodied Foundation Models (EFMs) in robotics, inspired by the reinforcement learning (RL) post-training paradigm established in LLMs. The approach leverages web-scale pretrained multimodal models and augments standard behavioral cloning (BC) with a steps-to-go prediction objective, enabling the extraction of dense, data-driven reward functions and robust success detectors. The second stage, termed "Self-Improvement," utilizes these learned signals for online RL, allowing fleets of robots to autonomously practice and acquire novel skills with minimal human supervision. The framework is validated across both simulated and real-world robotic domains, demonstrating significant improvements in sample efficiency, policy success rates, and generalization capabilities beyond the scope of imitation datasets.

Figure 1: Overview of the proposed two-stage fine-tuning approach, combining supervised imitation and steps-to-go prediction with online self-improvement via RL.

Methodology

Stage 1: Supervised Fine-Tuning (SFT)

EFMs are initialized from web-scale pretrained multimodal models (PaLI-3B in this work) and fine-tuned on imitation datasets using two objectives:

• Behavioral Cloning (BC): Standard supervised learning to maximize the likelihood of dataset actions conditioned on observations and goals.
• Steps-to-Go Prediction: The model predicts the number of timesteps remaining to achieve a specified goal, given the current observation. This auxiliary objective is critical for enabling reward extraction and success detection in Stage 2.

Mathematically, the steps-to-go prediction is formulated as:

$$\mathcal{L}_{\text{steps-to-go}} = -\mathbb{E}_{(o_t, a_t, g_{t'}) \sim \mathcal{D}} \left[ \log p^{EFM}_{\text{steps-to-go}(t' - t | o_t, g_{t'})} \right]$$

Stage 2: Self-Improvement via Online RL

The second stage leverages the steps-to-go predictions to define a dense, well-shaped reward function:

$$r(o_t, a_t, o_{t+1}, g) = d(o_t, g) - d(o_{t+1}, g)$$

where $d(o, g)$ is the expected steps-to-go to goal $g$ from observation $o$.

A robust success detector is defined as:

$\text{success}(o, g) = \mathbb{1}[d(o, g) \leq s]$

with $s$ a small threshold.

The RL update uses the REINFORCE loss with Monte Carlo returns computed from the above reward, and policy updates are performed on collected trajectories. Notably, the reward and success signals are derived from a frozen Stage 1 checkpoint, decoupling reward inference from policy learning and eliminating the need for ground-truth instrumentation.

Figure 2: Example trajectory and steps-to-go predictions in the Aloha Single Insertion Task, illustrating reward shaping and policy progress.

Figure 3: Visualization of steps-to-go prediction distributions, showing the model's sensitivity to subtle task events and recovery dynamics.

Experimental Results

Sample Efficiency and Policy Improvement

Experiments on LanguageTable and Aloha domains (both simulated and real) demonstrate that the two-stage framework yields policies with substantially higher success rates and greater sample efficiency compared to scaling imitation data alone. For instance, in LanguageTable, 10% additional robot time for Self-Improvement increases success rates from 45% to 75%, whereas an 8× increase in imitation data only yields a modest improvement to 60%.

Figure 4: Pointmass Navigation domain results, showing rapid policy improvement via Self-Improvement compared to BC.

Figure 5: Stage 2 Self-Improvement results across domains, highlighting significant gains in success rates with minimal additional data.

Real-World Applicability

The framework is validated in real-world LanguageTable experiments, where a single human operator supervises multiple robots. Self-Improvement with only 3% additional episodes boosts success rates from ~62% to ~87%, outperforming BC policies trained with 80% of the imitation dataset.

Figure 6: Real-world and simulated LanguageTable environments used for large-scale experiments.

Figure 7: Success rate plots during real-world Self-Improvement, demonstrating robust and reproducible policy gains.

Ablations: Importance of Multimodal Pretraining

Ablation studies reveal that the sample efficiency and robustness of Self-Improvement are critically dependent on web-scale multimodal pretraining. Reward models initialized from scratch or unimodal pretraining (Uni-PaLI) underperform significantly, especially in low-data regimes.

Figure 8: Ablation results showing the impact of multimodal pretraining on Self-Improvement efficacy and domain transfer.

Generalization: Autonomous Acquisition of Novel Skills

The combination of online Self-Improvement and foundation model pretraining enables policies to generalize beyond the imitation dataset, acquiring novel behavioral skills. In the BananaTable task, policies trained on LanguageTable data rapidly learn to manipulate a banana—a previously unseen object—achieving a success rate increase from 63% to 85% in 8 hours.

Figure 9: Strong generalization to BananaTable, with policies acquiring new manipulation strategies for out-of-distribution objects.

Theoretical Insights

The steps-to-go reward function is shown to be a form of implicit reward shaping, regularizing policy updates towards regions of the state space where the imitation policy is proficient. The mathematical formulation aligns the reward with the value function of the dataset policy, providing dense feedback and reducing variance in policy gradient estimation.

Implementation Considerations

• Model Architecture: The framework is agnostic to the underlying foundation model, but benefits from large-scale multimodal pretraining (PaLI-3B used here).
• Tokenization: Actions and steps-to-go are discretized and mapped to the model's token space, enabling seamless integration with transformer-based architectures.
• Compute Resources: Training requires substantial compute (64–128 TPUv3/v4 for SFT; reduced resources for RL), but reward inference is decoupled and can be scaled independently.
• Infrastructure: Distributed actor-learner architectures are employed, with local inference for high-frequency control tasks.
• Human Supervision: Minimal, limited to periodic resets; all reward and success signals are model-derived.

Limitations and Future Directions

• Reward Model Expressivity: Handling out-of-distribution failure states remains challenging; future work may leverage broader data sources or chain-of-thought reasoning for reward inference.
• Hierarchical Skill Chaining: Steps-to-go prediction naturally supports hierarchical control, but scalable episode boundary annotation is an open problem.
• RL Algorithms: The current framework uses on-policy REINFORCE for stability; off-policy methods may further improve sample efficiency.
• Pretraining Curricula: As robot-specific multimodal corpora grow, designing curricula that balance physical reasoning and semantic knowledge will be critical.
• General-Purpose Post-Training: Extending Self-Improvement to zero-shot generalization across tasks is a promising direction.

Conclusion

The proposed two-stage post-training framework for EFMs—combining supervised imitation and steps-to-go prediction with online Self-Improvement—enables autonomous, sample-efficient skill acquisition in robotics. The approach eliminates the need for manual reward engineering, leverages the generalization capabilities of web-scale foundation models, and unlocks the ability for robots to practice and acquire novel skills beyond their training data. These results establish a new paradigm for scalable, autonomous robot learning and highlight key avenues for future research in embodied AI.