A Tractable Inference Perspective of Offline RL

Comment #1: whether beam search is important in Trifle

we 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.

As 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.

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. 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.

Moreover, 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.

Ablations over beam search/rejection sampling

For 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:

• Trifle consistently outperforms TT across all beam search hyperparameters and is more robust to variations of both planning horizon $H$ and beam width $W$.
• (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).
• 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.

Comment #2: the downside of a more costly inference and requires training multiple models.

Trifle 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).

To 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.

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

[1] Liu, Anji, Kareem Ahmed, and Guy Van den Broeck. "Scaling Tractable Probabilistic Circuits: A Systems Perspective." arXiv preprint arXiv:2406.00766 (2024).

Comment #3: define the RvS abbreviation earlier

We thank the reviewer for the suggestion and will revise the paper accordingly.

Comment #4: details on the formulation, structure, and training algorithm of the TPM model

We 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.

Thank you for your thoughtful feedback and appreciation of our response, and we appreciate your recognition of our work as solid. We will revise the paper to include the additional experimental details and a more comprehensive description of TPMs, as suggested. If you have any further questions, we are pleased to discuss them with you.

Comment #1: comparison with models augmented by Q-function learning components (e.g., QDT)

We 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.

Additionally, 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.

Comment #2: definition of the inference-time optimality score

We 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.

The 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.

Comment #3: the correlation between the optimality score and the actual returns

(Following the above response) The 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.

In 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.

Comment #4: comparison with DD

We 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.

On 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.

Next, 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.

Comment #5: connection with methods that train an entropy-regularized policy such as DPO

We thank the reviewer for their insightful comment.

In 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.

In 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.

Comment #6: how to choose t'

We choose $t' = t+3$. See the comments below for a detailed elaboration.

Comment #7: why is the additional rejection sampling step necessary in lines 206-207

The 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).

Comment #8: training time of Trifle and baseline methods

It 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.

Comment #9: detailed elaboration of the proposed algorithm

Please refer to the official comments for a detailed response to this comment, thanks very much!

The specific steps for calculating the inference-time optimality score for a given inference step $t$, given $s_t$ and a policy $p(a_t \mid s_t)$, are as follows:

1. Given $s_t$, sample $a_t$ from $p_{TT}(a_t \mid s_t)$, $p_{DT}(a_t \mid s_t)$, or $p_{Trifle}(a_t \mid s_t)$.
2. Compute the state-conditioned value distribution $p^s(V_t \mid s_t)$.
3. Compute $R_t := \mathbb{E}_{V_t \sim p^a(RTG_t \mid s_t, a_t)} [V_t]$, which is the corresponding estimated expected value.
4. Output the quantile value $S_t$ of $R_t$ in $p^s(V_t \mid s_t)$.

To approximate the distributions $p^s(V_t \mid s_t)$ and $p^a(V_t \mid s_t, a_t)$ (where $V_t = RTG_t$) in steps 2 and 3, we train two auxiliary GPT models using the offline dataset. For instance, to approximate $p^s(V_t \mid s_t)$, we train the model on sequences $(s_{t-k}, V_{t-k}, \ldots, s_t, V_t)$.

Intuitively, $p^s(V_t \mid s_t)$ approximates $p(V_t \mid s_t) := \sum_{a_t} p(V_t \mid s_t, a_t) \cdot p(a_t \mid s_t)$. Therefore, $S_t$ indicates the percentile of the sampled action in terms of achieving a high expected return, relative to the entire action space.

Thank the reviewer for the suggestion. We will include a subsection to describe the algorithm and also provide an algorithm table in the next version of the paper.

The high-level idea of Trifle is that it can be built upon many RvS algorithms (e.g., TT, DT) to solve two key problems that are widely observed in the literature: (i) sample actions condition on high expected return, and (ii) estimate state-action values under stochastic transition dynamics. In the following, we first explain what modifications Trifle applied to solve both challenges; we then describe how the modifications are used to design algorithms upon existing RvS algorithms.

Scenario 1: sampling actions condition on high expected return

As 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)). 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.

In general, we first train a PC (the adopted TPM) to fit the joint distribution over $p(s_{t-k}, a_{t-k}, V_{t-k},...,s_t, a_{t}, V_{t})$ from the offline dataset. This step is agnostic to the base RvS algorithm Trifle is built on. Then, given any component that computes or samples from $p(a_t | s_t, V_t = v)$ (DT) or $p(a_t | s_t)$ (TT), we replace it to a good approximation of $p(a_t | s_t, \mathbb{E}[V_t] \geq v)$ by augmenting these proposal distributions with a correction term.

Specifically, we utilize the pretraind PC to compute the per-dimension correction term $p_{TPM} (V_t \geq v | s_t, a_{t}^{\leq i})$ by marginalizing out unseen action variables $a_t^{i+1},...,a_t^k$ and this forms an exponentially large action space, aiming to bias the actions generated by the base policy towards high expected return $\mathbb{E}[V_t]$. Notably, we query this conditional probability for each dimension $a_t^i$, as the inference proceeds, the variable $i$ will increase progressively, resulting in different marginal queries. However, thanks to PCs' tractabiltiy, we can use a single model to answer arbitrary marginalization queries. Moreover, the computation can be done exactly without any approximation and scales linearly to the number of action variables, which is highly efficient. In Appendix C.2, we present the concrete algorithm of how PCs (the adopted TPM) compute conditional probabilities.

Note that the rejection sampling methods via beam search steps is 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 to ensure fair comparision.

Scenario 2: estimating state-action values under stochastic transition dynamics

Many existing RvS algorithms requires explicitly evaluating $p(V_t | s_{0:t}, a_{t})$. However, the single-step value estimate $V_t = \mathrm{RTG}_t$ will be very inaccurate under stochastic environments. 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.

While 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 Eq. (3). This allows us to enjoy the better accuracy of the multi-step value estimates while mitigating the inefficiency to compute them.

For m-Trifle in stochastic Taxi environment, we largely follow the Algorithm 3 of Appendix 3.1, here's a more concise and complete version:

(1) choose t'>t and compute $p_{GPT}(a_t | s_t)$, $p_{TPM}(V_{t'}>v | s_t, a_t)$

(2) sample N action candidates from $p_{Trifle}(a_t|s_t) = p_{GPT}(a_t | s_t)\times p_{TPM}(V_{t'}>v | s_t, a_t)$

(3) for each action candidate $a_t$, sample future actions $a_{t+1},...,a_t'$ autoregressively from $p_{Trifle}$ for future evaluation.

(4) for each action candidate and each $h \in [t+1,t']$ , compute $p_{TPM} (r_h\ | s_t, a_{t:h})$ by marginalizing over intermediate states $s_{t+1:h}$ and also $p_{TPM} (V_{t'}\ | s_t, a_{t:t})$ , where $V_{t'} = RTG_{t'}$

(5) for each action candidate, compute： $E[ V_t^{m} ] = \sum_{h=t}^{t'} \gamma^{h-t} E_{r_h \sim p_{TPM} (\cdot | \tau_{\leq t}, a_{t+1:h})} [ r_h ] + \gamma^{t'+1-t} E_{RTG_{t'} \sim p_{TPM} (\cdot | \tau_{\leq t}, a_{t+1:t'})} [ V_{t'} ]$

(6) select the action with maximum $\mathbb{E}[ V_t^m ]$ to excute in the environment.

We choose $t' = t+3$.

Comment #1: connection between rejection sampling and the correction terms, and the corresponding computational complexity

We 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.

In 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.

Challenge 1: Sampling actions conditioned on high expected return

As 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.

To 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.

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

The 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.

While 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.

Ablation studies on the rejection sampling/beam search

For 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:

• Trifle consistently outperforms TT across all beam search hyperparameters and is more robust to variations of both planning horizon $H$ and beam width $W$.
• (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).
• (a) Trifle w/ beam search > Trifle w/ naive rejection sampling > Trifle w/o rejection sampling >> TT w/ naive rejection sampling. (b) Trifle w/ naive rejection sampling $\approx$ TT with beam search. 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.

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

Dear Reviewer DwFR,

Thank you again for your valuable feedback and suggestions. In our previous response, we have provided a detailed explanation and supplemented our arguments with theoretical analysis and experimental validation. To help you better understand the improvements we made, I would like to summarize the key points:

1. One of our main contributions: Our method introduces per-action-dimension correction terms that can be exactly and efficiently calculated by TPMs to enhance the action sampling procedure. These correction terms significantly improve the performance, and our method does not necessarily need to be combined with rejection sampling or beam search to be effective. For example, DT-based Trifle achieves significant performance gains compared to DT without using any rejection sampling, proving that the per-action-dimension correction terms play a critical role in guiding the action generation process.

2. Additional ablation study over rejection sampling: We provided additional ablation studies (Table 1 in the attached PDF) to show that

   (i) The proposed correction term works well across different base algorithms (DT and TT). Specifically, as indicated by the last row of Table 1(b) in tha attached PDF, TT w/o rejection sampling suffers a significant performance drop, while Trifle works well even w/o rejection sampling (competitive results as TT w/ beam search).

   (ii) Trifle can be effectively combined with rejection sampling or beam search to further boost performance, which leads to state-of-the-art results in 7 out of 9 Gym-MuJoCo environments.

3. Another equally important contribution: Our method enables exact and efficient multi-step value estimation, which significantly enhances performance in stochastic environments as shown in Section 6.2.

We believe we have addressed the concerns raised, but if you have any further questions or issues, we would be happy to discuss them with you. We look forward to your feedback.

Thank you again for your time and consideration.

the adaptive thresholding mechanism is adopted when computing the term $p_{TPM} (V_t \geq v | s_t, a_{t}^{\leq i})$ of Equation (2), where $i \in \{1, \dots, k\}$, $k$ is the number of action variables and $a_t^{i}$ is the $i$th variable of $a_t$. Instead of using a fixed threshold $v$, we choose $v$ to be the $\epsilon$-quantile value of the distribution $p_{TPM}(V_t|s_t,a_{t}^{< i})$ computed by the TPM, which leverage the TPM's ability to exactly compute marginals given incomplete actions (marginalizing out $a_{t}^{i:k}$). Specifically, we compute $v$ using $v = max_r\{r\in\mathbb{R}|p_{TPM}(V_t\geq r|s_t,a_{t}^{< i})\geq 1-\epsilon\}$. Empirically we fixed $\epsilon$ for each Gym-MuJoCo environment and $\epsilon = 0.2$ or $0.25$, which is selected by runing grid search on $\epsilon \in [0.1,0.25]$.

Comment #1: practical implications of the theoretical guarantees in Secs. 4.1 and 4.2

Thanks 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.

Sampling actions conditioned on high expected return

As 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.

Specifically, 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.

Estimating state-action values under stochastic transition dynamics

As 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.

Comment #2: effect of the beam search hyperparameters

We conduct additional ablation studies on both the adaptive thresholding mechanism and rejection sampling/beam search.

Ablation studies on the adaptive thresholding mechanism

We 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.

We 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.

Ablation studies on the rejection sampling/beam search

As 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.

For 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:

• Trifle consistently outperforms TT across all beam search hyperparameters and is more robust to variations of both planning horizon $H$ and beam width $W$.
• (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.
• 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.

Comment #3: computational complexity analysis

Thanks for the valuable suggestion.

First, we conduct a more detailed runtime analysis, and the main results are shown in Figure 1 of the attached PDF.

Figure 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.

Moreover, 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).

Note that there are recent breakthroughs [1] on designing efficient PC implementations, which can significantly speed up training and inference of Trifle.

[1] Liu et. al. "Scaling Tractable Probabilistic Circuits: A Systems Perspective." arXiv preprint arXiv:2406.00766 (2024).

Comment #4: evaluation on the Antmaze task

We 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.

Comment #5: potential negative societal impact

Thanks. We will include the following discussion:

This 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.

The scalability of Trifle to more complex tasks is also an interesting question. Some of the key challenges to scaling up an offline RL algorithm include (i) high-dimensional action spaces, and (ii) more complex transition dynamics and policy. 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. The second largely depends on the performance of modeling complex sequences. While we do not have a definite answer on the scalability of the adopted TPMs, there are abundant recent papers that significantly improve the expressive power of TPMs.

Thank you very much for supporting our paper and increasing the score! And thanks for your valuable suggestions during the rebuttal period to help us improve the paper. We will incorporate the additional details in the final version according to your comments. We sincerely appreciate your feedback!