Technical Summary: Dual-Latent Space Reinforcement Learning for Generative Robot Policy

Problem Statement

Pretrained generative robot policies, such as Diffusion Policies and flow-matching models, leverage large-scale datasets to learn expressive action priors. However, these models often struggle when downstream tasks deviate from the pretraining distribution, particularly in long-horizon tasks where small spatial or temporal errors lead to failure. While online adaptation is necessary, full fine-tuning of billion-parameter models is computationally prohibitive.

Existing parameter-efficient adaptation methods, such as Reinforcement Learning (RL) based steering, typically operate only at the boundaries of action generation. For instance, methods like Diffusion Steering via Reinforcement Learning (DSRL) learn an observation-conditioned initial-noise policy to steer a frozen generator. While effective, this approach only influences the starting point of the generation process. It fails to modulate intermediate action representations, limiting its ability to make precise corrections to contact locations, action magnitudes, or local trajectories during the generation process. Conversely, residual policies apply corrections after action generation, failing to shape the internal evolution of action representations. There is a lack of an RL mechanism that combines global behavior selection with direct modulation of intermediate action representations within a frozen generator.

Methodology: Dual-Latent Space Reinforcement Learning (DLSRL)

To address these limitations, the authors propose Dual-Latent Space Reinforcement Learning (DLSRL). This framework introduces a novel control interface that complements initial-noise steering with representation-level control inside a frozen generative policy.

1. Dual-Latent Actor

Unlike traditional noise-space RL methods that only predict an initial noise vector ($z_t$), the DLSRL actor predicts two distinct latent variables conditioned on the current observation ($o_t$):

• Initial-noise latent variable ($z_t$): Steers the overall behavior generation, similar to existing methods.
• Action-representation latent variable ($u_t$): Designed for intermediate feature modulation.

The actor is parameterized as $\pi_\theta(\cdot | o_t)$, outputting the joint distribution $(z_t, u_t)$.

2. Adapter Feature Mapping and Residual Injection

The action-representation latent variable $u_t$ is processed through a lightweight adapter feature mapper $A_\omega$ to produce adapter features $f_t$:
$$f_t = A_\omega(u_t)$$

These features are then injected into the frozen base policy $G_\phi$. Specifically, the adapter features are scaled by an injection strength $\lambda_{inj}$ and residually added to the hidden states of action tokens within selected frozen Transformer blocks:
$$\tilde{H}^{(l)}_{t,k} = H^{(l)}_{t,k} + \lambda_{inj} f_t$$
where $H^{(l)}_{t,k}$ represents the action-token hidden states at layer $l$ and generation step $k$. This design allows for direct adjustment of action representations without updating the base policy parameters or altering the network architecture. The base policy generates the final action chunk as:
$$a_t = G_\phi(o_t, z_t, f_t)$$

3. Dual-Latent Policy Optimization

The training framework employs a value-distillation mechanism to optimize the dual-latent actor without backpropagating through the frozen generator:

• Action-Space Critic ($Q^{act}_\psi$): Trained on online interaction data using standard off-policy temporal-difference learning to estimate the return of action chunks.
• Latent-Space Critic ($Q^{lat}_\nu$): Optimized via a distillation objective to match the values predicted by the action-space critic. This critic learns the joint effect of $z_t$ and $f_t$ on task returns.
  $$L_{distill} = \mathbb{E}[(Q^{lat}_\nu(o_t, \hat{z}_t, \hat{f}_t) - \text{sg}[Q^{act}_\psi(o_t, \hat{a}_t)])^2]$$
• Actor Update: The dual-latent actor is updated using an entropy-regularized objective that maximizes the latent-space value:
  $$L_{actor} = \mathbb{E}[\alpha \log \pi^z_\theta(z_t | o_t) - Q^{lat}_\nu(o_t, z_t, f_t)]$$

During deployment, only the actor and the adapter mapper are retained; the critics are discarded.

Key Contributions

The paper outlines three primary contributions:

1. Novel Framework: The introduction of a latent space RL framework that simultaneously optimizes initial-noise and action-representation latent variables, extending policy guidance from the sampling initialization phase to the intermediate representation phase.
2. Lightweight Modulation Mechanism: The development of a residual modulation mechanism that maps action-representation latents into adapter features aligned with action tokens, injecting them into Transformer hidden states without requiring base policy updates.
3. Empirical Validation: A series of experiments across diffusion and flow-matching policies demonstrating that DLSRL accelerates online adaptation speed while maintaining competitive final performance.

Experimental Results

Experiments were conducted on the RoboMimic (Lift, Can, Square) and LIBERO (six tasks including Stove-On, Bowl-Drawer-to-Plate, etc.) benchmarks using frozen Transformer-based Diffusion Policies and the $\pi_0$ flow-matching policy.

• Adaptation Speed: DLSRL consistently demonstrated faster performance improvement during early and intermediate training stages compared to baselines like DSRL, DPPO, and JSRL.
• Success Rates:
  • On the Can task (RoboMimic), DLSRL achieved approximately 99% success in the later stages, outperforming DSRL (~90%).
  • On LIBERO tasks, DLSRL reached saturated performance with fewer environment interactions than DSRL.
• Efficiency: Table I shows that DLSRL achieved shorter average episode lengths across all six LIBERO tasks, with reductions ranging from 15.27% to 25.40% compared to DSRL.
• Ablation Studies:
  • Removing the initial-noise steering (DLSRL-Rep) still improved success rates from ~27% to ~52% on the Can task, confirming that representation-level modulation alone contributes to adaptation.
  • Analysis of injection strength ($\lambda_{inj}$) revealed that while higher values ($0.09$) offered rapid initial gains, they caused instability. A moderate value ($0.06$) provided the best balance, achieving 100% success.

Significance and Claims

The authors claim that DLSRL represents the first work to integrate initial-noise steering with action-token hidden-state modulation. The significance of this work lies in its ability to provide an additional control interface inside a frozen generator, enabling precise corrections to contact locations and trajectories without the computational cost of fine-tuning large generative models.

The paper concludes that DLSRL effectively accelerates online robot policy adaptation. While the method shows superior adaptation speed, the authors modestly note that the advantage is more pronounced in tasks requiring high action precision (e.g., Can, Square) compared to simpler tasks where other methods may eventually converge to similar success rates. Future work is planned to evaluate the framework in real-world environments and investigate the effects of specific architectural choices (injection layers, steps, strengths) on performance.