Bounding-Box Inference for Error-Aware Model-Based Reinforcement Learning

Talvitie, Erin J., et al. “Bounding-Box Inference for Error-Aware Model-Based Reinforcement Learning.” Reinforcement Learning Conference, 2024.

Info

I replicated the results presented in this paper. The codebase can be found in this GitHub repository (in python).

The original codebase for the paper can be found in this GitHub repository (in C++).

warning

This content is a summary and includes my personal takeaways from the paper "Bounding-Box Inference for Error-Aware Model-Based Reinforcement Learning" by Talvitie et al. It reflects the points I found most relevant, along with my own conclusions on the topics discussed.

For a more comprehensive understanding of the topics and to align with the authors' perspectives, please READ THE PAPER.

Introduction

Model-Based Reinforcement Learning (MBRL) aims to learn a model of the environment dynamics to reduce sample complexity and potentially improve policy performance. Let $\mathcal{M} = (\mathcal{S}, \mathcal{A}, p, r, \gamma)$ be a Markov Decision Process (MDP), where $\mathcal{S}$ is the state space, $\mathcal{A}$ is the action space, $p(s'|s,a)$ is the transition distribution, $r(s,a)$ is the reward function, and $\gamma \in [0,1]$ is the discount factor. The goal is to find a policy $\pi(a|s)$ that maximizes the expected discounted return $G_t = \sum_{i=0}^{\infty} \gamma^i R_{t+i+1}$.

In MBRL, the agent simultaneously learns a model $(\hat{p}, \hat{r})$ that approximates $(p, r)$, and uses it to simulate state-action transitions. Naively treating simulated rollouts as ground truth can mislead the value function and degrade the policy. This problem is particularly acute when there is insufficient model capacity, inadequate exploration of the hypothesis space, or limited training samples. The paper defines:

• Aleatoric uncertainty as tied to the nature of the stochastic process, specifically the randomness or noise in the environment. This type of uncertainty is irreducible, meaning that no amount of additional data or model refinement can eliminate it, as it is intrinsic to the system being modeled.

• Epistemic uncertainty as tied to the lack of knowledge or uncertainty about the environment or model parameters. Unlike aleatoric uncertainty, it is reducible, meaning it can be decreased by gathering and training the model on additional data.

• Model inadequacy refering to a model's limitations in accurately capturing the true dynamics or complexity of the environment. This uncertainty occurs when the model is structurally flawed, overly simplistic, or based on incorrect assumptions, leading to systematic prediction errors. Unlike epistemic uncertainty, model inadequacy cannot be fully resolved by acquiring more data; it requires fundamentally revising or improving the model's structure to better reflect the underlying system.

Selective planning aims to mitigate this by weighing model-based value targets according to their estimated reliability. Let $\hat{q}(s,a)$ be an estimate of the state-action value function. Model-based Value Expansion (MVE) computes multi-step targets:

$$\hat{\rho}_1 = r_{t+1} + \gamma \max_{a'}\hat{q}(S_{t+1}, a'),$$

and for $h>1$,

$$\hat{\rho}_h = r_{t+1} + \sum_{i=2}^{h}\gamma^{i-1}\hat{r}_{t+i} + \gamma^h \max_{a'} \hat{q}(\hat{S}_{t+h}, a'),$$

where $(\hat{s}_{t+i}, \hat{r}_{t+i})$ are generated by the model $(\hat{p}, \hat{r})$. By weighting these targets based on uncertainty, selective planning can exploit accurate predictions while ignoring harmful ones.

The paper present a novel uncertainty measure for MVE, Bounding-Box Inference (BBI). BBI infers upper and lower bounds on the multi-step targets, and weights down the targets according to the gap between the bounds.

Problem Setting and Mathematical Formulation

We consider an MDP with unknown dynamics $p$ and reward $r$. At each time $t$, the agent observes a state $S_t \in \mathcal{S}$, selects an action $A_t \in \mathcal{A}$, and receives a reward $R_{t+1}=r(S_t,A_t)$. The environment transitions to a new state $S_{t+1} \sim p(\cdot|S_t,A_t)$.

The agent maintains a state-action value function $\hat{q}(s,a)$ aiming to approximate $q_{*}(s,a) = \max_\pi q_{\pi}(s,a)$, where $q_{\pi}(s,a) = \mathbb{E}[\sum_{i=0}^{\infty}\gamma^i R_{t+i+1}|S_t=s,A_t=a,\pi]$. The agent also learns a model $(\hat{p}, \hat{r})$ approximating $p$ and $r$.

Let $w_h \geq 0$ be the weight assigned to the $h$-step target. The MVE update:

$$\hat{q}(S_t,A_t) \leftarrow \hat{q}(S_t,A_t) + \alpha\left(\sum_{h=1}^H\frac{w_h \hat{\rho}_h}{\sum_{h'=1}^H w_{h'}}- \hat{q}(S_t,A_t)\right).$$

The challenge is to define the uncertainty over the simulated targets, $\hat{\rho}_h$.

Uncertainty Measures

Several uncertainty measures are compared in the paper:

1. One-Step Predicted Variance (1SPV):
   One-Step Predicted Variance estimates uncertainty by aggregating the predicted variances of one-step model predictions over the course of an $h$-step rollout. It assumes that the total uncertainty at each horizon $h$ can be represented as a linear sum of independent variances from individual state and reward dimensions. Specifically, the uncertainty at horizon $h$ is given by:

   $$u_h = \sum_{j=0}^{h-1} \left( \sum_d \hat{\sigma}_d^2(s_{t+j}, a_{t+j}) + \hat{\sigma}_r^2(s_{t+j}, a_{t+j}) \right)^2,$$

   where $\hat{\sigma}_d^2(s, a)$ represents the predicted variance of the next state dimension $d$, and $\hat{\sigma}_r^2(s, a)$ represents the predicted variance of the reward at state $s$ and action $a$.

   While simple and computationally efficient, this method has key limitations. The assumption that state uncertainties directly approximate TD target uncertainties can lead to overly conservative estimates, especially when state dimensions with high variance are irrelevant to the value function.

2. Monte Carlo Target Variance (MCTV):

   Monte Carlo Target Variance takes a direct approach to estimating uncertainty by sampling multiple rollouts from the model for the same initial transition. Given $k$ sampled TD targets $\{\hat{\rho}^j_h\}_{j=1}^k$ at horizon $h$, the mean target is defined as:

   $$\hat{\rho}_h = \frac{1}{k} \sum_{j=1}^k \hat{\rho}^j_h,$$

   and the uncertainty is estimated as the sample variance:

   $$u_h = \frac{1}{k-1} \sum_{j=1}^k (\hat{\rho}^j_h - \hat{\rho}_h)^2.$$

   This method provides a principled estimate of uncertainty directly over the TD targets. However, it is sensitive to the underlying model’s predicted probability distribution. If the model misrepresents the dynamics or assigns inappropriate probabilities to certain transitions, the sample variance may fail to reflect true model inaccuracies. Moreover, the accuracy of this method depends on the number of Monte Carlo samples ($k$), which introduces a trade-off between computational cost and the reliability of the estimate.

3. Monte Carlo Target Range (MCTR):

   Monte Carlo Target Range refines the idea of using samples from the model but replaces variance with range as the measure of uncertainty. Given the same $k$ sampled TD targets $\{\hat{\rho}^j_h\}_{j=1}^k$ at horizon $h$, the uncertainty is defined as:

   $$u_h = \max_j \hat{\rho}^j_h - \min_j \hat{\rho}^j_h.$$

   This range-based approach directly captures the spread of possible outcomes without relying on assumptions about the distribution’s shape. However, like MCTV, MCTR still depends on the model’s predicted distribution. If the model assigns negligible probabilities to critical transitions, the range may underestimate true uncertainty. Additionally, the estimate can be sensitive to the number of samples ($k$) used, as a small $k$ may fail to capture the full range of possible TD targets.

4. Bounding-Box Inference (BBI):

   A novel uncertainty measure that relies on bounding-box approximations of possible outcomes. Consider a bounding-box over states $\overline{\underline{s}} = [\underline{s}, \overline{s}]$ and over actions $\overline{\underline{a}}$. Define:

   $$\underline{q}(\overline{\underline{s}}_t,\overline{\underline{a}}_t) = \inf_{s\in\overline{\underline{s}}_t,a\in\overline{\underline{a}}_t}\hat{q}(s,a), \quad \overline{q}(\overline{\underline{s}}_t,\overline{\underline{a}}_t)=\sup_{s\in\overline{\underline{s}}_t,a\in\overline{\underline{a}}_t}\hat{q}(s,a).$$

   Similarly, given the current state and action $(S_t,A_t)$ and the actual next transition $(S_{t+1}, R_{t+1})$, we define a bounding-box model rollout:

   $$\overline{\underline{r}}_{t+i}, \overline{\underline{s}}_{t+i}, \overline{\underline{a}}_{t+i}, \quad \text{for } i=2,\dots,h.$$

   The bounding-box TD target at horizon $h$:

   $$\underline{\rho}_h = R_{t+1} + \sum_{i=2}^h \gamma^{i-1}\underline{r}_{t+i} + \gamma^h \underline{v}(\overline{\underline{s}}_{t+h}),$$ $$\overline{\rho}_h = R_{t+1} + \sum_{i=2}^h \gamma^{i-1}\overline{r}_{t+i} + \gamma^h \overline{v}(\overline{\underline{s}}_{t+h}),$$

   where $\underline{v}(\overline{\underline{s}}) = \max_a \underline{q}(\overline{\underline{s}},a)$ and similarly for $\overline{v}(\overline{\underline{s}})$.

   The BBI uncertainty is $u_h = \overline{\rho}_h - \underline{\rho}_h$. This approach is distribution-insensitive and relies on conservative bounding, ensuring that no harmful scenario is overlooked. Though potentially conservative, BBI avoids catastrophic failures due to model misestimation.

Selective MVE Update

Let $u_h$ be the uncertainty measure associated with horizon $h$. We form weights via a softmin:

$$w_h = \frac{\exp(-u_h/\tau)}{\sum_{j=1}^H \exp(-u_j/\tau)},$$

where $\tau>0$ is a temperature parameter. As $\tau \rightarrow 0$, the update approaches Q-learning; as $\tau \rightarrow \infty$, it approaches uniform weighting.

Empirical Evaluation

WIP

GoRight Environment

Experimental Setup

Baselines and Models

• Q-learning:
• Perfect Model:
• Expectation Model:
• Sampling Model:
• Linear Model:
• Regression Tree:
• Neural Network:

Results

Discussion and Conclusions