Warp RL: Reshaping Base Policy Distributions for Dynamics Adaptation

Warp RL: Reshaping Base Policy Distributions for Dynamics Adaptation

Background and Motivation

Adapting reinforcement learning (RL) policies to changed dynamics is vital for reliable robotic deployment. Traditionally, residual RL techniques freeze a base policy and learn an additive correction, resulting in actions of the form $a = z + \delta_\theta(s)$ for base sample $z \sim \pi_{\text{base}}(\cdot|s)$. While effective in some domains, this additive structure is a translation operator: it shifts the distribution mean, but leaves the variance and higher-order structure invariant. This limits corrective flexibility under substantive dynamics shifts where mismatches are geometric or state-dependent, leading to failure modes such as incorrect distribution variance, miscalibrated action confidence, and the inability to modulate corrections non-uniformly over the action space.

Warp RL: The Framework

Warp RL replaces the additive correction with a parameterized, invertible, state-conditioned transformation $T_\theta(z, s)$, forming actions $a = T_\theta(z, s)$ with $z \sim \pi_{\text{base}}(\cdot|s)$. The transformation is constructed as a rational-quadratic (RQ) spline flow, ensuring monotonicity, analytic invertibility, exact density evaluation, and identity-initialization. Critically, additive residual policy learning becomes a special case within this more general warping framework.

Figure 1: ManiSkill3 evaluation tasks used for controlled simulation experiments and real-world adaptation.

This generic warping operator, applied on top of a frozen base policy, strictly generalizes prior additive residual techniques. The architecture’s expressivity allows scaling, skewing, and non-uniform reshaping of the base action distribution, thereby addressing the structural deficiencies that cripple additive correction under significant or nonlinear dynamics shifts.

Methodological Details

Correction Hierarchy and Instantiations

Warp RL unifies three classes of correction operators:

• Additive (Residual RL): $T(z, s) = z + \delta(s)$
• Affine bijection: $T(z, s) = \alpha(s) \odot z + \delta(s)$ with learnable per-dimension scaling
• Nonlinear bijection (Warp RL): $T(z, s) = f(z; \phi(s))$, such as a state-conditioned RQ-spline

Identity-initialization (the warp starts as the identity function) is critical for safe deployment: preliminary behavior matches the well-trained base policy, and any adaptation is performance-driven.

Training Algorithms

The proposed warping architecture is agnostic to the optimizer and supports:

• Policy gradients (PPO): Exploiting closed-form densities using the change-of-variables formula, enabling direct compatibility with standard RL updates.
• Off-policy algorithms (SAC): Using warped log-probabilities with maximum entropy objectives.
• Evolution Strategies (ES): Treating $T_\theta$ as a black-box and exploring parameter space via Gaussian perturbations.

The invertible structure of the RQ-spline supports analytic computation of action log-densities, critical for PPO and SAC.

Experimental Results

Warp RL is benchmarked against standard and ablated variants on four ManiSkill3 manipulation tasks—PushCube, PushT, LiftPegUpright, and PegInsertionSide—as well as on a real-world SO101 robot performing peg insertion under dynamics shift.

Figure 2: Learning curves for main correction methods under shifted dynamics, showing IQM of success across seeds. Warp-ES consistently achieves rapid and robust recovery, especially when adaptation requires non-translational corrections.

Warp RL demonstrates superior adaptability across a spectrum of task-specific dynamics shifts:

• Translational shifts: On PushCube, both residual and warping methods match or exceed the base policy, as expected when a mean-shift suffices.
• Variance and geometry shifts: On tasks requiring altered distribution variance or non-uniform strategy (e.g., PegInsertionSide, LiftPegUpright), additive corrections stall or degrade (sometimes below the unadapted base). Warp RL—in particular, when coupled with ES optimization—achieves significantly higher success rates.
• Aggregate recovery: Warp RL consistently enables higher rates of full performance recovery across task-seed combinations.

  Figure 3: Ablation under ES training: residual (additive), affine, unconstrained MLP, and spline (Warp) architectures. The structured, monotonic RQ-spline consistently yields superior and stable adaptation.

Ablation studies reveal that:

• Parameter-space adaptation alone (residual or unconstrained MLP) is insufficient, and can destabilize adaptation.
• The gains from Warp RL are attributable to the structured, monotonic spline parameterization, not simply model capacity or optimizer choice.

  Figure 4: Impact of exploration inheritance. Removing base policy exploration inheritance destabilizes residual adaptation, affirming the importance of structure-induced exploration for robust transfer.

Comparison to policy fine-tuning baselines further confirms Warp RL’s efficiency: output-space adaptation via warping consistently avoids the pitfalls and optimization inefficiencies of retraining the base network.

Real-World Peg Insertion

In a sim-to-real setting, both additive and splined corrections close much of the transfer gap, but the spline-based Warp RL reduces the median episode duration by 30% with no sacrifice in success rates:

• Residual correction: Improves success rates, but does not shorten cycle times relative to the base (reflecting mere mean shift).
• Warp correction: Shifts probability mass toward the effective region in the action space, focusing exploration and execution, yielding faster—though slightly more variable—task completion.

Additional Analysis

Figure 5: Learning dynamics for additional adaptation baselines, contrasting frozen-base corrections with full and last-layer fine-tuning. Warp-ES surpasses both, especially under challenging dynamics shifts.

Fine-tuning the whole policy (KL-regularized or last-layer only) is competitive but less reliable and efficient than output-space warping. The hierarchical, lightweight warp approach preserves base competency while efficiently tailoring to new domains.

Theoretical and Practical Implications

Warp RL reframes post-hoc policy adaptation by generalizing additive correction to arbitrarily expressive, invertible distribution transformations. This theoretically overcomes inherited structural limitations of residual RL, making adaptation more robust to a wide range of dynamics shifts—even when distributional mismatch cannot be resolved by translation or affine transformation alone.

Practically, Warp RL serves as a “drop-in” adapter, integrating seamlessly into existing RL and sim-to-real infrastructures. Its modularity ensures compatibility with frozen pre-trained policies (including those from behavior cloning, chunked-action, or potentially diffusion models), and its computational efficiency enables real-robot deployment with minimal intervention.

Limitations and Future Directions

While Warp RL expands the corrective space, several challenges remain:

• The RQ-spline operates on a bounded domain; mapping from unbounded Gaussians introduces interplay with boundary behavior, especially in tasks with saturated action commands.
• Cross-dimensional and sequence-level warping (e.g., autoregressive or temporal transforms) could further increase adaptation capacity for higher-dimensional or time-dependent policies.
• Real-robot studies are task-limited and would benefit from broader empirical validation and larger-scale benchmarking.

Further, integrating warping with non-Gaussian, behavior-cloned, or diffusion-based base policies is a promising direction to bridge the gap to broader classes of modern RL agents.

Conclusion

Warp RL introduces principled, highly expressive, identity-initialized, invertible warping of base policy distributions for robust adaptation under dynamics shifts. Across simulated and real-robot tasks, it consistently matches or exceeds conventional residual and fine-tuning baselines, improves practical deployment efficiency, and defines a general framework for next-generation policy adapters in RL and robotics.