A Tractable Inference Perspective of Offline RL
Authors: Xuejie Liu
Anji Liu
Guy Van den Broeck
Yitao Liang
Venue: NeurIPS 2024 poster (2024)
Keywords: Tractable probabilistic models
Probabilistic circuits
Primary area: probabilistic_methods

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. While it is still possible to approximate such queries, we observe that such crude estimates 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 generative models to bridge the gap between good sequence models and high expected returns at evaluation time. Empirically, Trifle achieves $7$ state-of-the-art scores and the highest average scores in $9$ Gym-MuJoCo benchmarks against strong baselines. Further, Trifle significantly outperforms prior approaches in stochastic environments and safe RL tasks with minimum algorithmic modifications.

A Tractable Inference Perspective of Offline RL
Xuejie Liu Anji Liu²* Guy Van den Broeck?, Yitao Liang't 'Institute for Artificial Intelligence, Peking University
*Cmmputer Science Department, University of California, Los Angeles School of Intelligence Science and Technology, Peking University xj xjliu@stu, pku. edu. en, guyvdb@cs. ucla. edu, yitaol@pku. edu.
liuanji@cs. ucla. edu
Abstract
A popular paradigm for offline Reinforcement Learning (RL) tasks to first fit the offline trajectories to 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 conditiional/constriiddd generation is required to elicit rewarding actions. While it is still possible to approximate such queries, we observe that such crude estimates 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 generative models bridge the gap between good sequence models and high expected returns at evaluation time. Empirically, Trifle achieves 7 state-of-the-art scores and the highest average scores in 9 Mo benchmarks against strong baselines. Further, Trifle significantly outperforms prior approaches in stochastic environments and safe RL tasks with minimum algorithmic modifications,
Introduction
Recent advancements in deep generative models have opened up the possibility of solving offline Reinforcement Learning (RL) (27 tasks with sequence modeling techniques (termed RvS approaches). Specifically, we first fit sequence model to the trajectories provided in an offline dataset. During evaluation. the model is tasked to sample actions with high expected returns given the current state. Leveraging modern deep generative models such as GPTss [5] and diffusion models [18), RvS algorithms have significantly boosted the performance on various RL problems
Despite its appealing simplicity, is still unclear whether expressive modeling alone guarantees good performance of RyS algorithms, and if so, on what types of environments. This paper discovers that many common failures of RvS algorithms are not caused by modeling problems. Instead, while useful information is encoded in the model during training, the model is unable elicit such knowledge during evaluation. Specifically, this issue is reflected in two aspects: (i) inability to accurately estimate the expected return of a state and corresponding action sequence be executed given near- cerrect learned transition dynamics and reward functions; (ii) even when accurate return estimates exist in the offline dataset and are learned by the model, could still fail to sample rewarding actions during evaluation At the heart of such inferior evaluation-tin time performance is the fact that highly
"qqual contribution
Cooresponning author
Our code is available at
Botth observations are supported by empirical evidence as illustrated in Section
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

non-trivial conditional generation is required to stimulate high-return actions B3 3. Therefore, other than expressiveness. the ability to efficiently and exactly answer various queries (e.g.. computing the expected returns), termed tractability, plays an equally important role in RvS approaches.
Having observed that the lack of tractability is an essential cause of the underperformance of RvS algorithms, this paper studies whether we can gain practical benefits from using Tractable Probabilistic Models (TPMs) [35 7 23y, which by design support exact and efficient computation of certain queries? ? We answer the question in its affirmative by showing that we can leverage a class of TPMs that support computing arbitrary marginal probabilities to significantly mitigate the inference-time suboptimality of RvS approaches. The proposed algorithm Trifle (Tractable Inference for Offline RL) has three main contributions:
Emphasizing the important role of tractable models in offline RI. This is the first paper that demon- strates the possibility of using TPMs on complex offline RL tasks. The superior empirical performance of Trifle suggests that expressive modeling is not the only aspect that determines the performance of RyS algorithms, and motivates the development of better inference-aware RvS approaches.
Competitive empirical performance. Compared against strong offline RL baselines (including RvS, imitation learning, and offline temporal-diifereenee algorithms), Trifle achieves the state-of-the-art result on 7 out of 9 Gym- MuJoCo benchmarks [14] and has the best average score.
Generalizability to stochastic environments and safe-RL tasks. Trifle can be extended to tackle stochastic environments as well as safe RI tasks with minimum algorithmic modifications. Specifi- cally, we evaluate Trifle in 2 stochastic OpenAl-Gym [ environments and MuloCo environments, and demonstrate its superior performance against all baselines.
2 Preliminaries
Offline Reinforcement Learning. In Reinforcement Learning (RL), an agent interacts with an environment that is defined by a Markov Decision Process (MDP) (S, A, R. P. do) to maximize its cumulative reward. Specifically, the is the state space, .A is the action space, R S x A + R is the reward function, P: S x A + S is the transition dynamics, and do is the initial state distribution. Our goal is to learn policy π (a|s) that maximizes the expected return E E 7/r,]. where 7 E (0, 1] is a discount factor and T is the maximum number of steps.
Offline RL [27 aims to solve RL problems where we cannot freely interact with the environment. Instead, we receive dataset of trajectories collected using unknown policies. An effective learning paradigm for offline RI. is to treat it as sequence modeling problem (termed RL via Sequence Modeling or RvS methods) 20 13 Specifically, we first learn a sequence model on the dataset, and then sample actions conditioned on past states and high future returns. Since the models typically do not encode the entire trajectory, an estimated value or return-to-go (RTG) (i.e., the Monte Carlo estimate of the sum of future rewards) is also included for every state-action pair, allowing the model to estimate the return at any time step.
Tractable Probabilistic Models. Tractable Probabilistic Models (TPMs) are generative models that are designed efficiently and exactly answer wide range of probabilistic queries [35 37). One example class of TPMs is Hidden Markov Models (HMMs) (36), which support linear time (w.r.t. model size and input size) computation of marginal probabilities and more. Probabilistic Circuits (PCs) 77] are a general class of TPMs. As shown in Figure PCs consist of input nodes @ that represent simple distributions (e.g., Gaussian, Categorical) over one or more variables as well as sum O and product @ nodes that take other nodes as input and gradually form more complex distributions. Specifically, product nodes model factorized distributions over their inputs, and sum nodes build weighted mixtures (mixture weights are labeled on the corresponding edges in Fig. over their input distributions. Please refer to Appx. B for more detailed introduction to PCs.
1.0
0.0
0.0
0.1
X
0.22
0.0
10
1.0
1.0
0.3
0.3
X.
Figure I: An example PC over boolean
variables X1, XI...., X4. Every node's s probability given input 212222344 is labeled in blue. ((1))225344) = 0.22.

---

Normal Rorzalized Actual Returrs
Inference-tine Optimality Score
0.6
Norma Ized Actual Return
100
60
wakerid HR wakerd
MiPhhstah-
walker2d-MR
Iker2d-M
walker2d-ME
60
ba Ferestah
we kardd NI
In Erreeth ME
40
teporria
en trGdN
0.8
20
Avorage Estimated Returns
walker/d.M
hopper-M
halfcheetah:M
Inforence-time Optimality
Figure 2: RvS approaches suffer from inference time suboptimality. Left: There 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 [6) 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-tine optimality scores and unfavorable environmental outcomes by showing strong positive correlation between them.
Recent advancements have extensively pushed forward the expressiveness of modern PCs 30 31 9
leading to competitive likelihoods on natural image and text datasets compared to even strong Variational Autoencoder (43) and Diffusion model (22) baselines. This paper leverages such advances and explores the benefits brought by PCs in offline RL tasks.
Tractability Matters in Offline RL
Practical RvS approaches operate in two main phases training and evaluation. In the training phase, a sequence model is adopted to learn a joint distribution over trajectories of length T: {(st, 0e; rt,, RTG,)) ) During evaluation, at every time step I, the model is tasked to discover an action sequence ar.7 = a, } (or just a,) that has high expected return as well as high probability in the prior policy p( ((a:Tl St). which prevents it from generating out-of-distribution actions:
1 Sp(err (st)
if Ev.-p((le. [V] 2 v,
p(a.7|s: 8t, B 2 v) = 7 10
(I)
otherwise,
where Z is normalizing constant, V/2 is an estimate of the value at time step t, and e is pre-defined scalar chosen to encourage high-return policies. Depending on the problem, V, could be the labeled RTG from the dataset (e.g., RTG;)) or the sum of future rewards capped with a value estimate (e.g..
ET=1 r, + RTGT) 13 20
The above definition naturally reveals two key challenges in RvS approaches: (i) training-time optimality (i.e., "expressivity") how well can we fit the offline trajectories, and (ii) inference-time optimality: whether actions can be unbiasedly and efficiently sampled from Equation (T). While extensive breakthroughs have been achieved to improve the training- time optimality 6 20 it
remains unclear whether the non-trivial constrained generation task of Equation ( hinders inference- time optimality. In the following, we present two general scenarios where existing RvS approaches underperform as a result of suboptimal inference-time performance. We attribute such failures the fact that these models are limited to answering certain query classes (e.g.. autoregressive models can only compute next token probabilities), and explore the potential of tractable probabilistic models for offline RL tasks in the following sections.
Scenario #1 We first consider the case where the labeled RTG belongs to a (near- (near-)optimas Joptimal policy. In this case, Equation 1 can be simplified to p(a. |st. [[[] 2 v) (choose V, == RTG,) since one- step optimality implies multi-step optimality. In practice, although the RTGs are suboptimal, the predicted values often match well with the actual returns achieved by the agent. Take Trajectory
'To minimize computation cost, we only model truncated trajectories of length K (K 2) in practice.
3

---

Transformer (TT) [20] as an example, Figure 2 (left) demonstrates a strong positive correlation between its predicted returns (x-axis) and the actual cumulative rewards (y-axis) on six MuJoCo 42 benchmarks, suggesting that the model has learned the "goodness" of most actions. In such cases, the performance of Rvš algorithms depends mainly on their inference-time optimality, i.e., whether they can efficiently sample actions with high predicted returns. Specifically, let ar be the action taken by a RvS algorithm at state $;, and R; := ERRTG, is the corresponding estimated expected value. We define a proxy of inference-time optimality as the quantile value of Re in the estimated state-conditioned value distribution p( Vilss) ( The higher the quantile value, the more frequent the RyS algorithm samples actions with high estimated returns.
We evaluate the inference-timet time optimality of Decision Transformers (DT) 6 and Trajectory Trans- formers (TT) [20]. two widely used RvS algorithms, on various environments and offline datasets from the Gym-MuJoCo benchmark suite [14. As shown in Figure 2 (middle) the inference-time optimality is averaged (only) around 0.7 (the maximum possible value is 1.0) for most settings. And these runs with low inference-time optimality scores receive low environment returns (Fig. 2 (right)).
Scenario #2 Achieving inference-time optimality becomes even harder when the labeled RTGs are suboptimal (e.g,, they come from random policy). In this case, even estimating the expected future return of an action sequence becomes highly intractable, especially when the transition dynamics of the environment are stochastic. Specifically, to evaluate a state-action pair (5t, az),
since RTG, is uninformative, we need to resort to the multi-step estimate ViTh := Σd- ry RTG, (t' >> t), where the actions ore are jointly chosen to maximize the expected return. Take autoregressive models as an example. Since the variables are arranged following the sequential order .,.,
St, a₁ Tr, RTG,, St+1...., we need to explicitly sample $t+1. before proceed to compute the
rewards and the RTG in Vith, In stochastic environments, estimating [[[""]] could suffer from high variance as the stochasticity from the intermediate states accumulates over time.
As we shall illustrate in Section 6.2] compared to environments with near-deterministic transition dynamics, estimating the expected returns in stochastic environments using intractable sequence models is hard, and Trifle can significantly mitigate this problem with its ability to marginalize out intermediate states and compute E[V?") efficiently and exactly.
4 Exploiting Tractable Models
The previous section demonstrates that apart from modeling, inference-time suboptimality is an- other key factor that causes the underperformance of RvS approaches. Given such observations, a natural follow-up question is whether/how more tractable models can improve the evaluation-time performance in offline RL tasks? While there are different types of tractabilities (i.e., the ability to compute different types of queries), this paper focuses on studying the additional benefit of exactly computing arbitrary marginal/condition probabilities. This strikes a proper balance between learning and inference as we can train such tractable yet expressive model thanks to recent developments in the TPM community 19 30). Note that in addition to proposing a competitive RvS algorithm, we aim highlight the necessity and benefit of using more tractable models for offline RL tasks, and encourage future developments on both inference-aware RyS methods and better TPMs. As direct response to the two failing scenarios identified in Section B] we first demonstrate how tractability could help even when the labeled RTGs are (near-)optings (Sec. 4.1 We then move on to the case where we need to use multi-step return estimates to account for biases in the labeled RTGs (Sec. 4.2)
4.1 From the Single-Step Case...
Consider the case where the RTGs are optimal. Recall from 3 our goal is to sample actions from p(ar|se, [[]] 2 v) (where V, = RTG,). Prior works use two typical ways to approximately sample from this distribution. The first approach directly trains model to generate return-conditioned actions: p(a. st. RTG,) 6). However, since the RTG given a state-action pair is stochastic sampling
'Due to the large action space, it is impractical to compute p( (Vi = E., p(Vi]st, a.) p(as(st). Instead, the following illustrative experiments, we train an additional GPT model p p(Vils.) using the offline dataset. "This is true unless (i) the policy that generates the offline dataset is deterministic, (ii) the transition dynamics is deterministic, and (iii) the reward function is deterministic.

---

from this RTG- CTG-conditi policy could result in actions with small probability of getting a high
return, but with a low expected return 32 3
An alternative approach leverages the ability of sequence models to accurately estimate the expected return (i.e., (.) of state-action pairs (20). Specifically, we first sample from prior distribution p((e)st), and then reject actions with low expected returns. Such rejection sampling-base typically work well when the action space small (in which we can enumerate all actions) or the dataset contains many high-rewarding trajectories (in which the rejection rate is low). However, the action could be multi-dimensian trajectories in practice, rendering the inference- iifrrencetiime optimality score low (cf. Fig. 2).
methods
and the dataset typically contains many more low-return
Having examined the pros and cons of existing approaches, we left with the question of whether tractable model can improve sampled actions (in this single-step case). We answer it with a mixture of positive and negative results: while computing p(ay |8t, [V] 2 u) is NP-hard even when p(a,, Vilss) [s,
follows a simple Naive Bayes distribution, we can design an approximation algorithm that samples high-return actions with high probability in practice. We start with the negative result.
Theorem 1. Let a {ai} be set of k boolean variables and Vg be a categorical variables with two categories O and 1. For some St, assume the joint distribution over 0t and V, conditioned on St follows a Naive Bayes distribution: p(a,, Velss) == pl V.lss) IT= p(aj) VadiVV,,, V1, 9t), where aj denotes the tth variable of at. Computing any marginal over the random variables is tractable yet conditioning on the expectation p(a,) St, 2[V]]2v) v) Na-hard.
The proof is given in Appx. A While it seems hard to directly draw samples from p(ar|st, E[V] v). we propose to improve the aforementioned rejection sampling-basee method by adding a correction term to the original proposal distribution p(ar) (st) to reduce the rejection rate. Specifically, the prior is often represented by an autoregressive model such as GPT: pGrt(ar|s) == II: PGPT (atls, a."). where i is the number of action variables and of is the ith variable of ag. We propose to sample every dimension of autoregressively following:
1
Vi E {l.....k} p(a"is., az af::) = z Papt(aalss az5) PTPM 2 v/s,, or
(2)
where Z is a normalizing constant and PrPM (V; 2 2V,s,,, is correction term that leverages the ability of the TPM to compute the distribution of V, given incomplete actions (i.e., evidence on a subset of action variables). Note that while Equation (2) is mathematically identical to p(ae|se, V, 2 v) when = PrPm = PGPT, this formulation gives us the flexibility to use the prior policy (i.e., PGPT(a)) of (²)) represented by more expressive autoregressive generative models.
As shown in Figure 2 (middle), compared to using p(aels:) (as done by TT), the inference-time optimality scores increase significantly when using the distribution specified by Equation (2 (as done by Trifle) across various Gym-MuJoCo benchmarks.
4.2 ...To the Multi-Step Case
Recall that when the labeled RTGs are suboptimal, our goal is to sample from p(ace |s,, (sc, [V | v).
where === r- RTG, is the multi-step value estimate. However, as shown in the second scenario in Section 3] it is hard even to evaluate the expected return of an action sequence due to the inability to marginalize out intermediate states $1410'. Empowered by PCs, we can solve this problem by computing the expectation efficiently as it can be broken down into computing conditional probabilities p(r,|st, axa' [[ (tST<( <t) and p(RTG (st, dww) (see Appx. C.2 for details):
B[V.T]- Σ T=t E, (p( [r=] + (R [RTG.:]-
(3)
We are now left with the same problem discussed the single-step case how to sample actions with high expected returns (i.e., [V. Similar to Equation (2). we add correction terms that bias the action (sequence) distribution towards high expected returns. Specifically, we augment the original action probability II p(a, Isc. acr) with terms of the form p(V7" 2 20/se, agr) This leads to:
plaro (): II p(a, 1sc, acriv).
5

---

Table 1: Normalized Scores on the standard Gym-MuJoCo benchmarks. The results of Trifle averaged over 12 random seeds (For DT--bsse and DT-TTrifle, we adopt the same number of seeds as 6 Results of the baselines are acquired from their original papers.
TT(+))
DT
Dataset
Environment
DD
IQL
CQL
%BC TD3(+BC)
base
Trifle
base
Trifle
base
Trifle
Med-Expert
HalfChectah
95.0
95.1=03
82.3+61
89,9
86,8#13
91.9119
90.6
86.7
91.6
92.9
90.7
Med-Expert Hopper
110.0s27
113.0
74.7463
78,5
101.6:11
/
111.8
91.5
105.4
110.9
98.0
Med-Expert Walker2d
101.9268
109.3
109.3223
109.6:02
108.
1
108.6205
108.8
109.6
108.8
109.0
110.1
Mediam
HaltCheetah
46.9204
49,5=n02
48.7aa0
48.9
42.6
44.2
49.1
47.4
44.0
42.5
48.3
Medium
Hopper
61.1
67.1=43
55.2438
57.8=19
67.6=10
/
79.3
66.3
58.5
56.9
59.3
Medium
Walker2d
79.0azs
83,1
82,24a5
84.7
74
81,3
82.5
78.3
72.5
75.0
83.7
Med-Replay
HalfChectah
41.9
45,0=03
48.2
48,9=03
36,6#n8
39.2
39.3
44.2
45.5
40.6
44.6
Med-Replay Hopper
91.5a1
97.8
83.4
87.6.6.1
82.7
/
100.0
94.7
95.0
75.9
60.9
Med- Replay Walker2d
82.6159
88.3=38
84.6144
90.6=43
66.6230
73.
75.0
73.9
77.2
62.5
81.8
Average Score
78.9
83.1
74.3
77.4
74.7
81.8
77.0
77.6
74.0
75.3
where p(a, St, acr: 1)) x p(a, act) p(V. 2 e/st agr), Ocr and agr represent at7-1 and agr. respectively In practice, while we compute p(V. 2 u/8t as²) using the PC, p(a, (8,, (cr)) Ex, [p(a, Sgr, acr)] can either be computed exactly with the TPM or approximated (via Monte Carlo estimation over St+1:r) using an autoregressive generative model. In summary, we approximate samples from p(arar B[V]] [ v) by first sampling from p(ars (s: v), and then rejecting samples whose (predicted) expected return is smaller than v.
5 Practical Implementation
The previous section has demonstrated how to efficiently sample from the sepected-value-codeee policy (Eq. (I). Based on this sampling algorithm, this section further introduces the proposed algorithm Trifle (Tractable Inference for Offline RL). The high-level idea of Trifle is to obtain good action (sequence) candidates from p(a, (t. [E[V] > v). and then use beam search to further single out the most rewarding action. Intuitively, by the definition in Equation (T. the candidates are both rewarding and have relatively high likelihoods in the offline dataset, which ensures the actions within the offline data distribution and prevents overconfident estimates during beam search.
Beam search maintains set of N (incomplete) sequences each starting as an empty sequence. For ease of presentation, we assume the current time step is O. At every time step f, beam search replicates each of the N actions sequences into À E Z copies and appends an action d, to ev- ery sequence. Specifically, for every partial action sequence act, we sample an action following p(ar|so. act, E[V.] 2 v). where V, can be either the single-step the multi-step estimate depending on the task. Now that we have A trajectories in total, the next step is to evaluate their expected return, which can be computed exactly using the PC (see Sec. 4.2 The N-best action sequences are kept and proceed to the next time step. After repeating this procedure for H time steps, we return the best action sequence. The first action in the sequence is used to interact with the environment. Please refer to Appx. C for detailed descriptions of the algorithm.
Another design choice is the threshold value v. While it is common to use a fixed high return throughout the episode, we follow [12 and use an adaptive threshold. Specifically, at state 8t, we choose v to be the c-quantile value of p(V]ssz), which is computed using the PC.
6 Experiments
This section takes gradual steps to study whether Trifle can mitigate the inference- time suboptimality problem in different settings. First, in the case where the labeled RTGs are good performance indicators (i.e., the single-step case), we examine whether Trifle can consistently sample more
'We approximate p(V 2 ulss, agr) by assuming that the variables frey..., 11 RTG,) are indepen- dent. Specifically, we first compute (p(r, 8&s, ag. - and p(KTG, (s, agv) ). and then sum up the random variables assuming that they are independent. This introduces no error for deterministic environments and remains a decent approximation for stochastic environments.

---

rewarding actions (Sec. 6.1). Next, we further challenge Trifle in highly stochastic environments, where existing RvS algorithms fail catastrophically due to the failure to account for the environmental randomness (Sec. 6.2p. Finally, we demonstrate that Trifle can be directly applied to safe RL tasks (with action constraints) by effectively conditioning on the constraints (Sec. 6.33. Collectively, this section highlights the potential of TPMs on offline RL tasks.
6.1 Comparison to the State of the Art
As demonstrated in Section 3 and Figure although the labeled RTGs in the Gym-MuloCo [14 benchmarks are accurate enough to reflect the actual environmental return, existing RvS algorithms fail to effectively sample such actions due to their large and multi-dimensional action space. Figure 2 (middle) has demonstrated that Trifle achieves better inference-ting optimality. This section further examines whether higher inference-time optimality scores could consistently lead to better perfor- mance when building Trifle on top of different Rvš algorithms, i.e., combining PTPM (cf. Eq. with different prior policies PGPT trained by the corresponding RvS algorithm.
Environment setup The Gym-MuloCo benchmark suite collects trajectories in locomotion environments (HalfCheetah, Hopper, Walker2D) and constructs 3 datasets (Medium-Exper Expert, Medium, Medium- Replayy) for every environment, which results in 3 x the main difference between the datasets is the quality of its trajectories. Specifically, the dataset "Medium" records I million steps collected from a Soft Actor-Critic (SAC) [66 agent. The "Medium- Replay" dataset adopts all samples in the replay buffer recorded during the training process of the SAC agent. The "Medium-Expert" dataset mixes million steps of expert demonstrations and million suboptimal steps generated by partially trained SAC policy or random policy. The results are normalized such that a rell-trained SAC model hits 100 and a random policy has a score.
tasks. For every environment,
Baselines We build Trifle on top of three effective RvS algorithms: Decision Transformer (DT) [6.. Trajectory Transformer (TT) [20 as well as its variant TT(+Q) where the RTGs estimated by summing up future rewards in the trajectory are replaced by the Q-values generated by well-trained IQL agent (24). In addition to the above base models, we also compare Trifle against many other strong basclines: (i) Decision Diffuser (DD) (.) which is also competitive RvS method; (ii) Offline
TD learning methods IQL 24 and CQL [26]; (iii) Imitation learning methods like the variant of BC
34
which only uses 10% of trajectories with the highest return, and TD3(+BC) [S]
Since the labeled RTGs are informative enough about the "goodness" of actions, we implement Trifle by adopting the single-step value estimate following Section 4.1] where we replace PGPT with the policy of the three adopted base methods, i.e., prr((a (8). PrT(+ (ar(sx) and prr (a:)))))
Empirical Insights Results are shown in Table First, to examine the benefit brought by TPMs. we compare Trifle with three base policies, as the main algorithmic difference is the use of the improved proposal distribution (Eq. (2) for sampling actions. We can see that Trifle not only achieves a large performance gain over TT and DT in all environments, but also significantly outperforms TT(+Q) where we have access to more accurate labeled values, indicating that Trifle enhance the inference-time time optimality of base policy reliably and benefit from any improvement of the training-time optimality. See Appx. E. for more results and ablation studies.
Moreover, compared with all baselines. Trifle achieves the highest average score of 83.1. It also succeeds in achieving 7 state-of-the-art scores out of 9 benchmarks. We conduct further ablation studies on the rejection sampling component and the adaptive thresholding component (i.e., selecting v) in Appx. F
6.2 Evaluating Trifle in Stochastic Environments
This section further challenges Trifle on stochastic environments with highly suboptimal trajectories as well as labeled RTGs in the offline dataset. As demonstrated in Section 3] in this case, it is even
"When implementing DT Triffi, we have to modify the output layer of DT 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 3 unreported hopper environments, the modified DT fails to achieve the original DT scores.
7

---

Frozenl
Methods
Taxi
0.3
0.5
0.7
m-Trifle
-57
0.61
0.59
0.37
8- s-Trifle
-99
0.62
0.60
0.34
TT (0
-882
0.63
0.25
0.12
DT 6
-388
0.51
0.32
0.10
DoC. 47
-146
0.58
0.61
0.23
(a)
(b)
(c)
Figure 3: (a) Stochastic Taxi environment; (b) Stochastic FrozenLake Environment; (c) Average
returns on the stochastic environment. All the reported numbers averaged over 1000 trials.
hard to obtain accurate value estimates due to the stochasticity of transition dynamics. Sectioo 4.2 demonstrates the potential Trifle to more reliably estimate and sample action sequences under suboptimal labeled RTGs and stochastic environments. This section examines this claim by comparing the five following algorithms:
(i) Trifle that adopts V2 = RTG, (termed single-step Trifle or s-Trifle): (ii) Trifle equipped with V
(ii) TT 20) T RTG, (termed multi-step Trifle m-Trifle; see Appx. E.2 for additional details); (iv) DT (6) (v) Dichotomy of Control (DoC) [47. an effective framework to deal with highly stochastic environments by designing a mutual information constraint for DT training, which is a representative baseline while orthogonal to our efforts.
We evaluate the above algorithms we choose the Taxi benchmark for a detailed analysis of whether and how Trifle could overcome the challenges discussed in Section 4.2 Among the first four algorithms, s-Trifle and DT do not compute the "more accurate" multi-step value, and TT approximates the value by Monte Carlo samples. Therefore, we expect their relative performance to be DT s-Trifle < TT m-Trifle.
stochastic Gym environments: Taxi and FrozenLake. Here
Environment setup We create a stochastic variant of the Gym-Taxi Environment 11 As shown Figure Sal a taxi resides grid world consisting of a passenger and destination. The taxi is tasked to first navigate to the passenger's position and pick them up, and then drop them off at the destination. There are 6 discrete actions available at every step: (i) 4 navigation actions (North, South, East, West), (ii) Pick-up, (iii) Drop-off. Whenever the agent attempts to execute navigation action, it has 0.3 probability of moving toward a randomly selected unintended direction. At the beginning of every episode, the location of the taxi, the passenger, and the destination initialized randomly. The reward function is defined as follows: (i) -1 for each action undertaken; (ii) an additional +20 for successful passenger delivery; (iii) -4 for hitting the walls; (iv) -5 for hitting the boundaries; (v) -10 for executing Pick-up or Drop-off actions unlawfully (c.g., executing Drop-off
randomly
when the passenger is not in the taxi).
Following the Gym-MuloCo benchmarks, we collect offline trajectories by running Q-learning agent 45 in the Taxi environment and recording the first 1000 trajectories that drop off the passenger successfully, which achieves an average return -128.
We first examine the accuracy of estimated returns for Trifle, m- Trifle, and TT. DT excluded since it does not explicitly estimate the value of action sequences. Figure Figure illlustrats the correlation between predicted and ground-truth returns of the three methods. First, s-Trifle performs the worst since it merely uses the inaccurate RTG, to approximate the ground-truth return. Next, thanks its ability to exactly compute the multi-step value estimates, m-Trifle outperforms TT. which approximates the multi-step value with Monte Carlo samples.
Empirical Insights
We proceed to evaluate their performance in the stochastic Taxi environment. As shown in Figure the relative performance of the first four algorithms is DT < TT < s-Trifle < m- Trifffe, which largely aligns with the anticipated results. The only "surprising" result is the superior performance of s- Trifle compared to TT. One plausible explanation for this behavior that while TT can better estimate the given actions, it fails to efficiently sample rewarding actions.
Notably, Trifle also significantly outperforms the strong baseline DoC, demonstrating potential in handling stochastic transitions. To verify this, further evaluate Trifle on the stochastic FrozenLake

---

s-Trifle
TT
m-Trifle
Actual Returns
Actual Returns
Actual Returns
50
-50
R=0.41
100
R=
R=0,67
-100
-150
-150
200
150
-200 -100
250
-200
-100 0
-50
-25
Average Estimated Returns
Average Estimated Returns
Average Estimated Returns
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. I 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.
10 the experiment design follows the DoC paper [47]. For data collection, we perturb the policy of a well-trained DQN (with an average return of 0.7) with the e-greedy strategy. Here e is a proxy of offline dataset quality and varies from 0.3),
environment. Apart from fixing the stochasticity level p
0.3 0.7. As shown in Figure Be when the offline dataset contains many successful trials all methods perform closely to the optimal policy. As the rollout policy becomes more suboptimal (with the increase of c). the performances of DT and TT drop quickly, while Trifle still works robustly and outperforms all baselines.
6.3 Sefion-paac--Constrannee Cym-MuJoCo Variants
This section demonstrates that Trifle can be read- ily extended to safe RL tasks by leveraging TPM's ability to compute conditional probabil- ities. Specifically, besides achieving high ex- pected returns, safe RL tasks require additional constraints on the action states to be satisfied. Therefore, define the constraint as c, our goal sample actions from p( (a) st. [E[] 2 e,c), which can be achieved by conditioning on e in the candidate action sampling process.
Table 2: Normalized Scores on the Action-Space- Constrained Gym-MuJoCo Variants. The results of Trifle and TT are both averaged over 12 random seeds, with mean and standard deviations reported.
Dataset
Eavironment
Trille
TT
Med-Expert Med-Expert Med-Expert
Halfcheetah
81.9
77.8154
Hopper
109.6124
100.0142
Walker2d
105.1
103.6149
Environment setup In MuJoCo environments, each dimension of ag represents the torque applied on certain rotor of the hinge joints at timestep We consider action space constraints in the form of "value of the torque applied to the foot rotor S A", where A = 0.5 is threshold value, for three MuJoCo environments: Halfcheetah, Hopper, and Walker2d. Note that there are multiple foot joints in Halfcheetah and Walker2d, so the constraint is applied to multiple action dimensions [1| For all settings, we adopt the "Med "xpe-Expert" offline dataset introduced in Section 6.1
Empirical Insights The key challenge in these action-constrained tasks is the need to account for the constraints applied to other action dimensions when sampling the value some action variable. For example, autoregressive models cannot take into account constraints added variable af when sampling at. Therefore, while enforcing the action constraint is simple, it remains hard to simultaneously guarantee good performance. As shown in Table owing to its ability to exactly condition on the action constraints, Trifle outperforms TT significantly across all three environments.
When the agent takes action, it has probability p of moving in the intended direction and probability 0.5(1 - p) of slipping to either perpendicular direction. 1l We only add constraints to the front joints in the Halfcheetah environment since the performance degrades significantly for all methods if the constraint added to all foot joints.

---

7 Related Work and Conclusion
In offline reinforcement learning tasks, our goal is to utilize a dataset collected by unknown policies to derive an improved policy without further interactions with the environment. Under this paradigm. we wish to generalize beyond naive imitation learning and stitch good parts of the behavior policy. To pursue such capabilities, many recent works frame offline RL tasks as conditional modeling problems that generate actions with high expected returns [6,
[ 22] or its proxies such as immediate rewards Recent advances in this line of work can be highly credited the powerful expressivity of modern sequence models, since by accurately fitting past experiences, we can obtain 2 types of information that potentially imply high expected returns: transition dynamics of the environment, which serves as a necessity for plamning in model-bases fashion (8). (ii) a decent policy prior which acts more reasonably than random policy to improve from 20
25 39 40
While prior works on model- based RL (MBRL) also leverage models of the transition dynamics and the reward function (21] 17 22 RyS approaches focus more on directly modeling the correlation between actions and their end-performance. Specifically, MBRL approaches focus on planning only with the environment model. Despite being theoretically appealing. MBRL requires heavy machinery to account for the accumulated errors during rollout [19 41 and out- out-of-distribbtiin problems [48] (38). All these problems add a significant burden on the inference side, which makes MBRL algorithms less appealing in practice. In contrast, while RvS algorithms can mitigate this inference-time burden by directly learning the correlation between actions and returns, the suboptimality of labeled returns could degrade their performance. One potential solution is to combine RyS algorithms with temporal- difference learning to correct errors in the labeled returns 44 46).
While also aiming to mitigate the problem caused by suboptimal labeled RTGs, our work takes a dif- ferent route - by leveraging TPMs mitigate the inference-time computational burden. Specifically, we identified major problems caused by the lack of tractability in the sequence models, and show that with the ability to compute more queries efficiently, we can partially solve both identified problems.
Limitations. One major limitation of Trifle is its dependency on expressive TPMs trained on sequential data a if the TPMs are inaccurate, then Trifle will also have inferior performance. Another limitation is that current implementations of PCs are not efficient as neural network packages, which could slow down the execution of Trifle.
Acknowledgements
This work was funded in part by the National Science and Technology Major Project (2022ZD01 14902). DARPA ANSR program under award FA8750-23-2-0066 the DARPA PTG Program under award HR001 12220005, and NSF grant 1943641.
References
[1] Anurag Ajay. Yilun Du, Abhi Gupta, Joshua B Tenenbaum, Tommi S Jaakkola, and Pulkit Agrawal. Is conditional generative modeling all you need for decision making? In The Eleventh International Conference on Learning Representations, 2022.
[2] Brandon Amos, Samuel Stanton, Denis Yarats, and Andrew Gordon Wilson, On the model-based stochastic value gradient for continuous reinforcement lcarning. In Learning for Dynamics and Control, pages 6 20. PMLR, 2021.
[3] David Brandfonbrener, Alberto Bictti, Jacob Buckman, Romain Laroche, and Joan Bruna. When does relurn-conditioned supervised learning work for offline reinforcement learning? Advances in Neural Information Processing Syxtems, 35:1 2022.
2022.
[4] Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym. arXiv preprint arXiv:1606.0155: 2016.
[5] Tom Brown, Benjamin Mann. Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are shot leamers. Advances in neural information processing systems, 33:1877- 2020.
10

---

[6] Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee. Aditya Grover. Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information prucessing xystems, 34:15084-15097, 2021.
[7] YooJung Choi. Antonio Vergari, and Guy Van den Broeck. Probabilistic circuits: A unifying framework for tractable probabilistic models. oct 2020.
[8] Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep reinforcement learning in handful of trials using probabilistic dynamics models, Advances in neural information processing systems, 31, 2018.
[9] Alvaro HC Correia, Gennaro Gala, Erik Quaeghebeur, Cassio Campos, and Robert Peharz. Continuous mixtures of tractable probabilistic models. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pages 7244-7252, 2023.
[10] Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre- Pre-training training of deep bidirec- tional transformers for language understanding, arXiv preprint 1rXiiv:810.049000 2018
[11] Thomas G Dietterich. Hierarchical reinforcement learning with the maxq value function decomposition. Journal of artificial intelligence research, 13:227-303, 2000.
[12] Wenhao Ding, Tong Che, Ding Zhao, and Marco Pavone. Bayesian reparameterization of reward- conditioned reinforcement learning with energy-based models. arXiv preprint arXiv:2305.11344 2023.
[13] Scott Emmons, Benjamin Eysenbach, Ilya Kostrikov, and Sergey Levine. Rvs: What is essential for offline rl via supervised learning? In International Conference on Learning Representations, 2021.
[14] Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning, 2020.
[15] Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without explo- ration. In International conference machine learning, pages 2052-2002. PMLR, 2019.
[16] Tuomas Haarnoja, Aurick Zhou, Kristian Hartikainen, George Tucker, Schoon Ha, Jie Tan, Vikash Kumar, Henry Zhu, Abhishek Gupta, Pieter Abbeel, et al. Soft actor-critic algorithms and applications. arXiv preprint arXiv: 8/2.05905 2018
[17] Nicolas Heess, Gregory Wayne, David Silver, Timothy Lillicrap, Tom Erez, and Yuval Tassa. Learning continuous control policies by stochastic value gradients. Advances in neural information processing systems, 28, 2015.
[18] Jonathan Ho, Ajay Jain, and Picter Ahbeel. Demoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-6851, 2020.
[19] Taher Jafferjee, Ehsan Imani, Erin Talvitie, Martha White, and Micheal Bowling. Hallucinating value: pitfall dyna-style planning with imperfect environment models. arXiv preprint arXiv:2006.0453 04363, 2020.
[20] Michacl Janner, Qiyang Li, and Sergey Levine. Offline reinforcement Icarning as one big sequence modeling problem. Advances in neural information processing systems, 34: 1273-1288,, 2021.
[21] Lukasz Kaiser, Mohammad Babacizadeh, Piotr Milos, Blazej Osinski. Roy Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, et al. Model-based rein- forcement learning for atari. arXiv preprint 2003v:/903.0033 2019.
[22] Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational difftusion models. Advances in neural information processing systems, 34:21696 22006622777711 2021.
[23] Doga Kisa, Guy Van den Broeck, Arthur Choi, and Adnan Darwiche. Probabilistic sentential decision diagrams. In Fourteenth International Conference on the Principles of Knowledge Representation and Reasoning, 2014.
[24] Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q-learning. arXiv preprint 20. 2021.
[25] Aviral Kumar, Xue Bin Peng, and Sergey Levine. Roward-conditioned policies. arXiv preprint
arXiv: 1912. 13465, 2019.
[26] Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. Advances in Neural Information Processing Systems, 33:1179- 1191, 2020.
11

---

[27] Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprini arXiv:2005.01644 2020.
[28] Anji Liu, Mathias Niepert, and Guy Van den Broeck. Image inpainting via tractable steering of diffusion models. In Proceedings of the Twelfth International Conference on Learning Representations (ICLR), 2024.
[29] Anji Liu and Guy Van den Broeck. Tractable regularization of probabilistic circuits. Advances Neural Information Procexsing Systems, 34:3558-3570, 2021.
[30] Anji Liu, Honghua Zhang, and Guy Van den Broeck. Scaling up probabilistic circuits latent variable distillation. In The Eleventh International Conference on Learning Representations. 2022.
[31] Xuejie Liu, Anji Liu, Guy Van den Broeck, and Yitao Liang. Understanding the distillation process from deep generative models to tractable prohabilistic circuits, International Conference on Machine Learning, pages 21838. PMLR, 2023.
[32] Keiran Paster, Sheila Mcllraith, and Jimmy Ba. You can't count on luck: rvs fail in stochastic environments, Advances in Neural Information Processing Systems, 35:38966-38979, 2022.
Why decision transformers and
[33] Robert Peharz, Robert Gens, Franz Pernkopf, and Pedro Domingos. On the latent variable interpretation in sum-product networks. IEEE trunsactions pattern analysis and machine intelligence. 39(10)22000-2044 2016.
[34] Dean A Pomerleau. Alvinn: An autonomous land vehicle in neural network. Advances in neural information processing systems, 1, 1988.
[35] Hoifung Poon and Pedro Domingos. Sum-product networks: a new deep architecture. In Proceedings of the Twenty-Seventh Conference on Uncertainty in Artificial Intelligence, pages 337-346, 2011.
[36] Lawrence Rabiner and Bringhwang Juang. An introduction to hidden markow models. ieee assp magazine, 3(1):4-16, 1986.
[37] Tahrima Rahman, Prasanna Kothalkar, and Vibhav Gogate. Cutset networks: A simple, tractable, and scalable approach for improving the accuracy of chow-liu trees. In Machine Learning and Knowledge Discovery in Databases: European Conference, ECML PKDD 2014, Nancy, France, September 15- 2014. Proceedings, Part 11 14, pages 630 6000-645.555 Springer, 2014.
[38] Mare Rigter, Bruno Lacerda, and Nick Hawes. Rambo-rl: Robust adversarial model-based offline reinforcement learning. Advances in neural information processing systems, 35:16022-16097, 2022.
[39] Juergen Schmidhuber. Reinforcement learning upside down; Don't predict rewards-just map them to actions. arXiv preprint arXiv:1912.02877 2019.
[40] Rupesh Kumar Srivastava, Pranav Shyam, Filipe Mutz, Wojciech Jaskowski, and Jürgen Schmidhuber. Training agents using upside-down reinforcement learning. arXiv preprint arXiv: 2019.
[41] Erin Talvitie. Self-correcting models for model-based reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 31, 2017.
[42] Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: physics engine for model-based control. In 2012 IEEF/RSJ international conference on intelligent robots and systems, pages 5026-5033. IEEE, 2012.
[43] Arash Vahdat and Jan Kautz. Nvac: A deep hicrarchical variational autoencoder. Advances in neural information processing systems, 33: 13:19967-11677 2020.
[44] Antonio Vergari, YooJung Choi, Anji Liu, Stefano Teso, and Guy Van den Broeck. A compositional atlas of tractable circuit operations for probabilistic inference. In Advances in Neural Information Processing Systems 34 (NeurIPS), dec 2021.
[45] Christopher JCH Watkins and Peter Dayan. Q-lcarning, Machine learning, 8:279 292, 1992.
[46] Taku Yamagata, Ahmed Khalil, and Raul Santos-Rodriguun Q-learning decision transformer: Leveraging dynamic programming for conditional sequence modelling in offline rl. In International Conference on Machine Learning, pages 99989-390077. PMLR, 2023.
[47] Mengjiao Yang. Dale Schuurmans. Pieter Abbeel, and Ofir Nachum. Dichotomy of control: Separating what you control from what you cannot. arXiv preprint 234.3,2200..335, 2022.
12

---

[48] Mingde Zhao, Zhen Liu, Sitao Luan, Shuyuan Zhang, Doina Precup. and Yoshua Bengio. A consciousness- inspired planning agent for model-based reinforcement learning. Advances neural information processing systems, 34:1569-1581, 2021.
[49] Qinqing Zheng, Amy Zhang, and Aditya Grover. Online decision transformer. In international conference on machine learning, pages -70599 PMLR, 2022.
13

---

Supplementary Material
Proof of Theorem
To improve the clarity the proof, we first simplify the notations in Thm. D define X the boolean action variables a, == {aije= and Y as the variable V,, which is categorical variable with two categories and We can equivalently interpret Y boolean variable where the category corresponds to F and 1 corresponds to T. Dropping the condition on S everywhere for notation simplicity, we have converted the problem into the following one:
Assume boolean variables == (Xi) and Y follow a Naive Bayes distribution: p(æ, y) = p(y) II, p(ally) We want to prove that computing p(ær Ely) v), which is defined as follows, is NP-hard.
p(x) if 1ay) p Ely] v) = 7
2 v.
otherwise.
By the definition of as a categorical variable with two categories 0 and 1, we have
By p(y T| pvy F(z) pvy = T)
Therefore, we can rewrite p(a: Ely] 2 v)) as
1
( El
u) =
p(x) 1ppyy T|x) v).
where 1[:) is the indicator function. In the following, we show that computing the normalizing constant 7. := Ex, p(z)- [p( x) 2 v) is NP-hard by reduction from the number partition problem, which is a known NP-hard problem. Specifically,, for set numbers 71..... nī (Vi, n4 E
2t), the number partition problem aims to decide whether there exists subset S C [k] (define [k] = (k....,k})
k}))
that partition the numbers into two sets with equal sums: Eies n2 Ejes n2.
For every number partition problem {n,} we define a corresponding Naive Bayes distribution p(x, y) with the following parameterization, p(y = T) = 0.5 and
e and p(:: = Tly F) = c 1 en n Vi E [M) p(r; = Tly T) = e m.
e nn
It is easy to verify that the above definitions lead to a valid Naive Bayes distribution. Further, we have
Vi [k], log pl = T = T)
p(æ, Fly - T) Fly
and log
(5)
p(z:
p(z;
We pair every partition in the number partition problem with an instance x such that Vi, mz = T 2/3, the normalizing constant Z can be written as
i S and z; = otherwise. Choose
Z Σ p(w) 1 [p(l T|æ) 2 2/3].
(6)
Recall the one-to-one correspondence between S and æ. we rewrite p(y = T|æ) with the Bayes formula:
ply T) II. PZZyy
p(y = Txx) plu = T) II, p(z,,p T) ply = p(
1 Σ. log playly=T] P(Rily=a
Note that assume the naive Bayes model is parameterized using log probabilities.
14

---

1 + -res ni- Dygss
where the last equation follows from Equation
After some simplifications, we have
1 [ply - T) 2/3
=
n4
Σ n
ES
igs
Plug back to Equation (6. we have
2 = Σ p(a:) 1[
Σ JØS
nj 2 21},
SC
SES
p(x)
Σ
n4
Σ igs
m1 + 0/
SCIR]
iES
where the last equation follows from the fact that (i) if x satisfy > Lies n4 -Eigs n2 2 then o has Ejes n2 must be an integer.
Eies - Ejgs n; I -1 and vise versa, and (ii) Eies ni Note that for every solution S to the number partition problem, Σies Σ jes n = 0 holds. Therefore, there exists solution the defined number partition problem if Z <
B Introduction to Probabilistic Circuits
Probabilistic circuits (PCs) represent a wide class of TPMs that model probability distributions with a parameterized directed acyclic computation graph (DAG). Specifically, a PC p(X) defines a joint distribution over a set of random variables X by single root node n2. A PC contains three kinds of computational nodes: imput C sum @ and product @ The example PC in Figure [ digfreedeeinns a joint distribution over 4 random variables X1, X2, X4. Each leaf node in the DAG serves as input node that encodes a univariate distribution (e-g., Guassian, Categorical), while sum nodes or product nodes are inner nodes, distinguished by whether they are doing mixture factorization over their child distributions (denoted in)n).. Formally, PCs define probability distributions in the following recursive way:
defines
fn(x))
if is an input unit, if n. is sum unit, if n. is product unit,
Pn (x) = Eeein(n) On,e' Pe(ar) Toein(n)) Pc(x)
where nnc represents the parameter corresponding edge (n, c) in the DAG. For nodes, have Eccin(n) On.c = 1, 1,0 2 0, and we assume w.l.o.g. that a PC alternates between the sum and product layers before reaching its inputs.
B.1 Tractability and Expressivity
The expressivity of PCs comes from the ability to combine simpler distributions with sum and product nodes to form more complex distributions. Therefore, increase the capacity of PC, we can add more nodes to its DAG or find better structure that is more tailored to the target data distribution. On the other hand, tractability, the ability to answer certain probabilistic queries efficiently and exactly, is guaranteed by certain structural properties of PCs. For example, with smoothness and decomposability defined in the following, PCs can compute arbitrary marginal and conditional probability in time linear with respect to its size (number of edges in its DAG).
Definition (Decomposabilityy) A PC is decomposable if for every product unit n, its children have disjoint scopes:
Vc,. E in(n) (a # c2), Φ(c2) n ((c2) = Ø.
Definition (Smoothness). PC smooth for every sum unit n, its children have the same scope:
Ver, in(n), d(e)) d(o2).
15

---

As key procedure used in Trifle, we describe how to compute marginal queries given a smooth and decomposable PC. First, we assign probabilities to all input nodes. For an input node defined on variable X, if evidence on X is provided in the query, we set the output probability following Equation (7). Otherwise, we set the output probability to Next, we do feedforward pass over all inner (sum and product) nodes following Equation The final output at the root node is the desired marginal probability.
In addition to marginal and conditional probabilities, PCs can efficiently and exactly compute other queries including maximize a-posterior and various information theorreti queries given some additional structural constraints. Please refer to [4] for comprehensive overview.
B.2 Adopted PC Structure And Parameter Learning Algorithm
For all tasks/offline datasets, adopt the Hidden Chow-Liu Tree (HCLT) PC structure proposed by 29 as it has been shown to perform well across different data types.
Following the definition in Equation 7 PC takes as input sample æ and outputs the corresponding probability Pr, (z)). Given dataset D. the PC optimizer takes the PC parameters (consisting of sum edge parameters and input node/distribution parameters) as input and aims to maximize the MLE objective ETED log Pm
(z)). Since PCs can be deemed as latent variable models with hierarchically nested latent space [33]. the -oxiiaaioo-Maiiiizaation (EM) algorithm is usually the default choice for PC parameter learning. We adopt the ball-bacch EM algorithm proposed in [35..
Before tuning the parameters with EM, we adopt the latent variable distillation (LVD) technique proposed in (30 to initialize the PC parameters. Specifically, the neural embeddings used for LVD are acquired by BERT-like Transformer [10] trained with the Masked Language Model task. To acquire the embeddings of subset of variables Φ, we feed the Transformer with all other variables and concatenate the last Transformer layer's output for the variables Φ. Please refer to the original paper for more details.
We use the same quantile dataset discretized from the original Gym-MuJoCo dataset as done by TT [20), where each raw continuous variable is divided into 100 categoricals, and each categorical represents an equal amount of probability mass under the empirical data distribution.
C Algorithm Details of Trifle
This section provides a detailed description of the algorithmic procedure of step value estimates.
Trifle with single-/multi-
C.1 Trifle with Single-/Multi-Step Value Estimates
Similar to other RvS algorithms, Trifle first trains sequence models given truncated trajectories {(st, 0z, r2, RTG,)). Specifically, we fit two sequence models: autoregressive Transformer following prior work [20] as well as a PC, where the training details are introduced in Appx. B.2
During the evaluation phase, at time step I, Trifle is tasked to generate a; given 8cx and other relevant information (such as rewards collected past steps). As introduced in Section [5] Trifle generally works in two phases: rejection sampling for action generation and beam search for action selection. The main algorithm is illustrated in Algorithm[] [ where we take the current state š, as well as the past trajectory T<t input, utilize the specified value estimate fo heuristic guide beam search, and output the best trajectory. Note that f. is a subroutine of our algorithm that uses the trained sequence models to compute certain quantities, which will be detailed in subsequent parts. After that, we extract the current action ag from the output trajectory to execute in the environment.
At the first step of the beam search, we perform rejection sampling to obtain a candidate action set at (line of Algorithm|]] []] The concrete rejection sampling procedure for s-Trifle detailed in Algorithm 2 The major modification of m-Trifle compared to s-Trifle is the adoption of multi- step value estimate instead of the single-step value estimate, which is also shown in Algorithm 3 Specifically, Algorithm
3 is used to replace the value function fo shown in Algorithm |
16

---

Algorithm Trifle with Beam Search
Input: past trajectory Tet, current state 8s, beam width beam horizon H. scaling ratio sequence model M, value function f.
> fo
E[Vi] for s-Trifle and B[V?")) for m Trifle
2: Output: The best action 0r
Let Xt concat(T<c, 8t) (t)T:c,,,))..shaaeee (ttiappo(,-1)rreppaa (N. dim 0)
Batchify the input trajectory D cf. Algorithm
4: Perform rejection sampling to obtain a, using Algorithm 2
5: Initialize Xu con
6: foreach t = 1, l,.,,
XI-1 Xe-eett (A, dim 0)
Scale the number of trajectories from N XN
< (concat(xx x) Vx-11 X1-1, sample x ~ PM( xx-1)} b. Candidate next-token
prediction
X, f xep, (fo(X), NN)
keep most rewarding trajectories
10: end foreach
11: Xm argmaxxeXи fe(X)
12: return a; Xm
Algorithm 2 Rejection Sampling with Single-step Value Estimate
1: Input: past trajectory T<r. current state 8s, dimension of action k; rejection rate > 0
2: Output: The sampled action of
3: Let di <- concat(T<e,.5)
4: iori=1,... k do
5: Compute pGrT (a) Zt. a2
Note that of = Ø.
Compute pTPM (V: 12, a2, Enne PTPMa(V,, th Ir, aj 2:h)
The marginal can efficiently
computed by PC in linear time. See the algorithm in Appx.
7: Compute U/ ma E val(Vi) pTPM (Vi 2 0 Iı, a
1 6), for each oj val(Ai)
Compute p(a me. "0 us) = ParrTaa Te, az)) PPrMM (V 2 us Te, 0z () Apply Equation
Sample of p(a 22, 0, v08)
10: end for return o,
Algorithm 3 Multi-step Value Estimate
1: Input: Tg: (Su, ao, St, at), sequence model M, terminal timestep t' >> t. discount
2: Output: The multi-step value estimate E [[];"]]
3: Sample future actions at+1,..., av from M
4: Compute prp (g(,a(41.)) Enara PIPM(TH, Tgg) for E [t 1, Marginalize over intermediate states 8641th
Compute PTPM(RRCC TTt, ((41:) E.,, Sag (RTG, rge)
6: Compute
B[v="] Σn A- " Pra.pa(a((l [na] t " (aa)) [RTG]]
hmi
return E [
C.2
Computing Multi-Step Value Estimates
In this section, we present an efficient algorithm that computes Equation (3). From the decomposition of Equation (3), we can calculate E [V:"] if we have the probabilities p(r- o::e) (t < t') and pRTG, St, de:t). simple approach would to compute each the - ++1 probabilities separately using the algorithm described in Appx. B. (recall that conditional probabilities are quotient of the corresponding marginal probabilities). However, this approach has an undesired time complexity that scales lincarly with respect to t 1+
17

---

Following [28, we describe an algorithm that can compute all desired quantities using single feedforward and a backward pass the PC.
The forward pass. The forward pass is similar to the one descrihed in Appx. B.1 Specifically, we set the evidence as St, Gt:t' and execute the forward pass.
The backward pass. The backward pass consists of two steps: (1) traverse all nodes parents before children to compute statistic termed flow for every node n: flow,; (ii) compute the target probabilities using the flow all input nodes. Recall that we assume without loss of generality that and product layers. We further assume that all parents of input nodes
PCs alternate between are product nodes. We define the flow of the root node as 1. The flow of other nodes defined recursively as (define Pn. as the forward probability of node n))
E Emeppt() (0m, Po/Pm) flowm nt a product node, - (mmppa()) flowm
flow, ==
input sum node,
where pa(n) the set of parents node n.
Next for every variable X E .R...... Re-1,, RTGG) we first collect all input nodes defined on X. Define the set of input nodes as S. We have that
p( asa) = Σ
flow, fx(x,,
Z
nES
where fn is defined in Equation 7 and Z is a normalizing constant.
D Inference-time Optimality Score
We define the inference-time optimality score as proxy for inference-time optimality. This score primarily defined over state-action pair (st, 0z) at each inference step. In Figure 2 (right), each sample point represents a trajectory, and the corresponding inference-time optimality score is defined over the entire trajectory by averaging the scores of all inference steps.
and
The specific steps for calculating the score for given inference step given 8t and policy p(a, 8t). are as follows:
1. Given st, sample a, from prr(a: st). PDr(a, st), or Prrine(a sc)-
2. Compute the state-conditioned value distribution p" (V, st).
3. Compute R, (RTG [V)], which is the corresponding estimated expected value.
4. Output the quantile value S, of R2 (V, 8t).
To approximate the distributions (V, 8t) and (V, 8t, 82,0,) az) (where Vú RTG () in steps and 3, we train two auxiliary GPT models using the offline dataset. For instance, to approximate p² pV(V. st), we train the model on sequences (s * k, Vi-k..... St. Vi).
Intuitively, (V, se) approximates p(V/ st) == E., p(Ve St, at) p(ar sy). Therefore, S indicates the percentile of the sampled action in terms of achieving high expected return, relative the entire action space.
E
Additional Experimental Details
E.1
Gym-MuJoCo
Sampling Details. We take the single-step value estimate by setting V = RTG, and sample o, from Equation 2h. When training the GPT used for querying PGPT(aj)ss a a, we adopt the same model specification and training pipeline as TT or DT. When computing pTPM( 2 v|st, af"), we first use the learned PC to estimate v(V, sx) by marginalizing out intermediate actions oes and select the c-quantile value of (VV])s)) as our prediction threshold for each inference step. Empirically we fixed for each environment and ranges from 0.1 to 0.3.
18

---

Beam Search Hyperparameters. The maximum beam width N and planning horizon H that Trifle uses across MuJoCo tasks are 15 and 64, respectively.
Comparison with Value- Based Algorithms. To shed light on how Trifle compares to methods that directly optimize the Q values while filtering actions by conditioning high returns (as done in RvS algorithms), we compare Trifle with Q-learning Decision Transformer [46], which incorporates contrastive Q-learning regime into the RvS framework. As shown in the table below, Trifle outperforms QDT in all six adopted MuJoCo benchmarks:
Table Normalized Scores of QDT and Trifle on Gym-MuJoCo benchmarks
Dataset
Environment
Trifle
QDT
Medium
Halfcheetah
49.5s0.2
42.320.4
Med-Replay
Halfcheetah
45.0a0.3
35.620.5
Medium
Hopper
67.1+4.3
66.526.3
Med- Rep-Ry
Hopper
97.8a0.3
52.1420.3
Medium
Walker2d
83.1
67.1232
Med-Replay
Walker2d
88.3s3.
58.2a5.1
E.2 Stochastic Taxi Environment
Hyperparameters. Except for s-Trifle, the sequence length K modeled by TT, DT, and m-Trifle all equal to 7. The inference algorithm of TT follows that of the MuloCo experiment and DT follows its implementation the Atati benchmark. Notably, during evaluation, we condition the pretrained DT on 6 different RTGs ranging from -100 to -350 and choose the best policy resulting from RTG=-300 to report Figure 3c Beam width = 8 and planning horizon = 3 hold for TT and m-Trifle.
Additional Results on the Taxi benchmark. Besides the episode return, we adopt two metrics to better evaluate the adopted methods: (i) #penalty: the average number of executing illegal actions within an episode; (ii) P(failure): the probability of failing to transport the passenger to the destination within 300 steps.
Table 4: Results on the stochastic Taxi environment. All the reported numbers are averaged over 1000 trials.
Methods
Episode return
# penalty
P(failure)
s-Trifle
-99
0.14
0.11
m-Trifle
-57
0.38
0.02
TT
-182
2.57
0.34
DT
-388
14.2
0.66
DoC
-146
0.28
dataset
-128
2.41
Ablation Study Regarding Action Filtering. In an attempt to justify the effeccriveness/ncessitty of exact inference, we compare Trifle with value-based action filtering/value estimation in the following:
To begin with, we implemented the traditional Policy Evaluation algorithm on the Taxi offline dataset described in Section 6.2 of the paper. The policy evaluation algorithm is based on the Bellman update:
Q(sr. 0z) + Q(si, 0z) + a afre41 + 7 ac+1) ((s az))
Then we use the obtained Q function, denoted Quai, to perform the following ablation studies. We still choose TT as our base RvS model. Recall that given s,, TT first samples a₂ from its learned prior policy (ar((a)ss), which are subsequently fed beam search procedure that uses learned value function PTT( PTT(Ve]s, az)) select the best action. Therefore, we consider ablations on two key components of TT: (i) the prior policy pTT PTr(Ve[st, az) used sample actions, and (ii) the value function PTT (V:) st, a.) used to evaluate and select actions.
19

---

1. TT + Quaxt action filtering: weigh the prior policy prr(ag)s; with each action's 's exponen- tiated Quai value (i.e., vux ($,, at))), but still adopt TT's value estimation. In other words, in this experiment, we only use Qiai to improve the sampling quality as s-Trifle does.
2. TT + Quaxi value estimation: replace PrT and selection, but still use TT's prior policy PrT (4z)8:).
(Vi (rrr(Viss
at) with Quxi( oz) for action evaluation
3. TT + full Quad: simultancously use (uxi (3e, at)) for action filtering and Quaxi (St, as) for action evaluation.
We present the results of these ablation studies follows:
Method
Score
TT
-182
TT + Quai action filtering
-157
TT Quai value estimation
147
TT full Quxi
-13
m-Trille s-Trifle
-58
-99
From these results, we draw the following conclusions:
m-Trifle and s- Trifee achieve the best performance.
The rank of scores: TT full Quaii TT Quxii value estimation TT Quaii action filtering > TT suggests that using Quxi for both action filtering and value estimation is beneficial; combining the two leads the best performance.
Specifically, the fact that s-Trifle outperforms the Quxi based action filtering demonstrates that our filtration with exact inference is much more effective. The superior performance of m-Trifle also provides strong evidence that explicit marginalization over future states leads to better value estimation.
Additional Experiments
F1
Ablation Studies on Rejection Sampling and Beam Search
The key insight of Trifle to solve challenges elaborated in Scenario #1 is to utilize tractable probabilis- tie models to better approximate action samples from the desired distribution p(a:) E[V] 2 v). We highlight that the most crucial design choice of our method for this goal is that: Trifle can effectively bias the per-action-dimension generation process of any base policy towards high expected returns, which is achieved by adding per-dimension correction terms PTPM (V₂ 2 u|st, ar (Eq. (2)
in the paper) to the base policy.
While the rejection sampling method can help us obtain more unbiased action samples through post value(expected (eturn--estimattin session, we only implement this component for TT-based Trifle (not for DT- bT-based Trifle) for fair comparison, as the DT baseline doesn't 't perform explicit value estimation or adopt any rejection sampling methods. Therefore, the success of DTT-ased Trifle strongly justifies the effectiveness of the TPM components. Moreover, the beam search algorithm also comes from TT. Although a more effective way to do rejection sampling, not the necessary component of Trifle, either.
20

---

Figure 5: Scaling Curves of Inference Time. (Fix beam width 32)
Inference Time Horizon
Runtime #Act Vars
2.3
Step-wise Inference Timels)
1.2
2.0
Runtimels)
1.0
1.5
0.8
0.2
Horizon
Vars
Table Ablations over Beam Search Hyperparameters on Halfcheetah Med-Replay, (a) With H 1, the beam search degrades to naive rejection sampling (b) With W 1, the algorithm doesn't 't perform rejection sampling. It samples a single action and applies it to the environment directly.
(a) Varying Planning Horizon
(b) Varying Beam Width
Horizon H
Width W
TT
Ti-based Trifle
Horizon H
Width W
TT
T--bused Tridle
41.91.2.5
45.0-10.3
41.9-1.55
45.0- 0.3
32
40.11:2.0
43.1.11.0
5
16
42.511..9
42.6-:66
41.62:13
42.61:.6
42.9=0.4
43.5=0.3
32
39.7122.5
42.8:0.5
5
38.7.10.3
43.4-0.3
naive rej sampling)
32
33.64:3.0
39.6+00.7
(wlo sampling)
31.2+3.4
367--188
For TT-based Trifle, we adopted the same beam search hyperparameters as reported in the paper. We conduct ablation studies beam search hyperparameters in Table 5 investigate the effectiveness of Trifle's 's each component. From Table 5 we can observe that:
Trifle consistently outperforms TT across all beam search hyperparameters and is more robust to variations of both planning horizon and beam width.
(a) Trifle w/ naive rejection sampling TT w/ naive rejection sampling (b) Trifle w/o rejection sampling TT w/o rejection sampling. In both cases, Trifle positively guide action generation.
Trifle w/ beam search Trifle w/ naive rejection sampling Trifle w/o rejection sampling TT w/ naive rejection sampling. Although other design choices like rejection sampling/beam search help better approximate samples from the fnt-eexpeceed-euu--ooondiii distribution, the per-dimension correction terms computed by Trifle play a very significant role.
action
F.2 Computational Efficiency Analysis
Since TPM-related computation consistently requires 1.45s computation time across different horizons, the relative slowdown of Trifle is diminishing as we increase the beam horizon. Specifically, the Gym-Mujuco benchmark, the time consumption
one step (i.e., one interaction with environment) of TT and Trifle with different beam horizons are listed here (Figure (left) also plots an inference-time scaling curve of Trifle vs TT with varying horizons):
Table 6: The one-step inference runtime of the Gym-MuJuch benchmark
Horizon
TT
Trifle
0.5s
1.5s
15
1.2s
1.8s
21

---

Moreover, Figure (right) shows that Trifle's s runtime (PPM-related) scales linearly w.r.t. the number of action variables, which indicates its efficiency for handling high- dimeeeiinnii action spaces.
Trifle also efficient in training. It only takes 30-60 minutes 20s per epoch, 100-200 epochs) to train PC
GPU for each Gym-Mujuco task (Note that single PC model can be used to answer all conditional queries required by Trifte). In comparison, training the GPT model for TT takes approximately 6-12 hours (80 epochs).
F.3 Ablation Studies on the Adaptive Thresholding Mechanism
The adaptive thresholding mechanism is adopted when computing the term PTPM (V; 2 t/8,, a²") of Equation (2), where E (1....,k: k). & is the number action variables and aj is the ith variable of a, Instead of using fixed threshold
choose be the c-quantile value of the distribution (V: St, a²²) computed by the TPM, which leverage the TPM's ability to exactly compute marginals given *#incomplete* actions (marginalizing out a. Specifically, we compute v using max,(r {r E prPM (Vi 2 r|st, aa*) 2 1 - c). Empirically we fixed c for each Gym-MuJoCo environment and 0.2 or 0.25, which is selected by running grid search E [0.1, 0.25|.
Table 7: Comparison of Adaptive and Fixed Thresholding Mechanisms
(b) Performance of Fixed Thresholding (Varying
(a) Ablations over Adaptive Thresholding (Varying
<) on Halfcheetah Med- Med-Replay
Halfcheetah Med-Replay
Walker2d Med-Expert
Method
Score
adaptive
45.0=0.3
109.3:10.1
90
44.8
109. 1::0.2
TT
41.9:12.55
39,54228
108. 108.9100.2
TT-based Trifle (e 0.25)
45.0:10.3
70
44.9-0.3
108.4+:0.4
TT-based Trifle (e 0.2)
44,24-0.4
60
42.62:16
105.8:10.3
TT-based Trifle (c 0.15)
44.4:10.3
107.5++1.5
TT-based Trifle (e 0.1)
42.6++11.6
40
42.621.6
107.5+1.44
30
44.020.4
98.3:55.4
We report the performance of TT-based Trifle with variant e vs TT on Halfcheetah Med-Replay benchmark in Table 7a We can see that Trifle is robust to the hyperparameter and consistently outperforms the base policy TT.
We also conduct ablation studies comparing the performance of the adaptive thresholding mecha- nism with the fixed thresholding mechanism on two environments in Table 7b Specifically, given that Ve is discretized to be a categorical variable with 100 categoricals (0-99), we fix to be 600,70,00,50,40000
respectively.
The table shows that the adaptive approach consistently outperforms the fixed value threshold in both environments. Additionally, the performance variation of fixing v is larger compared to fixing as different v can be optimal for different states.
G Potential Negative Societal Impact
This paper proposes new offline RL algorithm. which aims to produce policies that achieve high expected returns given pre-collected dataset of trajectories generated by some unknown policies. When there are malicious trajectories in the dataset, our method could potentially learn to mimic such behavior. Therefore, we should only train the proposed agent in verified and trusted offline datasets.
22

---

NeurIPS Paper Checklist
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
Justification: All claims regarding empirical performance are justified by the results in
Section 6
The claims regarding the importance of tractability are justified by empirical
evidence shown in Section 3
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We have discussed limitations of our work in Section
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: For Thm. I we have elaborated the full assumptions in its main body. A proof
is provided in Appx. A
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: Our code is available at have provided full algorithm details (including algorithm tables) Section Appx. C Adopted hyperparameters are described in Section and Appx. well as all trained models will be released this paper gets accepted.
Moreover, we
and and
E The code as
Guidelines:
The answer NA means that the paper does not include experiments.
23

---

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. w
For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification:
Our code is available httpss//gihhbb.com/llbee////////
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (nttps://nips.co
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
/ 7/nips. /ttppp.cc/pblic/guideeee http:/ooeeosmiiiiooooooooooo
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
24

---

Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
Detailed settings of the experiments can be found in Section and Appx.
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented in the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We reported mean and standard deviation over 12 or more runs most applicable experiments, c.g., Table 1 and and Figure Gc
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2-simma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time exccution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We have a discussion regarding runtime in Appx.
F2
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
25

---

The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics httpa: httpp:::///euipp... http:pubbic/Bhiicscuieeeee
Answer: [Yes]
Justification: We have read the NeurIPS Code of Ethics and the research conducted in the paper conforms
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: The paper studies specific type of offline RL algorithm. The proposed method itself only tested on simulated environments such as games and thus has no immediate societal impact.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of MI.).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (c.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
26

---

Justification:
The paper does not release new data and does not involve large language models.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (c.g., code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification:
We have cited the owners/authors of the benchmarks we used.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi www.wethhde...m caree.comdaie com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re- re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: Not applicable.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/mooel as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
27

---

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Question:
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
28


Reviews:
[{"review_id": "kUlHh57HzM", "reviewer": "Reviewer_jZtp", "summary": "The paper introduces Trifle (Tractable Inference for Offline RL) that leverages Tractable Probabilistic Models (TPMs) to enhance the performance of offline RL tasks. The paper emphasizes that beyond the expressiveness of sequence models, tractability--efficiently answering probabilistic queries--is crucial for accurate evaluation in offline RL, particularly in environments with high stochasticity and action constraints. In particular, Trifle uses a type of TPMs called Probabilistic Circuits (PCs), which supports the computation of arbitrary marginal and conditional probabilities in linear time with respect to their size. In practice, Trifle uses a mixture of single-step (return-to-go) and multi-step value estimates to condition action generation, allowing Trifle to handle both optimal and suboptimal return-to-gos effectively. Trifle enhances traditional rejection sampling by incorporating TPMs to adjust the proposal distribution, improving the likelihood of sampling high-return actions. It employs beam search to maintain and refine high-probability action sequences, ensuring actions are within the offline data distribution and likely to yield high rewards. Additionally, Trifle uses an adaptive thresholding mechanism to dynamically select expected return thresholds, maintaining robust performance across various evaluation settings.\n\nThe paper provides comprehensive empirical comparisons across nine Gym-MuJoCo benchmarks, a stochastic Taxi environment, and action-space-constrained tasks. Trifle consistently outperforms state-of-the-art baselines, including Trajectory Transformers (TT), Decision Transformers (DT), and various imitation learning and offline TD learning methods. These results demonstrate Trifle's robustness and effectiveness, particularly in challenging stochastic and constrained environments, underscoring its potential to advance offline RL methodologies.", "questions": "Related to the above’s weaknesses:\n1. How does the adaptive thresholding mechanism and beam search hyperparameters affect the performance of Trifle across different environments and datasets? Having some ablation studies on this would be helpful to suggest practical values for future work. \n2. Could you provide more intuitive explanations and discuss the practical implications of the theoretical guarantees provided in Sections 4.1 and 4.2? While the theoretical results are well-presented, further clarity on how these results guide the design and implementation of Trifle would help readers better understand the theoretical contributions and their impact on performance. Specifically, discussing practical scenarios where these theoretical guarantees are most beneficial would strengthen the paper.\n3. Have you considered evaluating Trifle on more complex tasks like the Antmaze benchmark?\n\nPost rebuttal: increased my score from 6 to 7.", "limitations": "The authors have addressed some limitations of their work, specifically the dependency on expressive TPMs and the current inefficiency of PCs compared to neural network packages.  However, there is no discussion on the potential negative societal impact of their work as the method was tested on simulated environments with no immediate societal impact.\n\nOne possible limitation to discuss is also how Trifle scales in high-dimensional action spaces, which might introduce computational overhead and impact performance.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "**Originality**: The paper presents a novel perspective by highlighting the importance of tractability in offline RL. By introducing TPMs into this domain, it addresses a gap not thoroughly explored by existing approaches, which typically focus on the expressiveness of sequence models.\n\n**Quality**: The methodology is robust and well-supported by theoretical insights and extensive empirical evaluations. The experiments demonstrate that Trifle significantly outperforms strong baselines on diverse benchmarks, including challenging stochastic environments and safe RL tasks with action constraints.\n\n**Clarity**: The paper is well-written and organized, making complex concepts accessible. The inclusion of detailed experimental setups and comprehensive results aids in understanding the advantages of Trifle over existing methods. I appreciate the overview on TPMs and the build up from section 3 to 5 from the theoretical to the practical implementation.\n\n**Significance**: By demonstrating that tractability can greatly enhance evaluation-time performance, the paper paves the way for future research to develop more inference-aware RL approaches. The results on diverse and challenging benchmarks highlight the practical benefits of integrating TPMs into RL algorithms.", "weaknesses": "1. Limited evaluations on complex tasks (e.g. in the D4RL benchmark beyond the nine studied). There is perhaps a challenge (computational?) to scale up the PCs on more complicated environments. Even if the authors do not provide the results for these, perhaps it would be worth discussing any challenges on scaling up Trifle to these more complex tasks. \n2. More detailed ablation studies. The paper includes some ablation studies, such as comparisons of Trifle variants with and without Q-value-based action filtering, demonstrating the effectiveness of exact inference. However, it would benefit from additional detailed ablations to isolate the contributions of other components, such as the adaptive thresholding mechanism and beam search strategy, to offer deeper insights into which elements are most critical to Trifle's performance improvements.\n3. Scalability / Efficiency Analysis. The paper lacks a thorough analysis of Trifle's computational complexity and scalability. Discussing the trade-offs between the benefits of tractability and the computational overhead introduced, especially in high-dimensional action spaces, would provide a clearer picture of Trifle's practical feasibility. Including detailed runtime comparisons with other methods would be beneficial. The current runtime analysis indicates a 1.5 to 3 times increase in inference time compared to base models, but a more comprehensive analysis would be useful (e.g. plotting out a scaling law curve for Trifle vs TT in table 5 with more horizon, etc.)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "### Comment #1: practical implications of the theoretical guarantees in Secs. 4.1 and 4.2\n\nThanks for the constructive comment. The theoretical results are used to elaborate two key inference-side challenges. Both challenges are introduced in the general comment, and we provide further details in the following. We will incorporate the following discussion into the next version of the paper.\n\n**Sampling actions conditioned on high expected return**\n\nAs described in Section 3, our goal is to sample from $p(a_t | s_{0:t}, \\mathbb{E}[V_{t}] \\geq v)$. However, Theorem 1 illustrates that it is NP-hard to exact compute even when $p(s_{0:t}, a_{t}, V_{t})$ follows a simple Naive Bayes distribution, which guides Trifle to find *good approximations*.\n\nSpecifically, we utilize TPM to compute the correction term $p_{TPM} (V_t \\geq v | s_t, a_{t}^{\\leq i})$ to bias the actions generated by any prior policy towards high expected return $\\mathbb{E}[V_t]$. Note that we compute the correction term exactly for each dimension $a_t^i$, which is highly non-trivial as we have to marginalize out unseen action variables $a_t^{i+1},...,a_t^k$ and this forms an exponentially large action space. We have **more ablation studies** on the effect of the TPM-provided terms in the response to comment #2.\n\n**Estimating state-action values under stochastic transition dynamics**\n\nAs discussed in the general comment, we need to use multi-step value estimates under stochastic environments for better accuracy. While it is generally hard to obtain low-variance estimates, with TPMs' ability to compute arbitrary conditional distributions, we can efficiently compute the multi-step value estimate in Eq. (3). This allows us to enjoy the better accuracy of the multi-step value estimates while mitigating the inefficiency to compute them.\n\n### Comment #2: effect of the beam search hyperparameters\n\nWe conduct additional ablation studies on both the adaptive thresholding mechanism and rejection sampling/beam search.\n\n**Ablation studies on the adaptive thresholding mechanism**\n\nWe report the performance of TT-based Trifle with variant $\\epsilon$ vs TT on Halfcheetah Med-Replay in **Table 1(a)** of the rebuttal PDF. Trifle is robust to $\\epsilon$ and consistently outperforms TT.\n\nWe also conduct ablation studies comparing the performance of the adaptive thresholding mechanism with the **fixed thresholding mechanism** on two environments in **Table 1(b)** by fixing different $v$s. The table shows that the adaptive approach consistently outperforms the fixed value threshold in both environments.\n\n**Ablation studies on the rejection sampling/beam search** \n\nAs mentioned in the general response, for the Gym-MuJoCo benchmark, we only adopt rejection sampling/beam search for TT-based Trifle and not for DT-based Trifle. Therefore, the success of DT-based Trifle strongly justifies the effectiveness of the TPM components.\n\nFor TT-based Trifle, we adopted the same beam search hyperparameters as reported in the TT paper. We conduct ablation studies on beam search hyperparameters in **Table 2** of PDF to investigate the effectiveness of Trifle's each component:\n\n-  Trifle consistently outperforms TT across all beam search hyperparameters and is more robust to variations of both planning horizon $H$ and beam width $W$.\n- (a) Trifle w/ naive rejection sampling >> TT w/ naive rejection sampling (b) Trifle w/o rejection sampling >> TT w/o rejection sampling. In both cases, Trifle can positively guide action generation.\n- Trifle w/ beam search > Trifle w/ naive rejection sampling > Trifle w/o rejection sampling >> TT w/ naive rejection sampling. Although other design choices like rejection sampling/beam search help to better approximate samples from the high-expected-return-conditioned action distribution, the per-dimension correction terms computed by Trifle play a very significant role. \n\n### Comment #3: computational complexity analysis\n\nThanks for the valuable suggestion.\n\nFirst, we conduct a more detailed runtime analysis, and the main results are shown in Figure 1 of the attached PDF. \n\nFigure 1 (left) in the rebuttal PDF expands Table 5 and plots the step-wise inference-time scaling curve of Trifle vs TT with varying horizons. As we increase the horizon, the relative slowdown is mitigated. This is because, across different horizons, TPM-related computation consistently requires ~1.45s computation time, which additional computation overhead is diminishing as we increase the beam horizon. \n\nMoreover, Trifle is efficient in training. It only takes 30-60 minutes to train a PC (the adopted TPM) on one GPU for each Gym-Mujuco task (note that we only need one PC per task). In comparison, training the GPT model for TT takes approximately 6-12 hours (80 epochs). \n\nNote that there are recent breakthroughs [1] on designing efficient PC implementations, which can significantly speed up training and inference of Trifle. \n\n[1] Liu et. al. \"Scaling Tractable Probabilistic Circuits: A Systems Perspective.\" arXiv preprint arXiv:2406.00766 (2024).\n\n### Comment #4: evaluation on the Antmaze task\n\nWe tried to run the Antmaze task but have not successfully reproduce the performance of TT+Q (as reported by the original paper, the vanilla TT fails due to the sparsity of the rewards) since there is no official implementation on GitHub and we have not found third-party implementation for this. We will keep trying to reproduce during the discussion period.\n\n### Comment #5: potential negative societal impact\n\nThanks. We will include the following discussion:\n\nThis paper proposes a new offline RL algorithm, which aims to produce policies that achieve high expected returns given a pre-collected dataset of trajectories generated by some unknown policies. When there are malicious trajectories in the dataset, our method could potentially learn to mimic such behavior. Therefore, we should only train the proposed agent in verified and trusted offline datasets.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "iId8sCYMmh", "reviewer": "Reviewer_DwFR", "summary": "This paper studies offline reinforcement learning and argues that tractability, e.g., the ability to answer probabilistic queries, is important for performance improvement. As a result, the authors propose a model that utilizes modern tractable generative models to answer arbitrary marginal/conditional probabilities. Comprehensive experiments demonstrate that the new model can achieve better performance in most cases.", "questions": "See above", "limitations": "Yes", "rating": 4, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The paper is overall well-written and easy to follow. The considered problem is interesting and relevant. A comprehensive set of experiments also verify the merits of the proposed approach.", "weaknesses": "While the proposed approach seems sound, the novelty and significance might be limited. The key idea in Sec 4 seems to be based on classical rejection sampling methods, though with some additions of correction terms. \n\nCould the authors provide some discussions regarding computation complexity of the proposed methods?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "### Comment #1: connection between rejection sampling and the correction terms, and the corresponding computational complexity\n\nWe thank the reviewer for their constructive feedback. As mentioned in the general response, Trifle can be applied to many existing RvS algorithms (e.g., TT, DT) to mitigate the identified two challenges. To highlight, the per-action-dimension correction terms in Eq. (2) computed by TPM are crucial to enhance action sampling and mitigate the first challenge. And the rejection sampling methods via beam search steps, are only adopted by the specific TT-based Trifle algorithm described in Section 4. DT doesn't apply beam search as TT so we don't implement rejection sampling for DT-based Trifle. Therefore, the success of DT-based Trifle strongly justifies that the correction term can effectively bias the actions generated by DT towards higher expected returns.\n\n\nIn the following, we will first discuss **how Trifle addresses both challenges** and why TPM matters, as well as the **computation complexity of the proposed methods**. Then we conduct ablation studies on rejection sampling to confirm Trifle's superiority.\n\n**Challenge 1: Sampling actions conditioned on high expected return**\n\nAs described in Section 3, this problem can be formulated as: given a joint distribution over $p(s_{0:t}, a_{t}, V_{t})$, we want to query $p(a_t | s_{0:t}, \\mathbb{E}[V_{t}] \\geq v)$ (formally defined in Equation (1)). Note that the expectation over $V_t$ is very important because otherwise if we ignore the expectation operator and directly sample from $p(a_t | s_{0:t}, V_{t} \\geq v)$, we may sample actions with a low probability that lead to high returns. This causes significant performance degradation as discussed in [1]. Then we naturally want to know whether we can exactly compute this quantity. However, Theorem 1 illustrates that it is NP-hard to exactly compute $p(a_t | s_{0:t}, \\mathbb{E}[V_{t}] \\geq v)$ even when $p(s_{0:t}, a_{t}, V_{t})$ follows a simple Naive Bayes distribution, which guides Trifle to find *good approximations* of the query.\n\nTo achieve better approximations, we utilize TPM to compute per-action-dimension correction terms  $p_{TPM} (V_t \\geq v | s_t, a_{t}^{\\leq i})$ to bias the actions generated by any prior policy towards high expected return $\\mathbb{E}[V_t]$. Note that we compute the correction term exactly for each dimension $a_t^i$, which is highly non-trivial as we have to marginalize out unseen action variables $a_t^{i+1},...,a_t^k$ and this forms an exponentially large action space. This is why intractable models suffer and we need to leverage the ability of TPMs to compute arbitrary conditional probabilities. Notably, the computation for TPM is exact without any approximation and scales **linearly** to the number of action variables, which is highly efficient. We will later show ablation results that confirm **the superiority of the correction term both with and without the rejection sampling step**.\n\n**Challenge 2: Estimating state-action values under stochastic transition dynamics**\n\nThe second major challenge is the inaccuracy of the RTG labels when the environment transition dynamics are highly stochastic. This is because RTGs heavily depend on the randomness that determines state transitions. To combat this inaccuracy, it is common to use multi-step value estimates instead. TD(1) and TD(λ) are notable examples. TD(1) relies on full returns to update value estimates, while TD(λ) introduces a balance between n-step returns and full returns using a parameter λ, allowing for more accurate value estimation. An application of this is the Proximal Policy Optimization (PPO) algorithm, which uses the eligibility trace governed by TD(λ) to establish better value estimates. This approach helps PPO achieve more stable and accurate value estimation by effectively integrating future reward information over multiple timesteps.\n\nWhile theoretically promising, classical multi-step value estimates generally require careful Monte Carlo sampling of different future state sequences, which makes it hard to obtain low-variance estimates. However, as introduced in Appendix C.2, with TPMs' ability to compute arbitrary conditional distributions, we can efficiently (in one forward pass and one backward pass of the TPM ) compute the multi-step value estimate in Equation (3). This allows us to enjoy the better accuracy of the multi-step value estimates while mitigating the inefficiency of computing them.\n\n**Ablation studies on the rejection sampling/beam search**\n\nFor TT-based Trifle, we adopted the same beam search hyperparameters as reported in the TT paper (in the official GitHub repo https://github.com/jannerm/trajectory-transformer). We conduct ablation studies on beam search hyperparameters in **Table 2** of PDF to investigate the effectiveness of Trifle's each component:\n\n-  Trifle consistently outperforms TT across all beam search hyperparameters and is more robust to variations of both planning horizon $H$ and beam width $W$.\n- (a)Trifle w/ naive rejection sampling >> TT w/ naive rejection sampling \n  (b) Trifle w/o rejection sampling >> TT w/o rejection sampling. \n  In both cases,  Trifle's superior performance originates from that it can positively guide action generation (similar to DT-based Trifle vs DT).\n- (a) Trifle w/ beam search > Trifle w/ naive rejection sampling > Trifle w/o rejection sampling >> TT w/ naive rejection sampling. \n  (b) Trifle w/ naive rejection sampling $\\approx$ TT with beam search. \nAlthough other design choices like rejection sampling/beam search help to better approximate samples from the desired distribution $p(a_t | s_{0:t}, \\mathbb{E}[V_{t}] \\geq v)$, the per-dimension correction terms computed by Trifle used to guide per-dimension-action generation plays a very significant role. \n\n[1] Paster, Keiran, Sheila McIlraith, and Jimmy Ba. \"You can’t count on luck: Why decision transformers and rvs fail in stochastic environments.\"", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "EF08LtvsVB", "reviewer": "Reviewer_dmuC", "summary": "The paper provides empirical evidence that sequence models are able to identify promising actions, but that their policies at inference-time can be suboptimal. The paper proposes to use a tractable probabilistic model to bias the generated action sequences towards optimal ones. The paper provides fairly extensive empirical validation for its claims.\n\nI think the paper could benefit from merging and rephrasing Sections 4 and 5. In its current state, I found it difficult to gather the proposed algorithm. I would also consider adding a (reduced) algorithm box to the main paper.", "questions": "* In Figure 2, what is the definition of the \"optimality score\"?\n* Why is the formulation of eq. (1) preferred over training a policy $\\pi$ which maximizes expected return plus a (KL-)regularization term that encourages it to stay close to the data-generating distribution (such as common in the RLHF literature, e.g., [1])?\n* How do you choose $t'$?\n* Why is the additional rejection sampling step (lines 206-207) necessary given that $\\tilde{p}$ is already proportional to the probability of having large multi-step value?\n* How does the efficiency compare against conditional generative modeling (e.g., DD) and approaches such as QDT, given that TPMs can be slow to train?\n\n\n[1]: Rafailov, Rafael, et al. \"Direct preference optimization: Your language model is secretly a reward model.\" Advances in Neural Information Processing Systems 36 (2024).", "limitations": "See weaknesses and questions above.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "* The experimental evaluation of the claims is thorough and includes a wide selection of relevant baselines.\n* Trifle appears to be on par or slightly better than a number of other approaches.", "weaknesses": "* It is mentioned that a viable alternative to Trifle are sequence models augmented by components learning the Q-function (e.g., QDT). However, Trifle is not compared empirically against these methods. I believe such a comparison would significantly strengthen the paper.\n* I find Figure 2 (middle) slightly misleading. Why does it not directly plot the actual returns (and instead the \"optimality score\")? The correlation between \"optimality score\" and actual returns appears to be pretty weak in the relevant range. Why are middle and right not merged into one single figure?\n* In the Gym-MuJoCo benchmarks, DD (conditional generative modeling) seems to perform roughly as well as Trifle. I believe that a more extensive comparison between the approaches would benefit the paper. Why would Trifle be preferable?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "### Comment #1: comparison with models augmented by Q-function learning components (e.g., QDT)\n\nWe compare the TT-based Trifle with QDT in Appendix 4.1, and DT-based Trifle with QDT in **Table 3(a)** of the rebuttal PDF. The results demonstrate that DT-based Trifle significantly outperforms QDT, supporting our claim that in such scenarios, improving inference-time performance is more critical.\n\nAdditionally, since QDT enhances the training side with more accurate training labels and Trifle improves the inference side with better action sampling, both methods could be combined to achieve even better performance.\n\n### Comment #2: definition of the inference-time optimality score\n\nWe define the \"inference-time optimality score\" as a proxy for inference-time optimality. This score is primarily defined over a state-action pair $(s_t, a_t)$ at each inference step. In Figure 1 (middle) and Figure 1 (right), each sample point represents a trajectory, and the corresponding \"inference-time optimality score\" is defined over the entire trajectory by averaging the scores of all inference steps.\n\nThe inference-time optimality score at time $t$ is defined by the quantile value of $R_t := \\mathbb{E}_{V_t \\sim p^a(\\mathrm{RTG}_t \\mid s_t, a_t)} [V_t]$ ($a_t$ is the action selected by the agent) under the distribution $p(V_t \\mid s_t)$. This value will be high if the sampled action has a high estimated expected return. Please refer to the official comments for more details.\n\n### Comment #3: the correlation between the optimality score and the actual returns\n\n(Following the above response)\nThe above definition uses $p(V_t \\mid s_t)$ ($V_t = RTG_t$) to measure the quality of $a_t$. The validity of this proxy is based on two assumptions. First, in environments like gym-mujoco, the $ RTG_t $ labels provided in the offline dataset are of high quality and are a strong indicator of actual return (justified by Figure 1 (left)). Second, the $p(V_t \\mid s_t)$ fitted using GPT can accurately approximate $p(V_t \\mid s_t) := \\sum_{a_t} p(V_t \\mid s_t, a_t) \\cdot p(a_t \\mid s_t)$. Given the high-dimensional action space, the second assumption is challenging to verify directly. However, Figure 1 (right) shows a clear positive correlation between this proxy and the actual return, indicating that trajectories with higher scores often achieve higher final returns, with the small slope largely attributable to scale differences.\n\nIn summary, we have theoretically and empirically verified the effectiveness of the \"inference-time optimality score\" as a proxy. Figure 1 (middle) compares the inference-time performance of different policies by visualizing the distribution of the \"inference-time optimality score\" over multiple trajectories.\n\n### Comment #4: comparison with DD\n\nWe thank the reviewer for drawing our attention to the interesting connection between Trifle and DD. First, as shown in **Table 3(b)** of the PDF, we highlight that Trifle outperforms DD on 7 out of 9 MuJoCo benchmarks in Table 1 and is more robust with smaller stds.\n\nOn the methodology side, DD adopts diffusion models to fit state-action sequences and then generates rewarding sequences by conditioning on high return. From the inference side, similar to DT, DD aims to sample action conditions on high return, while Trifle aims to improve upon existing RvS algorithms (e.g., TT and DT) by allowing them to sample action conditions on the *expected* return. While DD achieves superior performance compared to some other baselines, their main technical contribution is orthogonal to that of Trifle.\n\nNext, in Section 4.2, we highlight Trifle's effectiveness in handling stochastic environments by efficiently and exactly computing the multi-step value estimate (Eq. (3)). This can significantly improve value estimation in stochastic environments. Since DD also requires a $p (V | s_{0:t}, a_{0:t})$ to guide action sampling, this value component of Trifle can be potentially incorporated into DD. We will include the above discussions in the next version of the paper.\n\n### Comment #5: connection with methods that train an entropy-regularized policy such as DPO\n\nWe thank the reviewer for their insightful comment.\n\nIn our humble opinion, we do not need to prefer one over the other as a policy trained with better objectives (e.g., that trained by some value-based offline RL methods) can be further utilized by RvS algorithms to obtain even better policies. For example, by using the Q values from a pre-trained IQL agent (which implies a better policy) for value estimation, TT can be significantly improved to be much better than both the vanilla TT and IQL.\n\nIn the case of Trifle, the question is whether we want to pay the additional inference-time overhead to achieve better performance. According to the general response, the additional components of Trifle take about 1.45s, which is negligible compared to the runtime of the base algorithm.\n\n### Comment #6: how to choose t'\n\nWe choose $t' = t+3$. See the comments below for a detailed elaboration.\n\n### Comment #7: why is the additional rejection sampling step necessary in lines 206-207\n\nThe necessity of rejection sampling comes from the hardness to draw exact action samples conditioned on high expected returns, which is shown in Theorem 1. For m-Trifle in stochastic envs, we need to compute $\\mathbb{E}[ V_t^{\\mathrm{m}} \\big ]$ using TPM via a post-value estimation step and perform rejection sampling to obtain more unbiased samples from the desired distribution in Eq. (1).\n\n### Comment #8: training time of Trifle and baseline methods\n\nIt takes 30-60 mins to train a PC (the adopted TPM) on one GPU for each Gym-Mujuco task. In comparison, training the GPT model for TT takes approximately 6-12 hours.\n\n### Comment #9: detailed elaboration of the proposed algorithm\nPlease refer to the official comments for a detailed response to this comment, thanks very much!", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "jRFz6cXZUg", "reviewer": "Reviewer_TgpP", "summary": "This paper considers the RvS setting, where a sequence model is learned and used for a control-as-inference policy extraction. Naive sequence models fail due to overly optimistic reward-to-go in stochastic environments. It is possible to alleviate this with rejection sampling, but this suffers from the curse of dimensionality. This paper proposes the use of Tractable Probabilistic Models, which allow the evaluation of probabilistic queries given various degrees of observability. The proposed method uses a beam search procedure to evaluate sequences of actions to locate the best sequence. Results are presented on D4RL Mujoco benchmarks and stochastic game environments.", "questions": "See weaknesses above.", "limitations": "A brief limitations section is present, which somewhat addresses the main limitations of the method.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The paper identifies a key issue in RvS methods, and proposes a new model to address it. It takes advantages of TPM models to evaluate returns under various degrees of masking.\n- The paper is written in a clear way.\n- Empirical results are promising, showing a consistent improvement in both deterministic and stochastic benchmarks.\n- The method is agnostic to the base policy distribution used, meaning it can be easily adapted.", "weaknesses": "- It seems like the strength of TPM comes from beam search conducted on trajectories sampled from a separate policy model (e.g. TT, DT). This comes with the downside of a more costly inference and requires training multiple models. Is there a baseline comparison that utilizes the same computational efficiency?\n- Is there a simple baseline to implement for the beam search value estimator that does not require a specialized TPM? To be thorough, it would be nice to decouple the beam search idea from the use of a TPM.\n- It would be good to define RvS abbreviation early on. \n- More detail on the formulation of the TPM model would clarify the method. How is the model trained and with what architecture?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "### Comment #1: whether beam search is important in Trifle\n\nwe will first clarify the relationship between Trifle and beam search. Then we conduct a comprehensive ablation study on beam search algorithms to confirm Trifle's superiority.\n\nAs mentioned in the general response, the key insight of Trifle to solve challenge #1 is to utilize tractable probabilistic models to better approximate action samples from the desired distribution $p(a_t | s_{0:t}, \\mathbb{E}[V_t] \\geq v)$. We highlight that the most crucial design choice of our method for this goal is that: Trifle can effectively bias the per-action-dimension generation process of any base policy towards high expected returns, which is achieved by adding per-dimension correction terms $p_{TPM} (V_t \\geq v | s_t, a_{t}^{\\leq i})$ (Eq. (2) in the paper) to the base policy.\n\nWhile the rejection sampling method can help us obtain more unbiased action samples through a post value(expected return)-estimation session, we only implement this component for TT-based Trifle (not for DT-based Trifle) for fair comparison, as the DT baseline doesn't perform explicit value estimation or adopt any rejection sampling methods. Therefore, the success of DT-based Trifle strongly justifies that the correction term computed by TPM can effectively bias the actions generated by DT towards higher expected returns.\n\nMoreover, the beam search algorithm comes from TT.  Although it is a more effective way to do rejection sampling, it is not the necessary component of Trifle.\n\n**Ablations over beam search/rejection sampling**\n\nFor TT-based Trifle, we adopted the same beam search hyperparameters as reported in the TT paper (in the official GitHub repo https://github.com/jannerm/trajectory-transformer). And we conduct ablation studies on beam search hyperparameters in **Table 2** of PDF to investigate the effectiveness of Trifle's each component:\n\n- Trifle consistently outperforms TT across all beam search hyperparameters and is more robust to variations of both planning horizon $H$ and beam width $W$.\n- (a)Trifle w/ naive rejection sampling >> TT w/ naive rejection sampling (b) Trifle w/o rejection sampling >> TT w/o rejection sampling. In both cases, Trifle's superior performance originates from that it can positively guide action generation (similar to DT-based Trifle vs DT).\n- Trifle w/ beam search > Trifle w/ naive rejection sampling > Trifle w/o rejection sampling >> TT w/ naive rejection sampling. Although other design choices like rejection sampling/beam search help to better approximate samples from the desired distribution $p(a_t | s_{0:t}, \\mathbb{E}[V_{t}] \\geq v)$, the per-dimension correction terms computed by Trifle used to guide per-dimension-action generation plays a very significant role.\n\n### Comment #2: the downside of a more costly inference and requires training multiple models. \n\nTrifle is efficient in training. It only takes 30-60 minutes (~20s per epoch, 100-200 epochs) to train a PC on one GPU for each Gym-Mujuco task (**Note that a single PC can be used to answer all conditional queries required by Trifle**). In comparison, training the GPT model for TT takes approximately 6-12 hours (80 epochs).\n\nTo evaluate inference-time efficiency, we conduct a more detailed runtime analysis, and the main results are shown in Figure 1 of the attached PDF. Figure 1 (left) expands Table 5 and plots the stepwise inference-time scaling curve of Trifle vs TT with varying horizons. We can see that: As we increase the horizon, the relative slowdown is mitigated. This is because, across different horizons, **TPM-related computation consistently requires ~1.45s computation time**, which additional computation overhead is diminishing as we increase the beam horizon. Figure 1 (right) shows that Trifle's runtime (TPM-related) scales **linearly** w.r.t. the number of action variables, which indicates its efficiency for handling high-dimensional action spaces.\n\nNote that there are recent breakthroughs [1] in designing efficient PC implementations, which can significantly speed up the computation of Trifle (both training and inference).\n\n[1] Liu, Anji, Kareem Ahmed, and Guy Van den Broeck. \"Scaling Tractable Probabilistic Circuits: A Systems Perspective.\" arXiv preprint arXiv:2406.00766 (2024).\n\n### Comment #3: define the RvS abbreviation earlier\n\nWe thank the reviewer for the suggestion and will revise the paper accordingly.\n\n### Comment #4: details on the formulation, structure, and training algorithm of the TPM model\n\nWe thank the reviewer for their suggestion. We provide a concise and intuitive introduction to the TPM we adopted: Probabilistic Circuits (PCs). We describe the representation and structure of PCs in Section 2 and provide a detailed and formal introduction in Appendix B. We adopted the Hidden Chow-Liu Tree PC structure and used EM to train the model. Please refer to Appendix B.2 for more details about the adopted PC structures and the parameter learning algorithm.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank all reviewers for their constructive feedback and for acknowledging our paper as novel, well-presented, and comprehensively evaluated. We summarize common questions and concerns raised by the reviewers in the following. 

**Key technical differences of Trifle compared to other RvS or offline RL algorithms**

As highlighted in our paper as well as in the literature (e.g., [1]), there are several challenges in RvS algorithms from a tractable inference perspective. This paper identifies two such challenges and proposes to solve them using **tractable probabilistic models**.

**Challenge #1**: Sampling actions condition on high **expected** return

Given a generative model of the joint probability distribution $p(s_{0:t}, a_{0:t}, V_{0:t})$, our goal is to sample actions $a_t$ condition on the current state $s_{0:t}$ and high expected future returns (i.e., high $\mathbb{E}[V_t]$), which can be concisely written as $p(a_t | s_{0:t}, \mathbb{E}[V_t] \geq v)$ (Eq. (1) in the paper). However, the condition on the expected future return makes it a hard inference task given generative models over the joint distribution. We show that this query is NP-hard to compute exactly even when the joint distribution follows Naive Bayes (Thm. 1).  We propose a promising approximate sampling algorithm for $p(a_t | s_{0:t}, \mathbb{E}[V_t] \geq v)$ that leverages the ability of TPMs to compute arbitrary conditional distributions. This component can be directly used to replace policy components such as $p(a_t | s_{0:t})$ (e.g., TT) and $p(a_t | s_{0:t}, V_t = v)$ (e.g., DT) used in existing RvS algorithms. Empirical evaluations demonstrate the addition of this component significantly improves over the base algorithm.

**Challenge #2**: Estimating state-action values under stochastic transition dynamics

Another major challenge comes from the inaccuracy of the RTG), an estimate of the return, provided in the offline datasets. This problem is even worse in stochastic environments, where the labeled RTGs have very high variation. To combat this inaccuracy, it is common to use multi-step value estimates such as TD(1) and TD(λ). Specifically, we want to compute $\mathbb{E} [r_t + \cdots + r_{t'} + V_{t'+1} | s_{0:t}, a_{t:t'}]$ give a joint distribution represented by a generative model. However, computing such values requires implicitly marginalizing out future states $s_{t+1:t'}$, which is computationally intractable for most models. We propose to use TPM to exactly compute such terms, which leads to a significant performance gain in stochastic environments.

**The relationship between Trifle and rejection sampling/beam search algorithm**

The key insight of Trifle to solve challenge #1 is to utilize tractable probabilistic models to better approximate action samples from the desired distribution $p(a_t | s_{0:t}, \mathbb{E}[V_t] \geq v)$. We highlight that the most crucial design choice of our method for this goal is that: Trifle can effectively bias the per-action-dimension generation process of any base policy towards high expected returns, which is achieved by adding per-dimension correction terms $p_{TPM} (V_t \geq v | s_t, a_{t}^{\leq i})$ (Eq. (2) in the paper) to the base policy. 

While the rejection sampling method can help us obtain more unbiased action samples through a post value(expected return)-estimation session, we only implement this component for TT-based Trifle (not for DT-based Trifle) for fair comparison, as the DT baseline doesn't perform explicit value estimation or adopt any rejection sampling methods. Moreover, the beam search algorithm also comes from TT. Although it is a more effective way to do rejection sampling, it is not the necessary component of Trifle, either. 

Next, following the suggestions of the reviewers, we conduct a comprehensive ablation study over beam search/rejection sampling in Table 2 of the attached PDF. The results:
- **Trifle with beam search > Trifle with naive rejection sampling > Trifle w/o rejection sampling >> Base policy w/o rejection sampling**
- **Trifle w/o rejection sampling >> TT with naive rejection sampling** 
- **Trifle with naive rejection sampling $\approx$ TT with beam search**

strongly justifies our claim. (We analyze the results in detail to each reviewer)

**Computational efficiency/scalability of Trifle**

For the inference-side efficiency, we conduct a more detailed inference-time evaluation, and the main results are shown in Figure 1 of the attached PDF. Specifically, Figure 1 (left) expands Table 5 and plots an inference-time scaling curve of Trifle vs TT with varying horizons, where we can draw the same conclusion as Appendix D.3 of the paper: Since TPM-related computation consistently requires ~1.45s computation time across different horizons, the relative slowdown of Trifle is diminishing as we increase the beam horizon. Figure 1 (right) shows that Trifle's runtime (TPM-related) scales **linearly** w.r.t. the number of action variables, which indicates its efficiency for handling high-dimensional action spaces. 

Trifle is also efficient in training. It only takes 30-60 minutes (~20s per epoch, 100-200 epochs) to train a PC on one GPU for each Gym-Mujuco task (**Note that a single PC model can be used to answer all conditional queries required by Trifle**). In comparison, training the GPT model for TT takes approximately 6-12 hours (80 epochs). 

Note that there are recent breakthroughs [2] in designing efficient PC implementations, which can **significantly speed up** the computation of Trifle (both training and inference). 

[1] Paster, Keiran, Sheila McIlraith, and Jimmy Ba. "You can’t count on luck: Why decision transformers and rvs fail in stochastic environments."

[2] Liu, Anji, Kareem Ahmed, and Guy Van den Broeck. "Scaling Tractable Probabilistic Circuits: A Systems Perspective."

Meta-review: The reviewers generally liked the paper and has unanimous positive reviews so we are recommending acceptance of this paper. That said, we would like to point out that it is unclear if the paper has sufficient comparisons to demonstrate the utility of the framework, for instance, results on the more challenging antmaze tasks are absent, on other gym environments the performance gains are tiny, so this is largely a simple framework. That said, the didactic experiments are nice. I would recommend acceptance conditioned on adding a more extensive set of results for the final version.
Decision: Accept (poster)