Source-linked AI summary

OptiDICE: Offline Policy Optimization via Stationary Distribution Correction Estimation

Jongmin Lee, Wonseok Jeon, Byung-Jun Lee, Joelle Pineau, Kee-Eung Kim

arXiv:2106.10783v1cs.LGcs.AI

TL;DR

Offline RL must optimize policies from static data despite distributional shift between target and behavior policies, which can cause overestimated action values in bootstrapping-based methods. OptiDICE estimates stationary distribution corrections through optimization that avoids sampling from the target policy, and it performs competitively with state-of-the-art methods on D4RL benchmarks. Its policy-extraction approach can work poorly when the optimal policy substantially diverges from the data-collection policy.

  • Problem

    Offline RL must address distributional shift between optimized target policies and data-collection policies because it causes overestimated action values in bootstrapping-based methods.

  • Method

    OptiDICE estimates optimal-policy stationary distribution corrections via minimax optimization without sampling from the target policy, reducing the problem to unconstrained convex optimization.

  • Results

    OptiDICE performs competitively with state-of-the-art methods on D4RL offline RL benchmarks.

  • Takeaways & Limitations

    OptiDICE optimizes policy in stationary-distribution space rather than Q-function or policy space, avoiding evaluation of out-of-distribution actions.

  • Takeaways & Limitations

    Policy extraction works poorly when the optimal policy substantially deviates from the data-collection policy because the shared support becomes scarce.

Abstract

from arXiv · show

We consider the offline reinforcement learning (RL) setting where the agent aims to optimize the policy solely from the data without further environment interactions. In offline RL, the distributional shift becomes the primary source of difficulty, which arises from the deviation of the target policy being optimized from the behavior policy used for data collection. This typically causes overestimation of action values, which poses severe problems for model-free algorithms that use bootstrapping. To mitigate the problem, prior offline RL algorithms often used sophisticated techniques that encourage underestimation of action values, which introduces an additional set of hyperparameters that need to be tuned properly. In this paper, we present an offline RL algorithm that prevents overestimation in a more principled way. Our algorithm, OptiDICE, directly estimates the stationary distribution corrections of the optimal policy and does not rely on policy-gradients, unlike previous offline RL algorithms. Using an extensive set of benchmark datasets for offline RL, we show that OptiDICE performs competitively with the state-of-the-art methods.

1. Introduction

Offline RL enables policy optimization from static data, avoiding costly or dangerous exploration, but distributional shift causes overestimated action values in bootstrapping-based methods. OptiDICE addresses this by estimating optimal stationary-distribution corrections and performs competitively on benchmark datasets.

  • Motivation: Offline RL trains policies from static data, avoiding the costly or dangerous exploration required by online reinforcement learning.Offline training addresses settings such as autonomous driving and clinical treatment, where exploratory behavior may be costly or dangerous.
  • Challenge: Distributional shift between the target and data-collection policies causes compounding action-value overestimation for model-free algorithms using bootstrapping.The overestimation arises when bootstrapped targets use predicted values for out-of-distribution actions.
  • Challenge: Prior offline RL methods mitigate overestimation with techniques that encourage action-value underestimation, adding hyperparameters that require proper tuning.
  • Approach: OptiDICE estimates stationary distribution corrections for the optimal policy rather than relying on the problematic evaluation of out-of-distribution actions.The method is presented as an offline RL algorithm designed to avoid the overestimation associated with such evaluations.
  • Evaluation: OptiDICE performs competitively with state-of-the-art methods on D4RL offline RL benchmarks.

2. Background

Offline RL optimizes policies from a static dataset, while prior model-free approaches often estimate Q-values with bootstrapping and out-of-distribution target-policy actions. This can produce unreasonably high Q-values through compounding estimation error.

  • RL formulation: An MDP specifies states, actions, transitions, rewards, an initial-state distribution, and a discount factor, while a policy maps states to action distributions.
  • RL formulation: The RL objective is to learn a policy that maximizes expected rewards under its stationary state-action distribution.
  • RL formulation: The action-value function Qπ and value function Vπ characterize discounted rewards under policy π, with Qπ defined through the Bellman equation.
  • Offline RL: Offline RL optimizes a policy using a static dataset collected before training, represented by its empirical distribution dD.
  • Offline RL: Prior offline model-free methods estimate Q-values, but bootstrapping with out-of-distribution target-policy actions often yields unreasonably high Q-values.

3. OptiDICE

OptiDICE estimates stationary distribution corrections directly, reformulating offline policy optimization into a stable convex problem that avoids evaluating out-of-distribution actions. It then extracts a policy from the corrected distribution, with limitations when target and data supports overlap only sparsely.

  • Stationary-distribution formulation: OptiDICE regularizes stationary distributions against the dataset distribution, balancing reward maximization with penalization of distributional shift.The formulation assumes positive dataset distribution and a strictly convex, continuously differentiable f-divergence generator.
  • Offline optimization: OptiDICE avoids out-of-distribution action evaluation, preventing the compounding error associated with bootstrapped estimates using such actions.Its objective uses expectations over the initial-state distribution and dataset distribution, enabling optimization from offline data.
  • Stationary-distribution formulation: The method estimates optimal stationary distribution corrections through minimax optimization without sampling from the target policy.The correction ratio represents the relationship between the optimal policy’s stationary distribution and the dataset distribution.
  • Stable optimization: A closed-form inner maximizer reduces the nested maximin problem to unconstrained convex minimization, improving numerical stability and convergence reliability.Strong duality permits reordering the optimization from maximin to minimax before deriving the closed-form solution.
  • Policy extraction: In the Four Rooms example, multiplying the learned correction by the dataset distribution produces an estimated stationary distribution for a near-optimal policy.The procedure first estimates ν and w, then forms the corrected distribution d̂π∗(s,a) = dD(s,a)w∗(s,a).
  • Stable optimization: The practical objective is a biased estimate that upper-bounds the target objective, with equality when the MDP is deterministic.The bias arises from nonlinearity and the double-sample problem; the upper-bound property formally justifies the practical objective.
  • Policy extraction: Importance-weighted behavioral cloning can fail when the target policy deviates substantially from the data-collection policy because their support intersection becomes scarce.OptiDICE therefore uses an information projection for policy training, replacing target state marginals with dataset state marginals.

4. Experiments

Experiments evaluate OptiDICE on tabular random MDPs and continuous D4RL benchmarks, showing strong mean performance, robustness, and competitive or state-of-the-art results. The method also exhibits stable optimization relative to approaches involving out-of-distribution actions.

  • OptiDICE uses χ2-divergence for tabular experiments and a softened version for continuous MDPs.
  • Random MDPs: In random MDPs, OptiDICE outperforms baselines in mean and 5%-CVaR when ζ = 0.9, and remains competitive in mean while leading CVaR when ζ = 0.5.The experiments use 10,000 runs and report mean and 5%-CVaR normalized performance.
  • D4RL benchmarks: OptiDICE achieves the highest normalized scores in 7 D4RL tasks against the best model-free baseline among eight algorithms.
  • D4RL benchmarks: OptiDICE achieves state-of-the-art performance on all Maze2D tasks and the best mean performance on 4 Gym-MuJoCo tasks.
  • D4RL benchmarks: On D4RL learning curves, OptiDICE strictly outperforms CQL on 6 tasks and performs on par with CQL on 4 tasks.Results use γ = 0.99 and mean scores with 95% confidence intervals from five runs per task.
  • OptiDICE directly estimates stationary distribution corrections and avoids out-of-distribution actions, simplifying optimization and improving stability over nested optimization.
  • Discount-factor robustness: OptiDICE remains stable as γ approaches 1, whereas CQL becomes unstable because its Q-function values become unbounded.

A. Proof of Proposition 1

The proof establishes a closed-form solution for the inner maximization by applying convex optimization, strong duality, and KKT conditions. Strict convexity of f makes the derivative invertible, enabling the explicit optimizer.

  • The original constrained optimization is convex because its objective is concave in d and its equality and inequality constraints are convex.
  • Strong duality permits changing the optimization order and reparameterizing the stationary distribution with w = d/dD.
  • For fixed ν, the inner maximization over w is solved using primal feasibility, dual feasibility, stationarity, and complementary slackness.
  • Strict convexity of f makes f′ strictly increasing and invertible, yielding the closed-form optimizer through (f′)^−1.
  • The KKT cases distinguish positive and zero optimal weights, with eν(s, a) compared against αf′(0).

B. Proofs of Proposition 2 and Corollary 3

The proofs show convexity of the reduced objective and characterize when the associated bound is tight. The tabular implementation then uses stationary-distribution corrections and second-order optimization to compute the solution.

  • Proposition 2: The reduced objective is convex in ν because the relevant transformed terms are convex and the Lagrangian is affine in its multipliers.
  • Proposition 2: The convexity argument uses strict convexity of f, monotonicity of (f′)^−1, and nonnegative second derivatives.
  • Corollary 3: The resulting inequality always holds, with equality when the MDP transition model is deterministic.
  • Tabular implementation: The tabular procedure constructs an MLE MDP, computes the data-collection stationary distribution, and optimizes stationary-distribution corrections.
  • Tabular implementation: The closed-form inner solution supplies w*, after which Newton’s method uses first and second derivatives to compute an optimal ν*.

D. Proof of Proposition 4

Proposition 4 derives a closed-form solution for the inner maximization with a normalization constraint using KKT conditions. The choice of f-divergence affects numerical stability, motivating fsoft-χ2 for the D4RL experiments.

  • Proof of Proposition 4: KKT conditions characterize the inner maximization through primal feasibility, dual feasibility, stationarity, and complementary slackness.The proof uses strong duality, making these conditions necessary and sufficient for the primal and dual solutions.
  • Proof of Proposition 4: The closed-form solution depends on whether the optimal correction w∗(s, a) is positive or zero.For positive w∗, complementary slackness gives µ∗(s, a) = 0; for zero w∗, dual feasibility determines µ∗(s, a).
  • f-divergence choice: fKL can become numerically unstable because its closed-form solution includes an exponential term.The instability arises for sufficiently large inputs to the exponential.
  • f-divergence choice: fχ2 can suffer from dying gradients because its ReLU-based solution gives zero gradients for negative inputs.This can make training ν slow or fail.
  • f-divergence choice: fsoft-χ2 combines the forms of fKL and fχ2 to avoid both numerical instability and zero-gradient behavior.Its solution is numerically stable for large inputs and always gives non-zero gradients; the authors use it for D4RL experiments.

F.1.2. DATA-COLLECTION POLICY CONSTRUCTION

The data-collection policy is constructed to achieve a specified ζ-optimality relative to uniform-random and optimal policies. The procedure softens the optimal policy and then perturbs optimal-action probabilities until the target performance is reached.

  • Data-collection policy construction: The procedure first computes π∗ and Q∗, initializes πsoft with π∗, and increases temperature in πsoft ∝ exp(Q∗(s, a)/τ).Temperature increases until the softened policy reaches an intermediate performance target.
  • Data-collection policy construction: The softened policy is perturbed by repeatedly discounting optimal-action probabilities at randomly selected states.After each perturbation, the action distribution is normalized.
  • Data-collection policy construction: The perturbation continues until the resulting policy reaches the desired ζ-optimality.The stopping condition compares its value with a ζ-weighted combination of optimal and uniform-policy values.
  • Experimental setup: The tabular experiments compare OptiDICE with BasicRL, RaMDP, RobustMDP, SPIBB, and BOPAH.Hyperparameters for several baselines follow their public-code settings, while OptiDICE uses α = N^-1.

F.2.1. TASK DESCRIPTIONS

The D4RL evaluation uses Maze2D navigation and Gym-MuJoCo continuous-control tasks, with datasets generated from random, medium, replay, and mixed expert/suboptimal trajectories. OptiDICE and baselines are trained with specified neural-network and hyperparameter-search settings.

  • TASK DESCRIPTIONS: Maze2D evaluates navigation to a fixed goal, with maze complexity increasing from umaze to medium to large.The objective is to find a shortest path using previously gathered trajectories.
  • TASK DESCRIPTIONS: Gym-MuJoCo covers hopper, walker2d, and halfcheetah continuous-control tasks.The datasets differ by how the behavior policy was generated.
  • TASK DESCRIPTIONS: The random dataset uses a randomly initialized policy, while the medium dataset uses an early-stopped SAC policy.The medium-replay dataset combines medium trajectories with samples collected during medium-policy training.
  • TASK DESCRIPTIONS: The medium-expert dataset combines equal amounts of expert and suboptimal trajectories.Suboptimal trajectories come from either a uniformly random or medium-performance policy.
  • Hyperparameter and training settings: OptiDICE and comparison methods use multilayer perceptrons, standardized observations and rewards, and task-specific searches over mixture count K and coefficient α.OptiDICE networks use two hidden layers, while the CQL setup uses three hidden layers with 256 units per layer.

G.1. Experimental results for γ = 0.99

Table 3 reports normalized performance for OptiDICE and offline-RL baselines across benchmark tasks. The caption states that OptiDICE achieves the best performance on 6 tasks compared with the listed baselines.

  • Experimental results for γ = 0.99: OptiDICE achieves the best performance on 6 tasks compared with BEAR, BRAC, AlgaeDICE, and CQL.The table also includes CQL results obtained by running the authors’ released code.
  • Experimental results for γ = 0.99: Table 3 reports normalized performance using baseline mean scores from the D4RL benchmark and separately reports CQL (ours).The caption specifies that the baseline methods are BEAR, BRAC, AlgaeDICE, and CQL.

G.2. Experimental results with importance-weighted BC

OptiDICE was evaluated with multiple policy-extraction methods on D4RL, and information projection performed better empirically than importance-weighted BC.

  • Figure 7 compares BC, OptiDICE with importance-weighted BC, and information projection on D4RL.
  • Information projection empirically performs better than importance-weighted BC for OptiDICE policy extraction.
  • For importance-weighted BC, α is selected from {0.0001, 0.001, 0.01, 0.1, 1} using mean performance.
  • The reported importance-weighted BC hyperparameters are summarized in Table 4.

G.3. Experimental results for γ ∈{0.99, 0.999, 0.9999, 1.0}

The experiments evaluate BC, CQL, and OptiDICE variants on D4RL across γ values from 0.99 through 1.0, with CQL unavailable at γ = 1.

  • At γ = 0.99, D4RL performance is compared among BC, CQL, OptiDICE-minimax, and OptiDICE-MSE.
  • At γ = 0.999, D4RL performance is compared among BC, CQL, OptiDICE-minimax, and OptiDICE-MSE.
  • At γ = 0.9999, D4RL performance is compared among BC, CQL, OptiDICE-minimax, and OptiDICE-MSE.
  • At γ = 1, D4RL performance is reported for BC, OptiDICE-minimax, and OptiDICE-MSE because CQL cannot handle this case.
Loading 2106.10783v1…