SlowGuess/NeurIPS_2023_OCR huggingface dataset:
3.2 Comparison to regularized and one-step offline RL
One of the most used methods in offline RL methods is adding policy or value regularization and constraints terms on top of a vanilla off-policy RL algorithm [5, 8, 12, 22, 23, 25, 26, 50, 51], referred to as regularized methods in this paper. Our method can be viewed as an alternative to using regularized losses. Instead, we enforce a non-Markov budget constraint on the policy class, and we argue this provides several unique advantages. First, compared with the coefficients of regularization terms, the budget parameter has a more clear physical explanation in most applications. Thus it is easier to tune the hyper-parameters and explain the decisions from the policy in practice. Second, the constraints on budget will never be violated in the test, it provides an additional level of safety. While regularization terms penalize the divergence during training, they cannot provide a guarantee on the divergence or distance between test policy and behavior policy.
Another type of offline RL that is related to our core idea is one-step RL [3, 49]. We propose and leverage the concept of limiting the number of counterfactual/off-policy steps. However, the "step" here refers to the decision steps, and in one-step RL it refers to the (training) iteration step. Although one-step RL only applies only one training step, the resulting policy can still be far away from the behavior policy in many different states, and give different decisions during the test. From another perspective, our method, although with a limited number (or even one) of counterfactual decision steps, still applies dynamic programming to find the best allocation of counterfactual steps.
3.3 Important implementation details
Algorithm 1 provided a general actor-critic framework to optimize policy value constrained by counterfactual decisions. To implement a practical offline deep RL algorithm, the first design choice is how we model the policy and choose the policy gradient estimator to  $\nabla_{\phi}\mathcal{L}_{\pi}$ . Algorithm 1 is compatible with any off-policy policy gradient methods. We implement Algorithm 1 using SAC-style [17] policy gradient for stochastic policy and TD3-style [13] policy gradient for deterministic policy, to provide a generic algorithm for both stochastic and deterministic policy models. As TD3 uses deterministic policy, the value of  $m$  is 1 in  $\widehat{T}_{\mathrm{CB}}$  and the actor loss becomes  $Q_{\theta}(s,b,\pi_{\phi}(s,b))$ . In both cases, we implement target updates following SAC/TD3. We defer more details to the appendix.
Both SAC and TD3 are originally proposed for online, off-policy RL. To better fit into the offline setting, previous offline RL methods based on these algorithms equip them with several key adaptations or implementation tricks. To eliminate orthogonal factors and focus on the budgeting idea in algorithm comparison, we follow these adaptations and describe them below. For SAC, prior work use twin Q functions [26], and a linear combination of the two  $Q$  values [8]. Prior work [8, 24] also sample  $m$  actions from the actor and take the maximum  $Q$  values in backup, instead of a standard actor-critic backup, which is a trick firstly introduced in [12, 15]. The entropy term in SAC is dropped in [8] as this term is mostly for online exploration. Previous TD3-based work [11] normalizes the state features and  $Q$  losses in TD3. All these adaptations are commonly used in state-of-the-art offline RL and are mostly considered as minor implementation details rather than major algorithmic designs. We refer to SAC and TD3 with all these adaptations as offline SAC and offline TD3 in this paper and refer to the two families of algorithms SAC-style and TD3-style.
One-step RL[3] and IQL [23] for offline RL is built on top of estimating the behavior value  $Q^{\mu}$  by SARSA. Unlike SAC or TD3, it is not applicable to apply the counterfactual budgeting idea on top of SARSA, since it is a native on-policy method and does not consider counterfactual decisions. Thus we focus on the implementation of BCOL with SAC and TD3.
Task Name	BC	10%BC	DT	RAMBO	ARMOR	IQL	Onestep	TD3 +BC	BCOL (TD3)	CQL	CDC	BCOL (SAC)
halfcheetah-m	42.6	42.5	42.6	77.6	54.2	47.4	55.6	48.4	45.0	46.1	62.5	50.1
hopper-m	52.9	56.9	67.6	92.8	101.4	66.3	83.3	59.4	85.8	64.6	84.9	83.2
walker2d-m	75.3	75.0	74.0	86.9	90.7	78.3	85.6	84.5	76.7	74.5	70.7	84.1
halfcheetah-mr	36.6	40.6	36.6	68.9	50.5	44.2	42.5	44.4	40.9	45.4	52.3	46.2
hopper-mr	18.1	75.9	82.7	96.6	97.1	94.7	71.0	50.1	83.4	92.3	87.4	99.8
walker2d-mr	26.0	62.5	66.6	85.0	85.6	73.9	71.6	80.2	49.7	83.7	87.8	86.0
halfcheetah-me	55.2	92.9	86.8	93.7	93.5	86.7	93.5	91.5	88.7	87.3	66.3	86.9
hopper-me	52.5	110.9	107.6	83.3	103.4	91.5	102.1	100.5	106.8	109.2	83.2	99.0
walker2d-me	107.5	109.0	108.1	68.3	112.2	109.6	110.9	110.1	108.5	109.9	103.9	110.9
mujoco total	466.7	666.2	672.6	753.1	788.6	692.4	716.0	669.2	685.6	713.0	699.0	746.0
antmaze-u	54.6	62.8	59.2	25.0	-	87.5	64.3	96.3	93.3	94.0	93.6	90.3
antmaze-u-d	45.6	50.2	53.0	16.4	-	62.2	60.7	71.7	68.0	47.3	57.3	90.0
antmaze-m-p	0.0	5.4	0.0	0.0	-	71.2	0.3	1.7	12.3	62.4	59.5	70.0
antmaze-m-d	0.0	9.8	0.0	0.0	-	70.0	0.0	0.3	14.0	74.3	64.6	72.3
antmaze-l-p	0.0	0.0	0.0	23.2	-	39.6	0.0	0.0	0.0	34.2	33.0	35.6
antmaze-l-d	0.0	6.0	0.0	2.4	-	47.5	0.0	0.3	0.0	40.7	25.3	37.6
antmaze total	100.2	134.2	112.2	67.0	-	378.0	125.3	171.3	187.7	352.9	333.5	396.0
Total	566.9	800.4	784.8	820.1	-	1070.4	841.3	840.2	873.3	1065.9	1032.5	1142.0
Table 1: Average normalized scores on D4RL tasks. Task names for MuJoCo: m=medium, mr=medium replay, me=medium expert. Task names for antmaze: u=umaze, m=medium, l=large, p=play, d=diverse. Bold numbers stand for the globally best and underlined numbers stand for the best with in the base method (SAC or TD3) group.
Looping over all budget values in the loss is not efficient in practice. We implement the  $Q$  function and policy with  $B$  output heads and tensorize the loop over  $b$ . One can further reduce the computation cost by sampling budget values instead of updating all  $B$  heads over every sample.
We evaluate our BCOL algorithm against prior offline RL methods on the OpenAI gym MuJoCo tasks and AntMaze tasks in the D4RL benchmark [9]. We compare the SAC-style and TD3-style implementation of BCOL with state-of-the-art offline RL algorithms, studying the effectiveness of budgeting counterfactual in offline RL. Our experiment results also reveal that behavior cloning with only one strategic counterfactual decision still works surprisingly well on MuJoCo tasks. Finally, we study the value of our dynamic programming methods on where to spend the counterfactual budget.
Baselines. We compared BCOL with regularized methods (policy and/or value regularization) based on SAC and TD3: CQL [26], CDC [8], TD3+BC [11]. Besides regularized methods, we also compare with one-step RL method [3] and IQL [23] as other state-of-the-art model-free offline RL methods, RAMBO [42] and ARMOR [52] as state-of-the-art in model-based offline RL, and behavior cloning, behavior cloning on the top  $10\%$  data, and decision transformer as imitation learning baselines. This covers a set of strong offline RL baselines. We are interested in both studying the effectiveness of counterfactual budget ideas in contrast to the closest regularized methods and state-of-the-art offline RL methods. We present a set of representative and the most performant offline RL baselines in this section, and defers the comparison against more offline RL baselines results to the Appendix due to the limit of space.
Benchmark. We report the results of BCOL together with baselines on 9 OpenAI gym MuJoCo tasks and 6 AntMaze tasks in D4RL. The gym MuJoCo tasks consist of the v2 version of medium, medium-reply, and medium-expert datasets in halfcheetah, walker2d, and hopper. MuJoCo tasks generally prefer imitation-type algorithm, as it contains a large portion of near-optimal trajectories. The 6 AntMaze tasks are considered harder tasks for offline RL as they contain very few or no near-optimal trajectories, and many previous methods lack performance reports on these tasks. We exclude the random datasets in MuJoCo tasks as they are not less discriminating and none of the offline RL algorithms learns a meaningful policy there. We exclude expert datasets in MuJoCo tasks since they can be solved simply by behavior cloning and mismatches the consideration in most offline RL algorithm designs.
How we report the results. Due to the inconsistency in the version of D4RL in the literature, it is necessary to clarify the source of baseline results here. We retrieve the results of baseline methods from the original papers if applicable (IQL, CQL, one-step RL, DT, RAMBO, ARMOR). We report the Rev. KL Reg variant of one-step RL as it shows the best performance on MuJoCo tasks [3]. For AntMaze tasks we report one-step RL and decision transformer results from the IQL paper as they
Figure 2: Total normalized score with different values of  $B$  and  $\omega$  in BCOL. The left two plots show MuJoCo average scores and the right two plots show AntMaze average scores.
are not reported in the original paper. We report the scores of behavior cloning and  $10\%$  behavior cloning. We report the scores of behavior cloning from the IQL paper. We report the scores of CQL from an updated version of CQL paper for D4RL v2 environments. We report the score of TD3+BC and CDC from our implementation since the original results are for D4RL v0 environments. The v2 results are generally better than the original v0 scores in the original papers and are consistent with TD3+BC's v2 scores from the IQL paper.
Table 1 shows the results of baseline algorithms as well as two variants of BCOL (TD3-style and SAC-style) on 15 D4RL tasks. For the scores from our experiment (BCOL, CDC, TD3+BC), we report the test episodic reward averaged over 300 test episodes as 10 test episodes per evaluation, the last 10 evaluations, and 3 runs with different random seeds. All algorithms run with 1M steps following the literature. We defer details like learning curves and standard deviations in the appendix.
Main results. As Table 1 shows, BCOL-SAC outperforms prior offline RL methods for both MuJoCo and AntMaze total scores. BCOL never falls behind the best methods by a large margin in any task. This indicates the effectiveness and robustness of the counterfactual budget in different scenarios. Especially for the total score of AntMaze, the harder tasks, BCOL-SAC outperforms most prior approaches by a large margin except IQL. We also find for both TD3-style and SAC-style algorithms, BCOL outperform the regularized method based on the same architecture, abating orthogonal factors other than the key idea about the counterfactual budget.
Hyper-parameters:  $B$  and  $\omega$ . Our algorithm only adds two hyper-parameters on top of SAC/TD3: budget  $B$  and  $\omega$ . We searched the value of  $B$  in  $\{1, 10, 50\}$  and the value of  $\omega$  in  $\{0, 1, 10, 100\}$ . We select one set of hyper-parameters for MuJoCo (SAC-style:  $B = 10, \omega = 10$ , TD3-style:  $B = 50, \omega = 10$ ) and one set for AntMaze (SAC-style and TD3-style:  $B = 50, \omega = 0$ ) based on the overall performance. This provides a fair comparison with baseline methods as all of them select their hyper-parameters either as this, or per task. Figure 2 shows how the total scores changes with different values of  $B$  and  $\omega$ . Full results with scores on each task are in the appendix.
Figure 2 also shows another surprising finding from our experiments on the MuJoCo tasks. With the budget being only one  $B = 1$ , BCOL (SAC) shows comparable performance to some recent offline RL work including CDC and IQL. With 10 steps out of 1000 steps, it is able to outperform most prior methods. This highlights how strategically adding a few important actions on top of behavior policy can provide strong performance in this benchmark. It also indicates that the MuJoCo tasks in D4RL, even the more mixed datasets, considerably prefer the imitation-type approach.
Ablation on how we realize the counterfactual budgeting constraints. The idea of budgeting counterfactual decisions naturally requires both the planning on budgeting in Algorithm 1 during the training and the planning on budgeting in Algorithm 2 during the testing. Ablation on the roles of budgeting itself and two planning parts is nonetheless helpful for us to understand how BCOL works. In Figure 3, we report the results of three ablation methods.
The first column shows the performance without budgeting constraints. This leads to offline TD3 or SAC algorithm with all implementation adaptations used by BCOL and others [11, 12, 26]. As we expected, vanilla offline RL without budgeting does not work well. The second column shows the performance with budgeting but without any planning on budgeting during either the training or the testing. This method realizes the budgeting by randomly assigning  $B$  decisions to the policy
<sup>1</sup>https://sites.google.com/view/cql-offline-rl
$^{2}\mathrm{TD3 + BC}$  paper does not include results and hyperparameters in AntMaze. We do the hyper-parameter search of their  $\alpha$ , within the range provided in the paper, and report the highest total score with  $\alpha = 3$ .
Figure 3: Percent difference of the performance on different budgeting methods compared with the full BCOL Algorithm (hc = HalfCheetah, hop = Hopper, w = Walker2d, am=AntMaze). The top row shows SAC-based experiments and the bottom row shows TD3-based experiments. TD3 plots do not include AntMaze-large tasks since the performances of BCOL are zero. No budgeting stands for offline SAC/TD3 without the budgeting constraints (equivalent to  $B \rightarrow \infty$ ). Budgeting without planning stands for randomly selecting  $B$  steps to follow from  $\pi$  and the rest from  $\hat{\mu}$  during the test, where  $\pi$  is learned by offline SAC/TD3. Budgeting without test-time planning stands for randomly selecting  $B$  steps (uniformly within the max horizon) to follow from  $\pi$  and the rest from  $\hat{\mu}$  during the test, where  $\pi$  is learned by Algorithm 1. In all settings,  $B$  is the same value as selected by BCOL.
learned by offline SAC/TD3, and the rest to the estimated behavior policy. Randomly stitching the counterfactual policy and behavior policy fails to fully realize the benefit of budgeting, since both training and inference are not budget-aware. The experiment in the third column studies if it is sufficient to be budget-aware during training. This ablation randomly selects  $B$  actions from  $\pi$  trained by BCOL and the rest from the estimated behavior policy. The setting is closest to BCOL, but the lack of planning on how to spend the budget during the testing still hurts the performance. The results on SAC-style and TD3-style implementations are consistent. This ablation shows that the effectiveness of BCOL relies on the collaboration of all three parts of our core idea: budgeting, planning on budgeting during the training, and planning on budgeting during the testing.
Most recently proposed offline reinforcement methods rely on some mechanism to force the learned counterfactual decision policy to stay close to the data support. One approach to this end is regularizing the policy by its divergence with data, either by parameterization [12, 15], constraints and projection [29, 32, 44], divergence regularization [10, 11, 25], or implicit regularization by weighted regression [10, 37, 39, 40, 50]. The other similar idea, which is often applied together, is regularizing the value estimate of policy in an actor-critic or Q learning architecture [5, 7, 8, 22, 26, 28, 33, 34, 51]. Similar regularization can also be realized by the ensemble of  $Q$  functions [1, 2, 14].
Some recent batch RL methods focus on more adaptive policy, similarly to this paper, but conditioning on history [16] or confidence level [19]. A similar idea to us of making a few key decisions in one trajectory was studied [18, 46], but they focus on stitching the logged trajectories in data.
In contrast to batch RL, imitation learning often does not involve dynamic programming on extrapolated actions. It has been used as a regularization in online RL [31, 38] and offline RL [11, 37]. Recently imitation learning methods using transformers and conditional training also achieves good performance on offline RL benchmarks [4, 56]. One-step RL [3] can also be viewed as an extension
of imitation learning with one-step policy extraction [49]. Although utilizing the behavior policy as well, Our key idea is different from these as we train a policy with awareness of behavior policy, rather than regularized by behavior cloning or imitation learning loss.
The form of budget considered in this work is counting the different decision distributions. A natural alternative is divergences between  $\pi$  and  $\mu$ . This provides a soft count of counterfactual decisions. However, it requires good calibration of both distributions and is more suitable to the discrete-action settings. We leave this as future work. On the theory side, a natural question about the value of budgeting counterfactuals is how the benefit of less counterfactual decisions is reflected in the theoretical properties of offline RL. We leave the investigation on this for future theoretical work.
In conclusion, this paper studies offline reinforcement learning which aims to learn a good counterfactual decision policy from a fixed dataset. We propose a novel idea of budgeting the number of counterfactual decisions and solving the allocation problem by dynamic programming. We provide strong empirical performance over offline RL benchmarks and optimality of the fixed point solution as theoretical justification.
[55] Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. Mopo: Model-based offline policy optimization. arXiv:2005.13239, 2020.
[56] Qinqing Zheng, Amy Zhang, and Aditya Grover. Online decision transformer. In International Conference on Machine Learning, pages 27042-27059. PMLR, 2022.
A Proof of Theorem 2
Proof. We first prove that  $\mathcal{T}_{\mathrm{CB}}$  is a  $\gamma$ -contraction with respect to  $\| \cdot \|_{\infty}$ . For any function  $Q_{1}, Q_{2}$  in the space, and any  $s \in S, b > 0$ :
$$
\begin{array}{l} \left| V _ {Q _ {1}} (s, b) - V _ {Q _ {2}} (s, b) \right| (10) \\ = \left| \max  \left\{\max  _ {a} Q _ {1} (s, b - 1, a), \mathbf {E} _ {a \sim \mu} Q _ {1} (s, b, a) \right\} - \max  \left\{\max  _ {a} Q _ {2} (s, b - 1, a), \mathbf {E} _ {a \sim \mu} Q _ {2} (s, b, a) \right\} \right| (11) \\ \leq \max  \left\{\max  _ {a} | Q _ {1} (s, b - 1, a) - Q _ {2} (s, b - 1, a) |, \mathbf {E} _ {a \sim \mu} | Q _ {1} (s, b, a) - Q _ {2} (s, b, a) | \right\} (12) \\ \leq \max  \left\{\| Q _ {1} - Q _ {2} \| _ {\infty}, \| Q _ {1} - Q _ {2} \| _ {\infty} \right\} (13) \\ = \| Q _ {1} - Q _ {2} \| _ {\infty} (14) \\ \end{array}
$$
The first line is from the definition of  $V_{Q}$ . The second line follows from the fact that  $|\max f(x) - \max g(x)| \leq \max |f(x) - g(x)|$ . The rest follows from the definition of infinity norm between  $Q$  functions. For  $b = 0$ , it is straightforward that  $|V_{Q_1}(s,b) - V_{Q_2}(s,b)| \leq \| Q_1 - Q_2\|_{\infty}$  as well. So  $\| V_{Q_1} - V_{Q_2}\|_{\infty} \leq \| Q_1 - Q_2\|_{\infty}$ . Now we have that
$$
\left\| \mathcal {T} _ {\mathrm {C B}} Q _ {1} - \mathcal {T} _ {\mathrm {C B}} Q _ {1} \right\| _ {\infty} = \| r + \gamma \mathbb {E} \left[ V _ {Q _ {1}} \right] - r - \gamma \mathbb {E} \left[ V _ {Q _ {2}} \right] \| _ {\infty} \tag {15}
$$
$$
\leq \gamma \| Q _ {1} - Q _ {2} \| _ {\infty} \tag {16}
$$
We finished the proof of  $\mathcal{T}_{\mathrm{CB}}$  is a  $\gamma$ -contraction operator. Thus there exists a unique fixed point of the operator  $\mathcal{T}_{\mathrm{CB}}$ . Now we prove the second part of the theorem that  $Q^{\star} = \mathcal{T}_{\mathrm{CB}}Q^{\star}$ .
First, we define
$$
V ^ {\star} (s, b) := \max  _ {\pi} \mathbb {E} \left[ \sum_ {t = 0} ^ {\infty} \gamma^ {t} r _ {t} \mid s _ {0} = 0, b _ {0} = b, \pi \right], \text {s . t .} b _ {t} \geq 0 \tag {17}
$$
In order to show  $Q^{\star} = \mathcal{T}_{\mathrm{CB}}Q^{\star}$ , we need to prove the following two equations:
$$
Q ^ {\star} (s, b, a) = r (s, a) + \gamma \underset {s ^ {\prime}} {\mathbb {E}} [ V ^ {\star} \left(s ^ {\prime}, b\right) ] \tag {18}
$$
$$
V ^ {\star} (s, b) = \left\{ \begin{array}{l l} \max  \left\{\max  _ {a} Q ^ {\star} (s, b - 1, a), \mathbb {E} _ {a \sim \mu} Q ^ {\star} (s, b, a) \right\} & b > 0 \\ \mathbb {E} _ {a \sim \mu} Q ^ {\star} (s, b, a) & b = 0 \end{array} \right. \tag {19}
$$
To prove Equation [18], we write  $Q^{\star}(s,b,a)$  as:
$$
\begin{array}{l} Q ^ {\star} (s, b, a) = \max  _ {\pi s. t. b _ {t} \geq 0} r (s, a) + \mathbb {E} \left[ \sum_ {t = 1} ^ {\infty} \gamma^ {t} r _ {t} \mid s _ {0} = s, a _ {0} = a, b _ {1} = b, \pi \right] (20) \\ = \max  _ {\pi s. t. b _ {t} \geq 0} r (s, a) + \gamma \mathbb {E} _ {s ^ {\prime}} \left[ \mathbb {E} \left[ \sum_ {t = 1} ^ {\infty} \gamma^ {t - 1} r _ {t} \mid s _ {1} = s ^ {\prime}, b _ {1} = b, \pi \right] \right] (21) \\ = r (s, a) + \gamma \mathbb {E} _ {s ^ {\prime}} \left[ \max  _ {\pi s. t. b _ {t} \geq 0} \mathbb {E} \left[ \sum_ {t = 1} ^ {\infty} \gamma^ {t - 1} r _ {t} \mid s _ {1} = s ^ {\prime}, b _ {1} = b, \pi \right] \right] (22) \\ = r (s, a) + \gamma \underset {s ^ {\prime}} {\mathbb {E}} \left[ V ^ {\star} \left(s ^ {\prime}, b\right) \right] (23) \\ \end{array}
$$
The first line holds by the linearity of expectation. The second line follows from the definition of conditional expectation. The third line follows from the fact that given  $s, a, s'$  does not depend on  $\pi$ . The fourth line follows from the definition of  $V^{\star}$  and the change of index  $t' = t + 1$ .
Next, we prove Equation 19. When  $b = 0$ , it is obvious that  $Q^{\star} = Q^{\mu}$  and  $V^{\star} = V^{\mu}$ . So  $V^{\star}(s,0) = \mathbb{E}_{a\sim \mu}[Q^{\star}(s,0,a)]$ . Now we consider the case of  $b > 0$ .
When  $b > 0$ , we have that
$$
\begin{array}{l} V ^ {\star} (s, b) = \max  _ {\pi s. t. b _ {t} \geq 0} \mathbb {E} \left[ \sum_ {t = 0} ^ {\infty} \gamma^ {t} r _ {t} \mid s _ {0} = s, b _ {0} = b, \pi \right] (24) \\ \leq \max  _ {\pi_ {0}} \max  _ {\pi s. t. b _ {t} \geq 0} \mathbb {E} _ {a \sim \pi (\cdot | s, b)} \left[ \mathbb {E} \left[ \sum_ {t = 0} ^ {\infty} \gamma^ {t} r _ {t} \mid s _ {0} = 0, a _ {0} = a, b _ {0} = b, \pi \right] \right] (25) \\ = \max  _ {\pi_ {0}} \mathbb {E} _ {a \sim \pi (\cdot | s, b)} \left[ \max  _ {\pi s. t. b _ {t} \geq 0} \mathbb {E} \left[ \sum_ {t = 0} ^ {\infty} \gamma^ {t} r _ {t} \mid s _ {0} = 0, a _ {0} = a, b _ {0} = b, \pi \right] \right] (26) \\ = \max  _ {\pi_ {0}} \mathbb {E} _ {a \sim \pi (\cdot | s, b)} \left[ \max  _ {\pi s. t. b _ {t} \geq 0} \mathbb {E} \left[ \sum_ {t = 0} ^ {\infty} \gamma^ {t} r _ {t} \mid s _ {0} = 0, a _ {0} = a, b _ {1} = b - I \left\{\pi_ {0} (\cdot | s, b) = \mu (\cdot | s) \right\}, \pi \right] \right] (27) \\ = \max  _ {\pi_ {0}} \mathbb {E} _ {a \sim \pi (\cdot | s, b)} Q ^ {\star} (s, b - I \left\{\pi_ {0} (\cdot | s, b) = \mu (\cdot | s) \right\}, a) (28) \\ = \max  \left\{\max  _ {a} Q ^ {\star} (s, b - 1, a), \mathbb {E} _ {a \sim \mu} Q ^ {\star} (s, b, a) \right\} (29) \\ \end{array}
$$
Next, we need to prove that
$$
\begin{array}{l} \max  _ {\pi s. t. b _ {t} \geq 0} \mathbb {E} \left[ \sum_ {t = 0} ^ {\infty} \gamma^ {t} r _ {t} \mid s _ {0} = s, b _ {0} = b, \pi \right] (30) \\ = \max  _ {\pi_ {0}} \max  _ {\pi s. t. b _ {t} \geq 0} \mathbb {E} _ {a \sim \pi (\cdot | s, b)} \left[ \mathbb {E} \left[ \sum_ {t = 0} ^ {\infty} \gamma^ {t} r _ {t} \mid s _ {0} = 0, a _ {0} = a, b _ {0} = b, \pi \right] \right] (31) \\ \end{array}
$$
$$
\pi^ {\star} = \underset {\pi s. t. b _ {t} \geq 0} {\operatorname {a r g m a x}} \mathbb {E} \left[ \sum_ {t = 0} ^ {\infty} \gamma^ {t} r _ {t} \mid s _ {0} = s, b _ {0} = b, \pi \right] \tag {32}
$$