Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

A Tractable Inference Perspective of Offline RL

Xuejie Liu, Anji Liu, Guy Van den Broeck, Yitao Liang

TL;DR

This work addresses offline RL by questioning the primacy of expressive sequence models and highlighting the role of tractable inference in achieving high returns. It introduces Trifle, a framework that combines Tractable Probabilistic Models with traditional sequence models to compute exact marginals and conditioning probabilities, enabling high-return action sampling even in multi-step, stochastic, or constrained settings. The approach uses per-dimension TPM-corrected sampling, beam search, and adaptive thresholds to bias actions toward high expected returns while staying within the offline data distribution, yielding state-of-the-art results on 7 of 9 Gym-MuJoCo benchmarks and strong performance in stochastic and safe RL tasks. The empirical results demonstrate that tractability can substantially improve inference-time optimality and overall performance, offering a practical path toward inference-aware offline RL; limitations include TPM accuracy dependency and computational considerations. Overall, Trifle advances offline RL by foregrounding tractable probabilistic reasoning as a key component of effective inference-time decision making.

Abstract

A popular paradigm for offline Reinforcement Learning (RL) tasks is to first fit the offline trajectories to a sequence model, and then prompt the model for actions that lead to high expected return. In addition to obtaining accurate sequence models, this paper highlights that tractability, the ability to exactly and efficiently answer various probabilistic queries, plays an important role in offline RL. Specifically, due to the fundamental stochasticity from the offline data-collection policies and the environment dynamics, highly non-trivial conditional/constrained generation is required to elicit rewarding actions. it is still possible to approximate such queries, we observe that such crude estimates significantly undermine the benefits brought by expressive sequence models. To overcome this problem, this paper proposes Trifle (Tractable Inference for Offline RL), which leverages modern Tractable Probabilistic Models (TPMs) to bridge the gap between good sequence models and high expected returns at evaluation time. Empirically, Trifle achieves the most state-of-the-art scores in 9 Gym-MuJoCo benchmarks against strong baselines. Further, owing to its tractability, Trifle significantly outperforms prior approaches in stochastic environments and safe RL tasks (e.g. with action constraints) with minimum algorithmic modifications.

A Tractable Inference Perspective of Offline RL

TL;DR

This work addresses offline RL by questioning the primacy of expressive sequence models and highlighting the role of tractable inference in achieving high returns. It introduces Trifle, a framework that combines Tractable Probabilistic Models with traditional sequence models to compute exact marginals and conditioning probabilities, enabling high-return action sampling even in multi-step, stochastic, or constrained settings. The approach uses per-dimension TPM-corrected sampling, beam search, and adaptive thresholds to bias actions toward high expected returns while staying within the offline data distribution, yielding state-of-the-art results on 7 of 9 Gym-MuJoCo benchmarks and strong performance in stochastic and safe RL tasks. The empirical results demonstrate that tractability can substantially improve inference-time optimality and overall performance, offering a practical path toward inference-aware offline RL; limitations include TPM accuracy dependency and computational considerations. Overall, Trifle advances offline RL by foregrounding tractable probabilistic reasoning as a key component of effective inference-time decision making.

Abstract

A popular paradigm for offline Reinforcement Learning (RL) tasks is to first fit the offline trajectories to a sequence model, and then prompt the model for actions that lead to high expected return. In addition to obtaining accurate sequence models, this paper highlights that tractability, the ability to exactly and efficiently answer various probabilistic queries, plays an important role in offline RL. Specifically, due to the fundamental stochasticity from the offline data-collection policies and the environment dynamics, highly non-trivial conditional/constrained generation is required to elicit rewarding actions. it is still possible to approximate such queries, we observe that such crude estimates significantly undermine the benefits brought by expressive sequence models. To overcome this problem, this paper proposes Trifle (Tractable Inference for Offline RL), which leverages modern Tractable Probabilistic Models (TPMs) to bridge the gap between good sequence models and high expected returns at evaluation time. Empirically, Trifle achieves the most state-of-the-art scores in 9 Gym-MuJoCo benchmarks against strong baselines. Further, owing to its tractability, Trifle significantly outperforms prior approaches in stochastic environments and safe RL tasks (e.g. with action constraints) with minimum algorithmic modifications.
Paper Structure (40 sections, 1 theorem, 19 equations, 5 figures, 11 tables, 3 algorithms)

This paper contains 40 sections, 1 theorem, 19 equations, 5 figures, 11 tables, 3 algorithms.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. Offline Reinforcement Learning.
  4. Tractable Probabilistic Models.
  5. Tractability Matters in Offline RL
  6. Scenario #1
  7. Scenario #2
  8. Exploiting Tractable Models
  9. From the Single-Step Case...
  10. ...To the Multi-Step Case
  11. Practical Implementation
  12. Experiments
  13. Comparison to the State of the Art
  14. Environment setup
  15. Baselines
  16. ...and 25 more sections

Key Result

Theorem 1

Let $a_t := \{a_t^i\}_{i=1}^{k}$ be a set of $k$ boolean variables and $V_t$ be a categorical variables with two categories $0$ and $1$. For some $s_t$, assume the joint distribution over $a_t$ and $V_t$ conditioned on $s_t$ follows a Naive Bayes distribution: ${p} (a_t, V_t \vert s_t) := {p}(V_t \v

Figures (5)

  • Figure 1: An example PC over boolean variables $X_1, \dots, X_4$. Every node's probability given input $x_1 x_2 \bar{x_3} x_4$ is labeled in blue. $p(x_1 x_2 \bar{x_3} x_4) = 0.22$.
  • Figure 2: RvS approaches suffer from inference-time suboptimality. Left: There is a strong positive correlation between the average estimated returns by Trajectory Transformers (TT) and the actual returns in 6 Gym-MuJoCo environments (MR, M, and ME denote medium-replay, medium, and medium-expert, respectively), which suggests that the sequence model can distinguish rewarding actions from the others. Middle: Despite being able to recognize high-return actions, both TT and DT chen2021decision fail to consistently sample such action, leading to bad inference-time optimality; Trifle consistently improves the inference-time optimality score. Right: We substantiate the relationship between low inference-time optimality scores and unfavorable environmental outcomes by showing a strong positive correlation between them.
  • Figure 3: (a) Stochastic Taxi environment; (b) Stochastic FrozenLake Environment; (c) Average returns on the stochastic environment. All the reported numbers are averaged over 1000 trials.
  • Figure 4: Correlation between average estimated returns and true environmental returns for s-Trifle (w/ single-step value estimates), TT, and m-Trifle (w/ multi-step value estimates) in the stochastic Taxi domain. $R$ denotes the correlation coefficient. The results demonstrate that (i) multi-step value estimates (TT and m-Trifle) are better than single-step estimates (s-Trifle), and (ii) exactly computed multi-step estimates (m-Trifle) are better than approximated ones (TT) in stochastic environments.
  • Figure 5: Scaling Curves of Inference Time. (Fix beam width = 32)

Theorems & Definitions (3)

  • Theorem 1
  • Definition 1: Decomposability
  • Definition 2: Smoothness