Source-linked AI summary

PLAS: Latent Action Space for Offline Reinforcement Learning

Wenxuan Zhou, Sujay Bajracharya, David Held

arXiv:2011.07213v1cs.ROcs.AIcs.LG

TL;DR

Offline RL must learn from static datasets despite incomplete support and extrapolation errors caused by out-of-distribution actions. PLAS learns the policy in a CVAE latent action space, implicitly constraining decoded actions while optionally permitting controlled perturbations. Across D4RL benchmarks and a real-robot cloth-sliding task, it achieves competitive performance and supports diverse datasets.

  • Problem

    Offline RL seeks to learn policies from fixed datasets, but out-of-distribution actions cause extrapolation errors and motivate keeping policy actions within dataset support.

  • Method

    PLAS trains a deterministic policy to output latent actions that a pretrained CVAE decoder maps into environment actions, with optional bounded perturbations for controlled out-of-distribution generalization.

  • Results

    PLAS achieves competitive performance across offline RL benchmarks and learns effectively from real-world cloth-sliding data, outperforming previous methods despite its simplicity.

  • Takeaways & Limitations

    PLAS provides a simple way to constrain policies within dataset support while allowing fine-grained control over in-distribution and out-of-distribution generalization.

  • Takeaways & Limitations

    The pessimistic-MDP framing assumes significantly small rewards for unvisited state-action pairs.

Abstract

from arXiv · show

The goal of offline reinforcement learning is to learn a policy from a fixed dataset, without further interactions with the environment. This setting will be an increasingly more important paradigm for real-world applications of reinforcement learning such as robotics, in which data collection is slow and potentially dangerous. Existing off-policy algorithms have limited performance on static datasets due to extrapolation errors from out-of-distribution actions. This leads to the challenge of constraining the policy to select actions within the support of the dataset during training. We propose to simply learn the Policy in the Latent Action Space (PLAS) such that this requirement is naturally satisfied. We evaluate our method on continuous control benchmarks in simulation and a deformable object manipulation task with a physical robot. We demonstrate that our method provides competitive performance consistently across various continuous control tasks and different types of datasets, outperforming existing offline reinforcement learning methods with explicit constraints. Videos and code are available at https://sites.google.com/view/latent-policy.

1 Introduction

Offline reinforcement learning learns policies from static datasets, but out-of-distribution actions create extrapolation errors that require careful policy constraints. PLAS instead learns in a latent action space to constrain actions implicitly and reports strong performance across benchmark and real-robot settings.

  • Motivation: Offline RL trains policies from static datasets, avoiding the millions of online environment interactions often required by real-world robotics.Static-data training also provides flexibility to account for safety and exploration during data collection.
  • Motivation: Out-of-distribution actions cause Q-function extrapolation errors, motivating constraints that keep policy actions within the dataset support without becoming overly restrictive.An overly restrictive policy can degenerate toward behavior cloning.
  • PLAS: PLAS learns a policy in a CVAE latent action space, whose decoder outputs actions while implicitly constraining them to the behavior-policy support.Unlike explicit divergence-based matching, the constraint is imposed through the latent representation.
  • Results: PLAS provides consistently good performance across continuous-control datasets with diverse actions and outperforms previous methods in D4RL and real-robot evaluations.The evaluations include D4RL benchmarks and deformable object manipulation with a physical robot.

2 Related Work

Related offline RL methods address static-data limitations through policy constraints, while imitation learning and generative action models provide distinct neighboring approaches.

  • Offline Reinforcement Learning: Offline RL methods use constraints such as KL-divergence or MMD to reduce out-of-distribution actions when learning from static datasets.These approaches differ in how the constraint enters policy or Q-function optimization.
  • Imitation Learning: Behavior cloning mimics dataset actions but does not reason about which recorded actions are better than others.Imitation-learning methods may also assume expert-policy access or interactive environment data.
  • Generative Models for Actions: Prior work used conditional variational autoencoders to model action sequences or latent action representations outside the offline RL setting.The cited applications include trajectory optimization and multi-agent reinforcement learning.

3 Background

Offline RL optimizes policies from incomplete transition datasets, where extrapolation errors motivate support constraints; PLAS uses latent representations and optional controlled perturbations to address this trade-off.

  • Preliminaries: An MDP is represented by state space S, action space A, transition probabilities P, reward function r, and discount factor γ.The RL objective is defined over returns in this environment model.
  • Offline RL: Offline RL uses a finite fixed transition dataset that may not cover the MDP’s continuous state and action spaces.Policy performance is therefore limited by knowledge inferred from the observed transitions.
  • Policy Constraints: A pessimistic MDP assigns significantly small rewards to unvisited state-action pairs, making its policy performance a lower bound for the true MDP.This provides a surrogate objective focused on transitions visited in the dataset.
  • Policy Constraints: Out-of-distribution actions can make bootstrapped Q-values arbitrarily wrong, with errors accumulating through policy updates.Constraining policies to known actions addresses this error, but the constraint should preserve flexibility within dataset support.
  • Policy Constraints: Existing methods impose policy constraints by behavior-policy sampling or divergence-based regularization, creating practical optimization trade-offs.BCQ samples from the behavior policy, whereas BEAR and BRAC use metrics such as KL-divergence or MMD.
  • Latent Action Space: PLAS maps states to latent actions and decodes them into environment actions, using the latent space as a natural support constraint.Its latent policy can choose latent actions flexibly rather than sampling from a fixed latent range.

4 Method

PLAS learns a policy over a CVAE’s latent action space, then decodes latent actions into environment actions to implicitly constrain behavior. A bounded latent policy can stay near the dataset support, while an optional residual enables controlled out-of-distribution generalization.

  • 4.1 Policy in Latent Action Space (PLAS): The CVAE models the behavior policy by reconstructing state-conditioned actions and learning a decoder pβ(a|s, z).The encoder and decoder parameters are α and β, respectively.
  • 4.1 Policy in Latent Action Space (PLAS): PLAS trains a latent policy that maps states to latent actions and uses a pretrained CVAE decoder to produce environment actions.The CVAE decoder is conditioned on the state and maps latent variables to actions.
  • 4.1 Policy in Latent Action Space (PLAS): Constraining each latent dimension to zi ∈[−σ, σ] keeps latent actions in a high-probability region of the N(0, 1) prior while preserving policy flexibility.The resulting decoded action is likely to have high probability under the behavior policy distribution.
  • 4.2 Generalization out of the dataset: A residual perturbation layer can relax the latent constraint by adding bounded action-space offsets within [−ϵ, ϵ].This keeps the final action close to dataset actions under the L∞ norm and is intended for cases where the Q-function generalizes well.
  • 4.3 Implementation Details: The training procedure first fits the VAE, then jointly trains latent policy and critic networks from dataset transitions using decoded next-state actions.Target networks and soft Clipped Double Q-learning stabilize critic training; the default maximum latent action is 2.
  • 4.3 Implementation Details: Figure 3(c) evaluates episode reward over five evaluation episodes every 150 training steps for PLAS and baselines in the real-robot cloth-sliding task.The figure accompanies a setup with a cloth fixed at the top-left corner and tactile sensing at the grasp point.

5 Experiments

Experiments evaluate PLAS on a physical cloth-sliding task and diverse d4rl locomotion datasets, alongside analysis of learned Q-function errors and optional action perturbations. PLAS outperforms baselines on the robot task and performs similarly to or better than the best baselines on medium-expert and medium-replay datasets.

  • D4RL Datasets: The d4rl evaluation covers Walker2d-v2, Hopper-v2, and Halfcheetah-v2 with random, medium, medium-expert, and medium-replay datasets.The broader benchmark suite also includes Adroit and Franka Kitchen, with locomotion results emphasized in the main text.
  • Real-Robot Experiment: PLAS outperforms all baselines on cloth sliding, with final performance similar to the expert policy.Behavior cloning fails on the low-quality dataset, Offline-TD3 performs worse than behavior cloning, BCQ remains similarly constrained, and BEAR declines after initially reasonable performance.
  • D4RL Datasets: Figure 4 averages each training curve over 3 seeds, with shaded regions showing one standard deviation across seeds.The figure reports training performance for medium-expert and medium-replay locomotion datasets.
  • Q-Function Errors: Q-function analysis uses MSE, overestimated-Q percentage, positive-error mean, and negative-error mean to assess overall error and bias direction.The estimates are compared with empirical returns from transitions during rollouts.
  • D4RL Datasets: PLAS consistently matches or exceeds the best baselines on medium-expert and medium-replay locomotion datasets.These datasets combine diverse states and actions from mixtures of policies, creating both opportunities for learning and challenges for policy constraints.
  • Q-Function Errors: PLAS maintains consistently low MSE and lacks significant overestimation or underestimation, while BEAR overestimates and BRAC underestimates.Although PLAS and BRAC reach similar final Walker2d medium-expert performance, PLAS converges faster and yields a better Q-function by MSE.
  • Perturbation Layer: The optional perturbation layer allows controlled out-of-distribution actions, with ϵ = 0 representing the latent policy alone.The layer is evaluated across different maximum perturbation limits on selected environments; the action range is (−1, 1), so ϵ = 0.5 permits a high perturbation range.

6 Conclusion

The conclusion presents PLAS as a simple offline-RL approach that constrains actions through the dataset’s latent action support. It reports effective learning from real-world data and competitive benchmark performance, while also studying controlled out-of-distribution generalization.

  • 6 Conclusion: PLAS implicitly constrains policy actions to the dataset support without being restricted by the dataset distribution’s density.The method also studies an additional perturbation layer for out-of-distribution Q-function generalization.
  • 6 Conclusion: PLAS learns effectively from real-world cloth-sliding data and achieves competitive performance across offline-RL benchmarks.The conclusion connects these results to more efficient use of data from various sources and future applications of RL on real robots.

A Implementation Details

The implementation trains a CVAE before optimizing a latent policy with its fixed decoder, using a deterministic tanh policy and an optional perturbation layer. Network sizes and training settings vary between MuJoCo datasets and the robot experiment.

  • Algorithm Implementation: The CVAE is trained first, then its fixed decoder is used while training a deterministic latent policy with tanh output.The latent-policy output is scaled by a maximum latent-action hyperparameter, and perturbation is disabled by default.
  • MuJoCo Hyper-parameters: For MuJoCo datasets, actor, critic, and CVAE learning rates are 1e-4, 1e-3, and 1e-4, respectively.The CVAE uses two hidden layers, with (128, 128) used for datasets smaller than 1e6 transitions to prevent overfitting.
  • Robot Hyper-parameters: For the robot experiment, actor and critic learning rates are 3e-4, the CVAE learning rate is 1e-4, and all networks use two hidden layers of size 64.The CVAE is trained for 15000 iterations, with τ = 0.005, λ = 0.75, and batch size 256 also specified.

B D4RL Results

PLAS consistently performs well on D4RL locomotion datasets, especially medium-expert and medium-replay datasets, while baseline methods succeed only on subsets. The max latent action controls how far the latent policy can move from the CVAE prior and behavior-policy support.

  • D4RL benchmark results: PLAS consistently achieves good performance on medium-expert and medium-replay datasets, whereas other baselines work well on some datasets but fail on others.Results are averaged over 3 seeds; PLAS denotes the latent policy without perturbation, and PLAS+P includes the perturbation layer.
  • Dataset construction: The reported experiments use hopper-medium-expert(b), formed by concatenating the medium and expert datasets.The results from Figure 4 and the other appendix experiments are based on this dataset variant.
  • D4RL benchmark results: The D4RL benchmark reports results using both average reward and normalized score across the evaluated datasets.
  • Sensitivity to max latent action: Larger max latent actions may produce out-of-distribution actions because they move the latent policy beyond regions with high probability under the CVAE prior.The max latent action limits latent-policy outputs to keep decoded actions likely to remain within the behavior-policy distribution.

D Ablation Study: Perturbation Layer

The perturbation layer is optional: it sometimes improves performance, but gains are usually small and sufficiently large perturbations typically reduce performance.

  • Perturbation-layer comparison: Adding the perturbation layer sometimes improves performance, but the improvement is usually not significant.The comparison evaluates max perturbations ϵ ∈ {0, 0.01, 0.05, 0.1, 0.2, 0.5}, with ϵ = 0 representing the latent policy alone.
  • Perturbation-layer comparison: Large max perturbations above a task-dependent threshold usually cause performance to drop.
  • Perturbation-layer comparison: The perturbation layer is therefore treated as an optional component of PLAS.
  • Experimental caveat: Walker2d-random experiments are unstable, so their perturbation-layer comparison may not be informative.

E Empirical Analysis on MMD Constraint

Simulations examine why sampled MMD constraints can mismanage out-of-distribution actions. The constraint is only roughly suitable for unimodal Gaussian policies and can prefer the gap between modes in multimodal distributions.

  • Unimodal distributions: For a behavior policy N(0,1) and agent policy N(0,x), the MMD loss keeps the desired threshold behavior only roughly with a Gaussian kernel and large sigma.Ideally, the loss would remain below a threshold for x ≤ 1.
  • Unimodal distributions: Sampled MMD constraints match entire distribution supports and can therefore be overly restrictive.
  • Multimodal distributions: For a behavior policy uniform on [−2,−1] ∪ [1,2] and agent policy N(x,0.5), the minimum MMD loss does not occur at the modes’ centers.With large sigma, the minimum occurs at x = 0, inside the behavior-policy distribution’s gap.
  • Physical-robot context: The robot experiment uses a Sawyer with tactile sensing to slide a gripped cloth along a fixture, but these setup details are separate from the MMD simulations.The task uses tactile readings, force, pose observations, and horizontal and vertical incremental actions.
Loading 2011.07213v1…