PaperlayerView exact revision on arXiv

Source-linked AI summary

Policy Agnostic RL: Offline RL and Online RL Fine-Tuning of Any Class and Backbone

Max Sobol Mark, Tian Gao, Georgia Gabriela Sampaio, Mohan Kumar Srirama, Archit Sharma, Chelsea Finn, Aviral Kumar

arXiv:2412.06685v1cs.LGcs.AI

TL;DR

RL adaptation across policy classes is difficult because deep RL methods often rely on assumptions tied to particular policy architectures. PA-RL improves sampled actions with a critic and trains policies on those actions using supervised learning; it reports a 13% aggregate advantage over the next-best approach and real-robot gains for diffusion policies and OpenVLA.

  • Problem

    Deep RL methods often entangle their training objectives and design with policy class, making effective fine-tuning across classes challenging.

    Source: §1 · p. 1

  • Method

    PA-RL samples policy actions, improves them through critic-based re-ranking and gradient ascent, then uses the optimized actions to supervise the policy.

    Source: §4 · p. 6, §4.1 · p. 7

  • Results

    PA-RL outperforms the next-best approach by 13% in aggregate across domains, and improves OpenVLA by 75% after zero-shot trials and 40 minutes of online fine-tuning.

    Source: §1 · p. 3

  • Takeaways & Limitations

    The authors report one actor-critic approach fine-tuning diffusion policies, autoregressive categorical transformers, and a large pre-trained robotic VLA policy.

    Source: §1 · p. 3

    PA-RL requires sampling multiple policy actions, which is expensive for large foundation policies.

    Source: §6 · p. 16

Abstract

from arXiv · show

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.

1. Introduction

Policy-agnostic RL (PA-RL) addresses the difficulty of adapting RL methods across policy classes by separating action improvement from policy training. It uses optimized actions for supervised policy updates and reports strong results across simulated and real-robot tasks.

1. Introduction

  • The paper motivates RL fine-tuning by the need for policies to adapt from autonomous, non-expert experience as tasks or environments change.It argues that imitation alone, whether done once or repeatedly, does not guarantee the most efficient learning in these scenarios.Source: §1 · p. 1
  • Most deep RL methods couple training design to a policy class; for example, SAC’s reparameterization suits Gaussian policies but can destabilize when used with diffusion policies.Adapting methods across policy classes can require weaker alternatives or modifications to other parts of the approach.Source: §1 · pp. 1–2
  • PA-RL fine-tunes different policy classes and backbones by improving actions directly, then training the policy to imitate those optimized actions with a supervised loss.It samples multiple candidate actions and takes critic-gradient steps to improve them before using them to train the policy.Source: §1 · p. 2
  • PA-RL outperforms the next-best approach by 13% in aggregate across domains, with its largest gains on long-horizon tasks with multimodal offline data.It improves diffusion policies by 80–100% on two real-robot manipulation tasks within 1–2 hours, and improves OpenVLA by 75% after zero-shot trials and 40 minutes of online fine-tuning.Source: §1 · p. 3
  • PA-RL applies one actor-critic approach to diffusion and autoregressive categorical transformer policies, including a large pre-trained robotic VLA policy.The authors report these as first results for the listed policy types and settings using a single approach.Source: §1 · p. 3

2. Related Work

Related work identifies a mismatch between policy-specific RL machinery and the variety of policy classes, while prior approaches often use restricted policy-improvement mechanisms. PA-RL instead combines critic-optimized policy samples with supervised training, and uses offline RL pre-training to train its proposal distribution.

2. Related Work

  • Gaussian-focused policy gradients can be unstable for diffusion policies, while supervised regression on dataset actions may fail to extract the best policy from autoregressive transformers.Prior diffusion approaches also rely on re-ranking, diffusion-specific gradients, or unstable reparameterized gradients, with differing limitations.Source: §2 · pp. 3–4
  • PA-RL combines global action re-ranking with local action-gradient optimization, whereas cited diffusion methods use one or the other and have distinct constraints.IDQL uses re-ranking without iterative distillation; DIPO and DDiffPG use action gradients online without pre-training.Source: §2 · pp. 3–4
  • For autoregressive categorical transformers, prior value-based fine-tuning required modifying value-function learning, whereas PA-RL reports success without such modifications.The authors report no prior value-based RL attempt for these models apart from Chebotar et al. [4].Source: §2 · p. 4
  • PA-RL samples new actions, optimizes them against the critic, and trains the policy on them with NLL, unlike prior RL-as-supervised-learning methods that regress on dataset or replay-buffer actions.The passage states this allows aggressive updates while retaining supervised learning’s simplicity; cited theory links on-policy weighted regression to mode-seeking behavior.Source: §2 · p. 4
  • Prior CEM-based methods optimize actions from a Q-function but do not use offline RL pre-training to train the proposal distribution.The paper says this pre-training matters because critic values outside the dataset support can be erroneous and hurt the approach.Source: §2 · p. 4

3. Problem Setup and Preliminaries

The section formalizes offline and offline-to-online RL fine-tuning, then describes the policy classes and actor-critic methods used in the experiments.

Problem setup

  • An MDP specifies state and action spaces, dynamics, rewards, an initial-state distribution, and a discount factor; the section also defines the policy Q-function and its neural-network estimate.Actions are d-dimensional continuous vectors in [-1, 1]^d.Source: §3 · p. 4

Problem setting

  • The study covers fully offline learning from a behavior-policy dataset and offline-to-online fine-tuning with autonomous interaction, aiming to optimize performance with few online samples.The offline setting uses a dataset of state, action, reward, and next-state experience; the online setting starts from an offline-learned policy.Source: §3 · p. 4

Policy classes and parameterizations

  • Experiments use continuous-action diffusion policies and transformer autoregressive policies that produce categorical action tokens.The autoregressive setup uses 128-bin discretization for simulation and the OpenVLA tokenizer for real-world experiments.Source: §3 · p. 5
  • Diffusion policies use a conditional DDPM denoiser and the standard diffusion training loss, which is not identical to negative log likelihood but is typically derived from a lower-bound approximation.Source: §3 · p. 5

Offline RL and online fine-tuning methods

  • The approach changes policy optimization while retaining the base algorithm’s critic-training procedure; experiments cover IQL and Cal-QL actor-critic methods.Cal-QL samples actions from the learned policy for critic training, whereas IQL regresses the Q-function to a higher expectile of the value function without querying new policy actions.Source: §3 · p. 5

Policy optimization

  • PA-RL optimizes candidate actions through critic re-ranking and gradient ascent, then trains the policy to imitate the most optimized action.Source: Figure 2 · p. 6

4. Policy Agnostic RL (PA-RL): Training Multiple Policy Classes with Actor-Critic RL

PA-RL separates action improvement from policy-parameter training: it optimizes actions, then uses supervised learning to train the policy to imitate them. Global and local action optimization supply those targets.

4. Policy Agnostic RL (PA-RL): Training Multiple Policy Classes with Actor-Critic RL

  • PA-RL replaces policy-gradient updates with supervised learning on optimized actions, avoiding policy-parameter Q-gradients and high-variance policy-gradient estimates.The approach targets difficulties with diffusion and autoregressive policies while retaining a policy-specific supervised learning loss.Source: §4 · p. 6

4.1. Stage I: Action Optimization

  • Global optimization samples multiple actions from the policy, ranks them by critic Q-values, and retains the top m candidates for refinement.The policy initializes action optimization to keep candidates near the policy’s support and preserve offline priors during fine-tuning.Source: §4.1 · pp. 6–7
  • Local optimization applies Q-gradient ascent directly to each retained action, complementing global ranking by making fine-grained changes nearby.Combining both steps addresses the respective risks of local minima and insufficiently fine-grained optimization.Source: §4.1 · p. 7

4.2. Stage II: Policy Training via Supervised Learning

  • The policy-training stage distills optimized actions using supervised likelihood objectives, with diffusion policies using a DDPM objective and autoregressive transformers using next-token cross-entropy.PA-RL can train against the highest-Q optimized action or a Q-weighted distribution over optimized actions.Source: §4.2 · pp. 7–8

4.3. Putting it All Together: Final PA-RL Algorithm

  • Using actions sampled from the current policy lets optimization move farther from the data manifold than AWR while remaining closer to it than CEM-based critic maximization.For IQL, PA-RL replaces the AWR update; for Cal-QL, optimized actions also supply Bellman-backup actions, while critic training remains unchanged. Local updates are accepted only when Q increases.Source: §4.2 · p. 8, §4.3 · p. 9

5. Experimental Evaluation

The evaluation tests PA-RL across simulated benchmarks, hybrid RL, real-robot fine-tuning, and controlled ablations. It reports strong simulated results, real-robot improvements, and evidence that global and local action optimization contribute differently across settings.

5.1. Results: Simulated Benchmarks from State and Image Observations

  • 13% higher fine-tuning performance than the next-best method in aggregate, with a 69% improvement in CALVIN; learning curves largely stay above or near the next-best method.PA-RL + Cal-QL also has the highest aggregate offline performance among the compared approaches.Source: §5.1 · p. 10, Figure 3 · p. 10, Table 1 · p. 11
  • In hybrid RL, PA-RL + RLPD improves a pretrained diffusion policy without critic pretraining, while PA-RL + IQL matches IDQL on FrankaKitchen.The RLPD result is reported after 200 episodes; the IQL comparison uses the FrankaKitchen domain.Source: Table 2 · p. 11, §5.1 · p. 11
  • 26% better average performance than tanh-Gaussian policies across three tasks when fine-tuning autoregressive categorical transformer policies.Each action dimension is discretized into 128 bins before autoregressive modeling.Source: §5.1 · p. 11

5.2. Results: RL Fine-Tuning of Robot Policies in the Real World

  • On two real-robot diffusion tasks, PA-RL improves offline-pretrained policies, including under distractor-object shift; the cup task succeeds after 50 episodes.Table 3 reports comparisons with iterated filtered behavior cloning on task (a) and performance under added distractors on task (b).Source: Table 3 · p. 12, §5.2 · p. 12, Figure 4 · p. 13

5.2.2. OpenVLA Fine-Tuning in the Real World

  • After 40 minutes of online fine-tuning, OpenVLA success rises from 40% in zero-shot trials to 70%.The fine-tuned policy also reduces wrong-object grasps near distractors and more securely grasps the target.Source: §5.2.2 · p. 13

5.3. Ablation Studies and Controlled Experiments

  • Global optimization is critical on diverse AntMaze data, while both global and local optimization matter on narrower CALVIN data.The authors recommend global optimization generally and local optimization when dataset action distributions are narrow.Source: §5.3 · p. 14, Table 4 · p. 14
  • CEM initialized from scratch performs poorly and overestimates Q-values; even with pretrained initialization, it trails PA-RL on low-coverage multimodal Kitchen and CALVIN tasks.The authors attribute this latter weakness to CEM averaging behavioral modes, which can produce actions unlikely under the pretrained policy.Source: §5.3 · pp. 14–15, Figure 5 · p. 15
  • Self-imitation learning using stale dataset or replay-buffer actions fails to achieve positive performance on antmaze-large-diverse-v2, where multiple base-policy samples are critical.This comparison disables local optimization and distills weighted likelihoods from one stale-policy action sample.Source: §5.3 · p. 15

6. Discussion and Conclusion

PA-RL fine-tunes policies across classes and parameterizations by improving action samples against a critic and using them to supervise the policy. The paper reports state-of-the-art online fine-tuning results across simulation tasks, while identifying costs and critic generality as remaining limitations.

6. Discussion and Conclusion

  • PA-RL uses critic-based re-ranking and gradient ascent to improve multiple action samples, then uses those actions to supervise policies of different classes and parameterizations.This is the paper’s actor-critic approach to policy fine-tuning.Source: §6 · p. 15
  • PA-RL achieves state-of-the-art online fine-tuning results across simulation tasks.The conclusion also reports results on real-robot tasks, but the supplied passage does not specify their number.Source: §6 · p. 16
  • Sampling multiple actions is costly for large foundation policies, and the critic used with OpenVLA is not a generalist model.The authors identify reducing sampling cost and developing generalist critics as future directions.Source: §6 · p. 16

A. Environment details

The appendix describes benchmark environments and datasets spanning navigation, kitchen manipulation, and image-based task sequences. It also specifies the CALVIN observation setup and the evaluation protocol.

A. Environment details

  • AntMaze covers medium and large mazes with play and diverse datasets; the diverse trajectories use random starts and goals, while play uses selected locations.The datasets have high state-action coverage.Source: A. Environment details · p. 21
  • FrankaKitchen provides complete, partial, and mixed tele-operated datasets, differing in task completion and the presence of unrelated behavior.The complete dataset is 37 times smaller than the other datasets.Source: A. Environment details · p. 21
  • CALVIN evaluates four sequential robot tasks using image observations only, without proprioception or object states.The tasks are OpenDrawer, TurnonLightbulb, MoveSliderLeft, and TurnonLED.Source: A. Environment details · p. 21
  • Results are averaged over five seeds and 32 evaluations per seed at each evaluation time-step, with scores scaled from 0–4 to 0–100.Plot shading represents standard errors over random seeds.Source: A. Environment details · p. 21

B.1. Details and hyperparameters for PA-RL

The appendix documents PA-RL’s action-optimization and model settings, baseline configurations, and a real-world OpenVLA fine-tuning example. It also reports that some optimization choices vary with the domain and candidate-value variation.

B.1. Details and hyperparameters for PA-RL

  • PA-RL typically samples 32 actions, keeps the top ten, then applies ten Q-function gradient steps with step size 3e-4.The authors note that larger action spaces or near-optimal datasets may call for different settings.Source: §B.1 · p. 22
  • Critic instability prompted distributional critics with HL-Gauss loss for AntMaze and FrankaKitchen, while CALVIN and real-robot critics used MSE.The distributional switch was made when seeds showed exploding Q-values during pre-training.Source: §B.1 · p. 22
  • For AntMaze, argmax candidate selection converged slightly faster; whether to distill only the argmax or all samples depended on candidate Q-value variance.Low variance favored argmax-only distillation; otherwise, sampling from the categorical distribution performed slightly better.Source: §B.1 · p. 22
  • The appendix specifies separate diffusion and transformer architectures, including image-domain ResNet-18 encoders and 128-bin discretization for transformer actions.The transformer uses four layers, hidden size 256, and eight heads; the diffusion policy uses five diffusion steps.Source: §B.1 · p. 23

B.2. Details and hyperparameters for baselines

  • Baseline settings include environment-specific DQL weight searches and shared Cal-QL hyperparameters, alongside configured IDQL, DPPO, and RLPD variants.IDQL’s AntMaze results are not directly comparable to its paper because the datasets differ (-v2 here versus -v0 there).Source: §B.2 · pp. 23–24

C. Filmstrips for Real-World Fine-Tuning of OpenVLA with PA-RL

  • OpenVLA’s success rate on “vegetable to sink” rises from 40% to 70% after 40 minutes of real-robot fine-tuning with PA-RL.The task involves identifying and placing a vegetable amid a distractor object.Source: Figure 8 · p. 24

D.1. Real Robot Fine-tuning on task (b)

The real-robot filmstrip reports that PA-RL changes a pot-placement policy from failing to completing the task after ten online episodes. A separate figure presents learning curves for transformer and diffusion policies paired with different RL methods.

D.1. Real Robot Fine-tuning on task (b)

  • After ten online fine-tuning episodes, PA-RL learns to complete the difficult pot-placement task that the offline initialization fails.The initial policy fails to grasp the pot and gets stuck trying to move it to the sink.Source: Figure 9 · p. 25

D.2. Learning Curves for Auto-Regressive Transformers and IQL with PA-RL

  • Figure 10 presents learning curves for autoregressive transformer policies with PA-RL and Cal-QL, and diffusion policies with PA-RL and IQL.The supplied caption identifies the compared policy and algorithm pairings but does not report their outcomes.Source: Figure 10 · p. 25

D.3. Local and Global Optimization Ablation Experiments

The ablations examine how action optimization choices affect PA-RL performance across tasks and policy settings. They compare optimization steps and sample counts, CEM variants, Bellman targets, and Gaussian policies.

D.3. Local and Global Optimization Ablation Experiments

  • 20% higher final performance on CALVIN follows increased local-optimization gradient steps, while antmaze-large-diverse-v2 reaches 96% success without them.The authors hypothesize that high action coverage lets global optimization with enough base-policy samples recover good actions on antmaze.Source: Figure 11 · p. 26
  • 33% higher normalized score on CALVIN comes from using 32 base-policy samples rather than one; antmaze-large-diverse-v2 performance sharply decreases below five samples.The default PA-RL setting uses 32 samples.Source: §D.3 · p. 26

D.4. CEM Optimizer + Random Initialization Comparisons

  • CEM initially predicts Q-values far above Monte Carlo returns on kitchen-complete-v2; regression-trained critics show larger overestimation than distributional critics.The distributional critic bounds predicted values by design; on kitchen-partial-v2, predicted Q-values become large positive numbers although rewards are non-positive.Source: Figure 14 · p. 27

D.5. CEM Optimizer + pre-trained policy initialization

  • 42% better offline-only performance across tested domains, PA-RL outperforms pre-trained-policy CEM by 66% on kitchen-mixed-v2 and 172% on CALVIN.CEM catches up on antmaze-large-diverse-v2, kitchen-partial-v2, and kitchen-mixed-v2, while PA-RL significantly outperforms it on kitchen-mixed-v2 and CALVIN.Source: Figure 15 · p. 28

D.6. Comparison with computing actions for Bellman Backup with the base policy

  • Using base-policy samples for Bellman targets makes antmaze-large-diverse-v2 fine-tuning initially less stable, despite ultimately reaching similar performance to optimized-action targets.Both approaches start from the same pre-trained critic checkpoints.Source: Figure 16 · p. 28

D.7. Learning curves for Gaussian Policies with PA-RL

  • Gaussian-policy PA-RL performs competitively with standard Cal-QL on kitchen tasks after base-policy behavior cloning and critic pre-training.The distillation learning rate is the only hyperparameter changed for Gaussian policies, set to 3e-4.Source: Figure 17 · p. 29

E. Training time discussion

PA-RL's training and inference costs depend on action sampling, critic optimization, and policy architecture. Caching can improve critic-training efficiency, while inference optimization adds memory and computation requirements.

E. Training time discussion

  • On CALVIN, PA-RL improves at a similar rate per wall-clock time as IDQL but reaches a higher performance value, while DQL remains largely flat.The methods used the same TPU v4 instance type and codebase.Source: Figure 18 · p. 29
  • An action cache can make critic training more efficient when base-policy sampling costs more than T critic gradient steps, as with OpenVLA or large diffusion policies.The cache is updated after each policy-training epoch to avoid staleness and keep the critic aligned with the optimal or on-policy value function.Source: E. Training time discussion · pp. 29–30
  • For diffusion policies, backpropagating critic gradients through the diffusion chain uses a larger memory footprint than PA-RL's DDPM objective by the number of denoising steps.PA-RL's supervised policy-distillation objective can be more efficient than policy improvement through reparameterization.Source: E. Training time discussion · p. 30
  • Inference-time action optimization can increase memory use through multiple policy samples and requires critic-gradient steps, though fewer samples may suffice than during training.Those gradient steps may be cheaper than full Q-function forward passes when image encoding dominates computation; their effects may diminish later in fine-tuning.Source: E. Training time discussion · p. 30
Loading 2412.06685v1…