Policy Agnostic RL: Offline RL and Online RL Fine-Tuning of Any Class and Backbone (2412.06685v1)

Abstract: Recent advances in learning decision-making policies can largely be attributed to training expressive policy models, largely via imitation learning. While imitation learning discards non-expert data, reinforcement learning (RL) can still learn from suboptimal data. However, instantiating RL training of a new policy class often presents a different challenge: most deep RL machinery is co-developed with assumptions on the policy class and backbone, resulting in poor performance when the policy class changes. For instance, SAC utilizes a low-variance reparameterization policy gradient for Gaussian policies, but this is unstable for diffusion policies and intractable for autoregressive categorical policies. To address this issue, we develop an offline RL and online fine-tuning approach called policy-agnostic RL (PA-RL) that can effectively train multiple policy classes, with varying architectures and sizes. We build off the basic idea that a universal supervised learning loss can replace the policy improvement step in RL, as long as it is applied on "optimized" actions. To obtain these optimized actions, we first sample multiple actions from a base policy, and run global optimization (i.e., re-ranking multiple action samples using the Q-function) and local optimization (i.e., running gradient steps on an action sample) to maximize the critic on these candidates. PA-RL enables fine-tuning diffusion and transformer policies with either autoregressive tokens or continuous action outputs, at different sizes, entirely via actor-critic RL. Moreover, PA-RL improves the performance and sample-efficiency by up to 2 times compared to existing offline RL and online fine-tuning methods. We show the first result that successfully fine-tunes OpenVLA, a 7B generalist robot policy, autonomously with Cal-QL, an online RL fine-tuning algorithm, improving from 40% to 70% in the real world in 40 minutes.

Policy Agnostic RL: An Insightful Examination

The paper "Policy Agnostic RL: Offline RL and Online RL Fine-Tuning of Any Class and Backbone" presents a methodological advancement in reinforcement learning (RL) that addresses the prevalent challenge of effectively fine-tuning various policy classes and architectures under the actor-critic RL framework. The authors introduce a novel approach coined as "Policy Agnostic RL" (abbreviated as PARL), which leverages a supervised learning paradigm to transcend conventional limitations associated with policy class dependencies.

Core Contributions

The primary contribution of this work is the introduction of a flexible and universal framework for policy improvement that detaches the process of policy optimization from the intricacies of specific policy architectures. The method posits that policy training can rely on a universal supervised learning loss, thereby circumventing the challenges of estimating complex gradients across diverse policy classes. PARL employs a two-stage policy improvement approach:

1. Action Optimization: This stage enhances actions sampled from the current policy through a combination of global optimization and local gradient-based optimization. The strategy involves ranking actions by their associated Q-values and further refining these actions to maximize expected returns.
2. Policy Training via Supervised Learning: In this stage, the policy is trained to mimic the "optimized" actions, thereby sidestepping the computational and stability challenges of backpropagation through deep policy networks.

Empirical Evaluation

The methodology undergoes rigorous empirical validation across varying domains, such as simulated benchmarks from the D4RL suite, including AntMaze tasks, FrankaKitchen tasks, and the CALVIN benchmark, alongside real-world robotics trials. The results demonstrate that PARL outperforms or matches state-of-the-art methods such as Implicit Diffusion Q-Learning (IDQL), Diffusion Q-Learning (DQL), and others in both offline RL performance and online fine-tuning efficiency. Notably, the approach shows a substantial improvement in learning efficiency, particularly in tasks requiring expressive policy classes with complex action distributions.

The paper also reports successful real-world applications, fine-tuning large generalist models such as OpenVLA on a physical robotic platform, yielding significant improvements in task success rates within a constrained interaction time window.

Implications and Future Directions

This work presents profound implications for the field of reinforcement learning. By abstracting policy optimization from specific architectures to a more universal framework, PARL promises enhanced adaptability and scalability in deploying RL solutions across diverse application domains. Practically, the method opens avenues for leveraging powerful policy architectures, such as diffusion models and transformer-based autoregressive policies, without the necessity for bespoke algorithmic modifications, which have traditionally entailed substantial engineering efforts.

Theoretically, this approach encourages a reevaluation of the boundaries between supervised learning and reinforcement learning, particularly regarding policy improvement strategies. It prompts further exploration into optimizing action samples and the potential of integrating global and local optimization techniques more seamlessly into existing RL frameworks.

Future research could focus on optimizing the computational aspects of PARL, particularly for large-scale applications and environments where inference efficiency is critical. Additionally, expanding the applicability to multi-agent setups or domains with even higher dimensional state-action spaces could further validate and enhance the methodology's robustness.

Conclusion

"Policy Agnostic RL" represents a substantive advancement in reinforcement learning methodologies, enabling efficient and stable fine-tuning across a spectrum of policy classes and architectures. Through innovation in action optimization combined with conventional supervised learning, the authors deliver a compelling solution to some of the most persistent challenges in RL, with profound implications for both theoretical exploration and practical deployment.