Transition Constrained Bayesian Optimization via Markov Decision Processes

W: Generally, I believe the paper’s presentation could be enhanced...

Indeed, the reviewer is right we could start by presenting the problem as a control problem with an unknown objective, and then derive a surrogate for the unknown objective that is refined over episodes. Our main audience is practitioners of Bayesian optimization; hence, we chose the route in the paper. We believe this comment reflects preferences in exposition, not on the correctness or soundness of the work.

The actions are formally defined in Sec. 2.2. We could have introduced them already in Sec. 2.1, but it would clutter the notation to no benefit in exposition. We will make sure to improve clarity in Sec. 2.1.

W: In the background section, it is not very obvious what the paper’s contribution is versus what constitutes background information. The citations are not very transparent.

Please see the general rebuttal.

W: I believe one of the most important related concepts, Bayesian Reinforcement Learning or Dual Control, is not discussed at all.

Indeed, these are related but also crucially different. For example, Bayesian RL uses a reward function that is linear in the state-action visitation, while our utility is non-linear but convex. Dual control goes more in the direction of an unknown transition system, which we assume to be known here but indeed constitutes a very reasonable further work. We can add a small discussion of both concepts.

W: Additionally, I would have appreciated more analysis on why the specific acquisition function was chosen over others. How would the algorithm perform with a different acquisition function?

The objective was motivated by a goal to minimize the probability of a mistake. It was then refined to be computationally tractable with only one inequality in the derivation that makes it not tight (see general rebuttal). Our framework with the acquisition function as we derive it has the advantage that the induced state-action distribution of a trajectory allows for adaptive resampling and is theoretically grounded.

If we had a different goal to maximizer identification, a different acquisition function would perhaps be more appropriate and could be used within our framework provided that it depends only on the state-action visitations akin to described in Mutny et. al. (2023) in Lemma 1. We leave this to future work.

Q: Is the utility function presented in the paper a contribution? Could the authors elaborate a bit on this?

Please see the general rebuttal.

Q: I am a bit confused about the notation in line 132. Could the authors please define $\tau_i$?

$\tau_i$ is the deployed trajectory at the $i$th episode. We will clarify this in the paper.

Q: In line 136 and 199 is $S$ a switch in notation?

This was a typo, thanks for pointing it out.

Q: Could the authors elaborate a bit on line 209. What exactly is sampled using Thompson Sampling?

We use Thompson Sampling to sample the elements of the candidate set of maximizers. In other words, it is used as a finite approximation of the uncertainty quantification provided by the continuous set $\mathcal{Z}_t$. That is, we create $K$ samples of the GP and use the optimum of each sample to create a set of $K$ potential maximizers.

Q: In line 259, why is this specific constraint of 0.5 chosen?

This is an arbitrary choice related to the code implementation, as we work in the search space $[-0.5, 0.5]$ (and normalize functions and data to those bounds). In reality, any bounds could be chosen, in which case the constraint becomes $x \in [a, b]$. Thanks for pointing this out; we will clarify this in the paper.

Q: I am a bit confused about line 263. I do not think MPC is necessarily convex. There is a lot of literature on non-linear MPC. Could the authors elaborate on this?

One thing is non-linear, another is non-convex. In fact, there is a lot of literature on both. Most MPC is studied with quadratic costs (non-linear), and with linear dynamics, then it is convex in the action variables. This is what is taught in graduate MPC courses. Of course, there is also non-convex non-linear MPC (common with mixed integer, for example), but this is more exotic and needs to be heuristically solved. We fall into this category as our cost is the utility which is non-convex in the action variables, but it can be heuristically solved.

Q: Section 2 is the background but there are no citations in section 2.1. Is this a new solution proposed by the paper, or is it based on existing work?

Please see the general rebuttal.

Q: Line 575: Please explain why the proportionality holds. Monotonicity does not imply proportionality (same ratio).

This is a typo. With the first $A \propto B$ we refer to there existing a finite constant such that $A(x) = C B(x)$. Our goal is to find $\arg\min_x A(x)$, so we can equivalently minimize $\min \log B(x)$ to identify the minimizer of $A(x)$ as well, and this is what the second $\propto$ refers to, but it should be replaced with $\stackrel{!}{=}$ and explained in the text for more clarity. We apologize for the confusion. In addition, later in line 578, there is a typo; instead of $\propto$, it should read $\approx$, see the preceding line for an explanation.

Q: Line 587: Please explain the two inequalities.

The first inequality is missing a square (a typo, see 574 for the correct definition). The step is incorrectly explained and in fact holds only for large $T \gg 0$, or when $\mathbf{V}_T \gg \mathbf{I}$, where then $(\mathbf{V}_T + \mathbf{I})^{-1}\mathbf{V} \approx \mathbf{I}$.

The second inequality follows by $\mathbf{V}_T = \mathbf{X}^\top \mathbf{X}$, push-forward identity, $\mathbf{V}_T \preceq \mathbf{V}_T + \mathbf{I}$ and monotonicity of the matrix inverse under Löwener ordering. The two are then substituted back to the original objective. Notice that the objective is as such only when $T$ is large (asymptotic regime).

Q: Would the authors please clarify equation (7)? What does it mean by the expectation of a state $X_t$? What is $\pi$ in $d_{\pi}$? $Z_t$ and GP posterior may change after every time step, so I have a hard time understanding why $U$ remains the same and the expectation is only taken over $X_t$ in equation (7). Furthermore, as the utility depends on the time step (the GP posterior and $Z_t$), two visitations to the same state at different time steps have two different utilities. Then, how do we combine all visitations into $d_\pi$ (lines 143-144).

In the discussion around Eq. (7), $X_t$ refers to the set of trajectories (see L132). We analyze and derive the algorithm in the episodic setting (trajectory feedback), where we execute a trajectory in full before updating the posterior (GP and $Z_t$). A trajectory is a random set, and therefore we must define a utility with respect to the expected trajectory when using a policy $\pi$. Each policy $\pi$ induces a state-action distribution, which tells us how likely we are to visit a state in a realization of a trajectory, which we denote $d_\pi$. To our benefit, as we show in the derivations, the objective depends only on $d_\pi$.

For trajectory feedback, $Z_t$ does not change at every time-step, only after the whole trajectory is executed. When planning, we keep $Z_t$ fixed. Later we introduce instant feedback, where GP and policy are recalculated at each time step. However, we do not provide a formal derivation of optimality for this setup. The utility is still defined properly where we optimize the remaining part of the trajectory and by extension $\pi$ or $d_\pi$ or the time-steps from $h$ to $H$. This draws a parallel to the MPC literature of using receding horizon re-planning.

Q: The pseudocode in Algorithm 1 is unclear. What is $d_\pi$ (there is no initialization of it)? Why does the update of $\hat{d}\_{t,h+1}(x)$ match equation (11)? Why are $GP_{t+1}$ and $Z_{t+1}$ only updated with $X_{t,H}$ and $Y_{t,H}$ (the observations from episode $t$, not all episodes from 1 to $t$)? What are $X_{t,H}$ and $Y_{t,H}$ (the time step is only until $H-1$).

Note $d_\pi$ does not require initialization as it is the variable of the utility; it denotes the state-action distribution of a policy. In the second line of the algorithm you can see we take an argmax over this variable.

Throughout the algorithm, we want to keep track of the visited states. We do this via the empirical state-action distribution, defined by the states we have visited so far. After visiting new state(s), we want to add them to the empirical state-action distribution which we can do by following the normalization procedure defined in eq. (11). For clarity, the empirical state-action distribution $\hat{d}\_{t,h}$ denotes past states, while the variable $d_\pi$ is representing future states we must choose.

The GP and $Z$ are updated using observations from all previous trajectories. In the paper, we use lower-case variables to denote individual observations, and upper-case variables to denote all trajectories up to the time-step specified by the index.

Q: The policy depends on the time step, so $\pi_{t,0}$ is only deployed at the first time step. How is the state visitation frequency of such a policy estimated well if it is only deployed at the 1st time step of each episode?

The state-action frequency of a given policy can be calculated exactly; there is no need to deploy the policy itself. We assume the transition dynamics are known.

Q: Would the authors please clarify equation (12) using the correct subscript of $d (d_{t,h})$?

Notice that $\hat{d}_{t,h}$ refers to the empirical visitation. In other words, the frequency of visited states in trajectory $t$ at time horizon $h$. This visitation changes as we execute the policies for a longer time and more trajectories—this is updated depending on the trajectory we have followed up to time-step $t, h$. The empirical visitation $\hat{d}$ and also the variable $d$ in optimization Eq. (12) is $d:\mathcal{S} \times \mathcal{A} \times H \rightarrow [0,1]$.

Q: Some minor remarks: In equation (11), should $j$ be from 1 to $t-1$? Line 122: $y(xh)$ should be $y(xh,ah)$.

Yes, thanks for pointing out the mistakes.

We would like to thank the reviewer for their prompt response.

The corrected proof resolves my concern

We are happy to hear that. Please consider raising your score, especially if you find the following clarifications useful.

However, the explanation regarding...

The reviewer's intuition is correct, let us explain in detail what this "constant" is and how we use it (in essence the constant is the total length of the trajectories):

Firstly, we would like to clarify that empirical state visitations are also normalized. This can be seen, for example, in Eq.(11) where we introduced them.

Let us also note that the posterior variance of a Gaussian process does not depend depend on observed values, but only the points $x$, hence we can forward plan reduction in uncertainty. This is why Eq. (4) is equal to Eq. (6). Now, we show that the relaxation of the objective to state distributions holds in the simplest setting, for notation we will use $X$ to refer to the trajectory and $S$ to refer to the state space (in the paper we use $\mathcal{X}$):

For simplicity assume we have a single trajectory of length $H$ and no observation noise:

$

V(X) = \sum_{x \in {X}} \Phi(X) \Phi(X)^T = H \sum_{x \in S} \hat{d_X} (x) \Phi(X) \Phi(X)^T:= H V(\hat{d_X})

$

Now we note that:

$

||\Phi(z) - \Phi(z')|| _{V(X)^{-1}} = (\Phi(z) - \Phi(z'))^T V(X)^{-1} (\Phi(z) - \Phi(z'))

$

$

= \Phi(z)^T V(X)^{-1} \Phi(z) + \Phi(z')^T V(X)^{-1} \Phi(z') - 2 \Phi(z)^T V(X)^{-1} \Phi(z')

$

$

= \text{Var}[f(z) | X] + \text{Var}[f(z') | X] - 2 \text{Cov}[f(z), f(z') | X] \

$

$

= \text{Var}[f(z) - f(z') | X]

$

Note that since $V(X) \propto V(\hat{d}\_X)$, then:

$

\text{Var}[f(z) - f(z') | X] \propto ||\Phi(z) - \Phi(z')||\_{V(\hat{d}\_X)^{-1}}

$

so the optimization of the objectives is equivalent, note that the proportionality constant is exactly the normalization constant. In Lemma D.1 we formalize this more carefully in the setting with many trajectories and observation noise. We have to change the regularization of the objective appropriately (e.g. we must scale the contribution of the prior by $1 / H$, when planning for one episode, to counteract the normalization of the state-distributions), i.e., for the objectives to be equivalent we must use:

$

V(d_\pi) = \left(\sum_{x \in S} d_\pi(x) \frac{\Phi(x)^T \Phi(x)}{\sigma^2} + \frac{1}{H} I \right) = \frac{1}{\sigma^2}\left(\sum_{x \in S} d_\pi(x) \Phi(x)^T \Phi(x)+ \frac{\sigma ^2}{H} I \right).

$

For future reference, there are still two minor typos in the paper that we found, Eq. (8) it should read:

$

V(d_\pi) = \left(\sum_{x \in S} d_\pi(x,a) \frac{\Phi(x)^T \Phi(x)}{\sigma^2(x,a)} + \frac{1}{TH} I \right).

$

where we plan for $T$ episodes (trajectories) to make it consistent with the correct result in Lemma D.1. And in the proof of Lemma D.1, in the last equation of L705 there should not be a $TH$ before the sum. We wanted to manipulate $TH$ out but forgot it inside the inverse -- an honest typo. We will fix the typos and add a clarification that we need the correct regularization for the objectives to be equivalent.

To be fully clear, our implementation of the objective in the code does include the correct regularization procedure (see file mdpexplore/functionals/bandit_functionals.py L242 where we multiply the identity by $1 - \alpha$).

Thank you the response and for carefully analyzing our derivation. We want to again mention that the objective used in our paper has been used in works of Fiesz et. al. (2019) published at NeurIPS. We hope there is no doubt that this is a meaningful objective. We only provide a Bayesian motivation for it.

Upon a more meticulous look of our derivation, we found we had an unnecessarily complicated path that led to the sign blunder and wrong fraction order that the reviewer points out (which ended up cancelling each other out). Allow us instead to provide a more elegant and simple derivation (which will be added to the revised paper). It follows from the Gaussian tail bound that states for any $X \sim N(\mu, \sigma^2)$ and $t > 0$ then: $\mathbb{P}(X \geq \mu + t) \leq e^{-\frac{t^2}{2 \sigma^2}}$ We can now bound the probability of making an error, indeed $\mu_T(z) - \mu_T(z^*) \sim N(a_z, b_z^2)$: $\mathbb{P}(\mu_T(z) - \mu_T(z^*) \geq 0) = \mathbb{P}(\mu_T(z) - \mu_T(z^*) \geq a_z + (-a_z) ) \leq e^{-\frac{a^2}{2 b^2}}$ with the caveat that the bound only holds asymptotically as $T \rightarrow \infty$, as in this regime $a_z \approx f(z) - f(x_f^*)$, due to asymptotic argument we made previously. Notice that by definition of $x_f^*$, $f(z) - f(x_f^*)\leq 0$, hence substituting for $t = -a_z$,

$

E_{f\sim GP}[\log P(\mu_T(z) - \mu_T(x_f^\star) \geq 0 | f)] \leq - \frac{1}{2} E_{f\sim GP}\left[\frac{a_z^2}{b_z^2}\right]

$

by taking first the log, and second the expectation on both sides. From here the arguments follow in a similar fashion to the original paper, with the caveat that originally we had an error sign and the numerator and denominator the wrong way around and so the errors cancelled out. For full clarity here is the argument:

Our aim is now to minimize making an error, so minimizing $-E_{f \sim GP}\left[\frac{a_z^2}{b_z^2}\right]$. It then follows that in the large $T$ limit (from the previously discussed bounds on $b_z$ and $a_z \approx f(z) - f(x_f^\star)$):

$

- E_{f\sim GP}\left[\frac{a_z^2}{b_z^2}\right] \leq - E_{f\sim GP} \left[ \frac{(f(z) - f(x_f^\star))^2}{k_{\mathbf{X}}(z, x_f^\star)} \right]

$

combining the two equations we obtain:

$

E_{f\sim GP}[\log P(\mu_T(z) - \mu_T(x_f^\star) \geq 0 | f)] \leq - \frac{1}{2} E_{f\sim GP} \left[ \frac{(f(z) - f(x_f^\star))^2}{k_{\mathbf{X}}(z, x_f^\star)} \right]

$

note the correct interpretation of the bound: the probability of an error is small if (i) the uncertainty is small, or (ii) the difference in values is large. Then we simply follow the arguments in Section C.3 to arrive to the objective:

$

- E_{f\sim GP} \left[ \frac{(f(z) - f(x_f^\star))^2}{k_{\mathbf{X}}(z, x_f^\star)} \right] \leq - \frac{E_{f\sim GP}[\text{gap(f)}]}{\text{Var}[f(z) - f(z') | \mathbf{X}]}

$

And finally, we note that due to the negative sign, and the denominator being constant, minimizing the RHS is equivalent to minimizing the objective introduced in Eq.(4), i.e. $\arg\min_{x}-g(x)$ is equivalent to $\arg\min_{x} \frac{1}{g(x)}$ when $g(x) > 0$.

W: The main weakness of the work would appear to be the lack of compelling evidence for why the authors' method should be chosen over competitor methods such as LSR or SnAKe on practical problems...

Please see the general rebuttal.

W: The code would benefit from a README describing how to reproduce the results of the paper.

We agree and we will add this after peer review. We did have a README but hastily deleted it to preserve anonymity!

W: On line 567 epsilon is homoscedastic but in the main paper the authors state that epsilon can be heteroscedastic?

For simplicity, the derivation is done for the homoskedastic setting, however the same derivation holds in the heteroskedastic setting by simply changing $\mathbf{V} = \sum_{i=1}^T \frac{1}{\sigma_i^2}\Phi(x_i)\Phi(x_i)^\top$ instead of for heteroskedastic least squares $\mathbf{V} = \sum_{i=1}^T \frac{1}{\sigma^2}\Phi(x_i)\Phi(x_i)^\top =: \mathbf{X} \mathbf{\Sigma}\mathbf{X}$. We can add this comment to the discussion.

W: Other minor points

Thanks for pointing out minor issues which we will address.

Q: In Equation 3, do the authors have any intuition...

It allows an interpretation as noise to signal ratio. In the numerator, we have uncertainty (covariance) between $f(z)$ and $f(x_f^*)$, while in the denominator we have the level of signal that is the gap between $f(z)$ and $f(x_f^*)$. The smaller the gap, the harder it is to know which of the two is maximal, increasing the probability of error. On the other hand, if the two states $x_f^*$ and $z$ are similar to each other, we do not need to put additional effort as they are very correlated, hence balancing these two is the right metric. We will add this intuition to the paper to improve clarity.

Q: In Equation 4, the authors state their objective as a maximization problem. However, in the text, the authors refer to minimizing the uncertainty among all pairs in $\mathcal{Z}$. Could the authors explain this?

The ultimate goal is to maximize an unknown function $f$. For exposition, as is common in BayesOpt, we talked about utility or acquisition function maximization, but then deriving the utility we find we need to minimize the uncertainty among all pairs. This lends itself to explain as minimization better. We could equally talk about maximization of utility $=$ 1/loss.

Q: In Section D.1 of the appendix, the authors give the formulation of the objective for general kernel matrices. The authors state that the objective may be computed by inverting a $|\mathcal{X}| \times |\mathcal{X}|$ matrix. How restrictive is this for the scale of problems that can be addressed with this formulation of the objective?

Indeed, this is a limitation seen also with some classical BO approaches such as Thompson sampling which need inversion of $|\mathcal{X}| \times |\mathcal{X}|$ matrix -- same complexity as our method.

Common remedies are some finite dimensional approximations of the kernels, leading to $m \times m$ inversions or other reformulations improving the efficiency of matrix inversion that we cite in our work. Akin to Thompson sampling, even combinations of different techniques such as feature approximation and different sampling, namely Matheron rule for sampling, could be a very practical remedy (Wilson et al. 2020). This issue is a fundamental difficulty associated with reasoning (planning) in non-parametric function spaces, not necessarily something specific to our method. Also notice that in continuous domains the reparametrization of $d$ eliminates this inversion but leads to potentially non-convex program in the reparametrization.

Q: In Section 5, what was the motivation for the decision to fix the GP hyperparameters?

Please see the general rebuttal.

Q: For Figures 2 and 3 why are there no error bars for the average regret curves?

We did not include them to avoid cluttering the graphics. The percentage of runs we identify the maximizer allows us to see the robustness of each method. Nonetheless, we are happy to include complete regret plots in the appendix.

L: The main limitation as described above would appear to be the lack of compelling evidence for problem settings under which the authors' algorithm outperforms other transition-constrained Bayesian optimization methods such as LSR and SnAKe.

Please see the general rebuttal.

Reviewer J8j4, could you please review the authors' response and see whether it addresses your open questions? Please acknowledge having done so in a comment together with a discussion of whether / how it has affected your assessment of the submission. Thanks.

W: The objective is non-convex and differs from conventional optimization in RL. This demands heuristics, which are not sufficiently discussed.

This is true for the continuous case and common to most setups in control theory. However, special to our setting, for all discrete environments in the experiments we are able to solve the sub-problems provably with convex solvers, see Frank-Wolfe updates.

W: In the experiments, the GP hyperparameters are fixed...

Please see the general rebuttal.

W: The results in Figure 4 appear to lack statistical significance. Additionally, there is no clear preference among the variants of MDP-BO, resonating with the challenge in conventional BO.

Please see the general rebuttal.

W: There is a series of works on applying BO to graphical structures...

Causal BO uses graphs that relate different variables in the optimization process according to a Bayesian causal model. Our graphs relate transitions between states. This is different in spirit, but we will add to the discussion.

Q: Could the author comment on the practice of limiting the extension of MDP-BO to either TS or UCB?

We obtained good results with both, so we did not try to use any other batch algorithms for estimating the maximizer candidate set. However, there is flexibility in the algorithm to use different estimation methods and could be the focus of future work.

Q: One minor question: in figure 3(a), is the legend "MDP-BO" written as "MDP-B0"?

Correct, this is an error, thanks for catching it!

Q: The objective defined in Section 2.1 and the general BO framework presented in Algorithm 1, beyond the MDP-based planning solver, are commonly seen in previous works. Li and Scarlett (2022) studied the batched version of this formulation, Zhang et al. (2023) studied the high-dimensional treatment, and Han et al. (2024) studied the formulation in game theory, among others. There seems to be potential to generalize the proposed MDP-BO framework into broader applications. It would be a bonus point to additionally discuss the corresponding potentials of MDP-BO.

Thanks for all the references, they are indeed relevant. We will include a discussion of them in the paper.