Expressive Modeling is Insufficient for Offline RL: A Tractable Inference Perspective

The authors examine the problem of offline RL via RvS approaches, and note that one (underexamined) issue is tractability - answering flexible probabilistic queries faithfully. Due to the nature of the dataset generation and stochastic MDP structure, this is particularly difficult. The authors use TPMs for RvS to address this problem, showing that they are better at achieving the requested returns, and perform competitively, particularly with suboptimal data. Furthermore, they are better suited for constrained-action RL.

The paper looks at the "tractability" (ability to answer probabilistic queries) issue in offline RL and uses TPMs (Tractable Probabilistic Models) for solving complex RL tasks.

The paper proposes to use tractable probabilistic models (TPM) in reinforcement learning via supervised learning (RvS) approaches such that the computation of the multi-step value estimate can be done in a more tractable manner (in replacement of the potentially high-variance Monte-Carlo estimate needed in normal autoregressive generative models). The authors showed that through thorough empirical analysis that obtaining the multi-step value estimate accurately and act according to it is crucial in achieving good performance at the inference time, and previous approaches have failed to do so to some extent. In contrast, the use of TPM readily addresses such issue and consequently results in performance improvement on a range of offline reinforcement learning tasks tested (nine original D4RL locomotion tasks, a modified gym-taxi task, and three action-constrained safe RL locomotion tasks).

The paper considers the problem of RL via sequence modeling, an offline RL paradigm given offline trajectories of the form (s, a, r, R) where R denotes the return-to-go. This is achieved by optimizing and sampling from p(a_{t: T} | s_t, E[V_t] >= v).

The authors highlight that there are two main challenges in training such policies. (1) Estimation accuracy: Ability to estimate the expected return of a state and a corresponding action sequence. (2) Tractability issues: Ability to efficiently sample from this distribution. The authors claim that both of these issues are jointly responsible for poor test time performance of RvS based methods.

While fixing the estimation accuracy part is already an active area of research, the authors focus on fixing the tractability part by using Tractable Probabilistic Models (TPMs) for sampling P(V_t >= v | s_t, a_t). The authors experimentally show that their method works well in practice and beats many other offline RL methods on various benchmarks.

Introduction to the TPM and PC.

We thank the reviewers for their valuable feedback and agree that including a section on TPMs and PCs in the main body of the paper would be beneficial. In the following, we provide a detailed elaboration on the fundamental concepts behind TPMs and PCs as well as the training algorithm of PC. Due to space constraints, we added a concise paragraph to Section 2：

• Probabilistic circuits (PCs) are a general and unified computational framework for tractable probabilistic modeling, i.e. a wide class of TPMs can be represented as a PC. Similar to neural networks, PCs are also computation graphs containing millions of computation units, where constraints are imposed on their structures to enable efficient computation of various probabilistic queries. Generally, computing probabilistic queries involves executing specific feedforward/backward algorithms over the PC-defined computation graph.

and put the full explanation in Appendix B and Appendix C. Concretely, we added the following contents:

• B Tractable Probabilistic Models

  A probabilistic model can be seen as a black box to answer queries about the quantities of interest of the joint probability distribution, such as computing marginal probability and performing maximum-a-posterior inference given some evidence. Tractable probabilistic models provide more guarantees when answering probabilistic queries: i) it can perform exact inference to the model’s distribution and no approximations are required. ii) the query computation can be carried out efficiently, that is, in time polynomial (linear in many cases) in the size of the model. Notably, tractability is defined for a family of models only w.r.t. a class of queries and not an absolute property. Indeed, a tractable representation for one query class might not admit polynomial time inference for another query class.

• C.1 We provide an example PC in Figure 4 to illustrate how PC encodes a joint distribution.

• C.2 We use the example PC to show how we execute marginalization inference in practice.

• C.3.1 We introduce the Expectation-Maximization parameter learning algorithm of PCs.

The performance of Trifle.

To better understand the significance of Trifle, we would like to draw the reviewer's attention to the improvements of Trifle compared to TT. That’s because empirically we directly adopt $p_\mathrm{TT} (a_t | s_t)$ as our prior policy $p_\mathrm{GPT} (a_t | s_t)$ in Equation (1), and the large performance gain over TT indicates that the correction term $p_\mathrm{TPM}(V_t>v|s_t, a_t)$ can enhance the inference-time optimality of base policy reliably.

We choose TT as our base policy for its simple yet effective nature as well as its competitive performance. However, we highlight that by replacing $p_\mathrm{GPT} (a_t | s_t)$ (cf. Equation (1)) with better prior distributions (generated by other RvS algorithms), Trifle could still yield a decent performance gain. To verify this, we adopt Trifle on top of the DT algorithm by replacing $p_\mathrm{GPT} (a_t | s_t)$ with $p_\mathrm{DT}(a_t|s_t)$ and run experiments on the Gym-Mujoco benchmarks. Results are provided in Table 3 of the updated paper and are copied below for your convenience. (We did not choose BR-RCRL because no official released code was found.)

As shown in the table, adding the TPM correction term to DT could also boost DT’s performance, which lives up to our expectations. Specifically, the architecture and hyperparameters of Trifle based on DT follow that of Trifle built on TT, while the results of DT are obtained from the original paper (we adopt the same number of seeds as the DT paper). Notably, we only report 5 results for the reason we have to modify the output layer of DT to make it combinable with TPM. Specifically, the original DT directly predicts deterministic action while the modified DT outputs categorical action distributions like TT. In the other unreported 4 environments, the modified DT fails to achieve the original DT scores, and we haven’t figured out the causes since the time is tight.

This paper studies reward-conditioned reinforcement learning, and highlights two problems with methods in this area: (1) estimation accuracy: ability to estimate the expected return of a state and a corresponding action sequence. (2) tractability issues: ability to efficiently sample from this distribution. The paper aims to tackle specifically the second problem by using tractable probabilistic models.

The reviewers liked the usage of TPMs, though there are concerns about evaluating on more complex tasks, such as Antmaze tasks, which were not studied in this paper (though they are a good case of Scenario #2 discussed in the paper). I also think the comparisons lack several other comparisons for an understanding of the approach: rather than running Equation 2 at inference time, why not train the distribution of action tokens with a loss that biases it towards high value? This is a method previously studied in the reward-conditioned policies paper, and regardless, would help us understand the tradeoff between estimation and inference errors.

The paper lacks baselines for instance the DoC baseline (https://openreview.net/references/pdf?id=C0SfmDkou) and other similar methods, that learn latent variables to address issues with DT and TT with stochastic environments. A comparison would be important to have in my opinion.

I also think that using Q-functions or advantage functions (for instance, the "TT+Q" method from the TT paper) should be compared to. To what extent such a better, well estimated RTG values (obtained in the form of Q or advantage functions) remove the need for inference-time tractability? These kinds of questions are important to understand the benefits of this approach in my opinion, especially given that on the tasks studied, the proposed method is not doing singnificantly better than prior approaches.

Overall I think while the ideas of changing inference is interesting and novel, it remains unclear if the method is really improving performance. Without a complete comparisons to other approaches (including those that bias the sampling of RvS type policies using estimates of Q-values or advantages at train time and those which handle stochasticity by learning other conditional variables for TT and DT), it is unclear if this paper can be accepted at this time.

Reject