# Wall Street Tree Search: Risk-Aware Planning for Offline Reinforcement
Learning

Dan Elbaz, 1, 2 Gal Novik, 1 Oren Salzman 2

###### Abstract

Offline reinforcement-learning (RL) algorithms learn to make decisions using a
given, fixed training dataset without online data collection. This problem
setting is captivating because it holds the promise of utilizing previously
collected datasets without any costly or risky interaction with the
environment. However, this promise also bears the drawback of this setting as
the restricted dataset induces uncertainty because the agent can encounter
unfamiliar sequences of states and actions that the training data did not
cover. To mitigate the destructive uncertainty effects, we need to balance the
aspiration to take reward-maximizing actions with the incurred risk due to
incorrect ones. In financial economics, modern portfolio theory (MPT) is a
method that risk-averse investors can use to construct diversified portfolios
that maximize their returns without unacceptable levels of risk. We propose
integrating MPT into the agent’s decision-making process, presenting a new
simple-yet-highly-effective risk-aware planning algorithm for offline RL. Our
algorithm allows us to systematically account for the _estimated quality_ of
specific actions and their _estimated risk_ due to the uncertainty. We show
that our approach can be coupled with the Transformer architecture to yield a
state-of-the-art planner, which maximizes the return for offline RL tasks.
Moreover, our algorithm reduces the variance of the results significantly
compared to conventional Transformer decoding, which results in a much more
stable algorithm—a property that is essential for the offline RL setting,
where real-world exploration and failures can be costly or dangerous.

##  1 Introduction

Reinforcement Learning (RL) is concerned with an agent learning how to take
actions in an environment to maximize the total reward it obtains. The RL
agent typically learns by trial and error, involving online interaction with
the environment to collect experiences (Sutton and Barto 1998). But learning
in the real world may be undesirable, as online data acquisition is often
costly, time-consuming, or even dangerous. Offline RL aims to bridge the gap
between RL algorithms and real-world systems by leveraging an existing dataset
or batch to learn how to make decisions in an offline stage without any online
interactions with the environment (Levine et al. 2020).

However, learning solely from offline data is a double-edged sword. On the one
hand, it enables applications in domains where online exploration is avoided,
e.g., in healthcare (Gottesman et al. 2018; Wang et al. 2018; Yu, Ren, and Liu
2019), autonomous driving (Sallab et al. 2017), and recommendation systems
(Strehl et al. 2010; Covington, Adams, and Sargin 2016). On the other hand, it
poses a major algorithmic challenge as we are faced with high levels of
uncertainty. As in offline RL, we typically want the learned policy to perform
better than the policy that collected the data. Consequently, we must execute
a different sequence of actions from the sequences stored in the batch. As a
result, the agent encounters unfamiliar state-action sequences, which induces
_subjective uncertainty_ (uncertainty that comes from ignorance due to the
limited size of the training batch). This uncertainty can lead to erroneous
value estimations, a phenomenon known as _distributional shift_ (Kumar et al.
2020), which is one of the central challenges of offline RL. Moreover,
_objective uncertainty_ , which is due to inherent system stochasticity,
further increases uncertainty and aggravates the offline RL problem,
preventing the agent from learning an optimal policy.

One promising way to solve the offline RL problem is using _model-based RL_
(MBRL). MBRL divides the RL problem into two stages: The first stage is
learning an approximate environment model (referred to as the transition or
dynamics model) with the data. In the second stage, this model is used for
decision-making (i.e., planning or policy search), usually via Model-
Predictive Control (MPC) (Richalet et al. 1978). An advantage of using MBRL is
that we can benefit from a convenient and powerful supervised learning
workhorse in the model-learning stage, which allows us to generalize to states
outside the support of the batch. However, due to the distribution shift, the
model becomes increasingly inaccurate as we get further from the points in the
batch. As a result, a planner that tries to obtain the highest possible
expected reward under the model without any precautions against model
inaccuracy can result in “model-exploitation” (Levine et al. 2020)—a situation
where the planner prefers predictions with higher returns than would be
obtained from the actual environment, resulting in poor performance.

Recent offline RL approaches suggest drawing from the tools of large-scale
language modeling, using a Transformer architecture (Vaswani et al. 2017) to
model distributions over trajectories (Janner, Li, and Levine 2021; Chen et
al. 2021). However, large-scale language models are trained on millions of
high-quality web pages; hence their decoding methods are not designed to be
immune to the effects of objective and subjective uncertainty prevalent in
offline RL. In this work, we aim to address the question of how we can modify
these decoding algorithms to tackle the uncertainty present in offline
problems.

In portfolio optimization theory, investors demand a reward for bearing risk
when making investment decisions. The economist Harry Markowitz introduced
this risk-expected return relationship in the mean-variance model in a 1952
essay (Markowitz 1952), for which he was later awarded a Nobel Memorial Prize
in Economic Sciences. The mean-variance model weighs the risk against the
reward and solves the portfolio-selection problem (formally defined in Sec.
4.3), i.e., decides how much wealth to invest in each asset by considering
only the expected return and risk, expressed as the variance.

Our main contribution is a new risk-aware planning algorithm for offline RL
that mitigates the distributional shift problem, which is inspired by modern
portfolio theory (MPT). Our planning algorithm is s best-first search
algorithm that treats the value associated with each candidate trajectory as
an asset with specified mean and variance, which acts as a proxy for the risk
indicating a distribution shift. When building the search tree, we treat the
limited node-expansion budget as wealth to invest in many different assets
(each asset corresponds to a candidate trajectory) and solve a portfolio-
optimization problem to determine how much wealth to invest in each of these
assets. We refer to our planner, which we formally introduce in Sec. 5, as
Wall Street Tree Search (WSTS) as an homage to the famous Monte Carlo Tree
Search (MCTS) (Kocsis and Szepesvári 2006) algorithm and evaluate it in
continuous control tasks from the widely studied D4RL offline RL benchmark (Fu
et al. 2020). We show, in Sec. 6, that WSTS matches or exceeds the performance
of state-of-the-art (SOTA) offline RL algorithms. At the same time, it is
substantially more reliable than the conventional decoding method resulting in
significantly reduced variance when used on the same environment model.

##  2 Related Work

Algorithms for offline RL typically fall under two categories: model-free and
model-based algorithms. Generally, each category addresses the uncertainty
(mainly due to the distributional shift) differently. In _model-free_
algorithms, the agent learns a policy or value function directly from the
dataset. Such algorithms typically address the distribution shift by
constraining the learned policy to avoid out-of-distribution behavior that the
dataset does not support (Fujimoto, Meger, and Precup 2019; Kumar et al. 2019;
Siegel et al. 2020; Wu, Tucker, and Nachum 2019; Peng et al. 2019) or use
uncertainty quantification techniques, such as ensembles, to stabilize
Q-functions (Agarwal, Schuurmans, and Norouzi 2020; Fujimoto, Meger, and
Precup 2019; Jaques et al. 2019; Wu, Tucker, and Nachum 2019).

The second category,model-based RL methods, is the approach we take in this
work, which is less explored for offline RL. Nonetheless, prior works have
demonstrated promising results of MBRL methods, particularly in offline RL
(Kidambi et al. 2020; Yu et al. 2020; Janner, Li, and Levine 2021; Chen et al.
2021). However, MBRL methods are subjected to the so-called model-exploitation
problem, which causes the decision maker (the planner) to choose erroneous
predictions with higher returns than would be obtained from the actual
environment. Argenson and Dulac-Arnold (2020) addressed the model-exploitation
problem by using ensembles and averaging the reward over all ensemble members.
Kidambi et al. (2020) addressed this problem by incorporating pessimism into
learned dynamics models. Yu et al. (2020) also studied the effects of
uncertainty in a model-based approach to offline RL and suggested optimizing a
policy using an uncertainty-penalized reward. In contrast, we propose an
explicit planning algorithm to account for the uncertainty.

Our approach to offline RL builds on the recent works by Janner, Li, and
Levine (2021); Chen et al. (2021) that frame RL as a sequence-modeling problem
and exploit the toolbox of contemporary sequence modeling. At the heart of
their approach is a sequence model based on the GPT-3 Transformer decoder
architecture (Brown et al. 2020), which they use for learning a distribution
over trajectories by jointly modeling the states and actions to provide a bias
toward generating in-distribution actions. To plan using the Transformer model
(decode its outputs), Janner, Li, and Levine coupled it with a minimally
modified version of beam search (BS) (Lowerre and Reddy 1976), and replaced
the log probabilities of transitions with the expectation of the predicted
reward signal to create a search strategy for reward-maximizing behavior.
However, decoding methods accounting only for the trajectories with the
maximum predicted rewards and ignoring the uncertainty may suffice for large-
scale language models, but might not be the best option in an offline RL
setting. We incorporate uncertainty into sequential decision-making to address
this issue and propose a new decoding algorithm.

##  3 Problem Definition

A _Markov decision process_ (MDP) is a tuple
$\mathcal{M}=\\{S,A,r,P,\rho_{0},\gamma\\}$ where $S$ and $A$ are the sets of
states and actions, $r:S\times A\rightarrow\mathbb{R}$ is the reward function,
$P$ is the system’s dynamics, which is a conditional probability distribution
of the form $P(s_{t+1}|s_{t},a_{t})$, representing the probability over the
next state given the current state and the applied action, $\rho_{0}$ defines
the initial state distribution and $\gamma\in(0,1]$ is a scalar discount
factor that penalizes future rewards.

A _trajectory_
$\upsilon=(s_{1},a_{1},r_{1},s_{2},a_{2},r_{2},...,s_{H},a_{H},r_{H})$ is a
sequence of $H$ states, actions and rewards. Here both the reward and the next
state of a given state-action pair adhere to some underlying MDP. A _policy_
$\pi:S\rightarrow A$ is a mapping between states and actions.

In the _offline RL_ problem, we are given a static, previously-collected
dataset $\mathcal{D}$ of trajectories collected from some (potentially
unknown) behavior policy $\pi_{B}$ (for example, $\pi_{B}$ can be human
demonstrations, random robot explorations, or both). The goal is to use
$\mathcal{D}$ to learn, in an offline phase, a policy $\pi$ for the
underlying, unknown MDP $\mathcal{M}$, to be executed in an online phase.

##  4 Algorithmic Background

We take a _model-based_ offline RL approach in this work. Such an approach
divides the problem into two stages S1: learning an MDP model
$\mathcal{M}^{\prime}$ that approximates $\mathcal{M}$ using the dataset
$\mathcal{D}$ and S2: using the learned model $\mathcal{M}^{\prime}$ to
extract the policy $\pi$. We will refer to the first and second stages as the
_offline model learning_ and the _online decoding_ stages, respectively, and
to the entire approach as _model-based RL_. We note that we use the same
Transformer architecture proposed by Janner, Li, and Levine (2021) to learn an
MDP model (S1), but change the decoding procedure (S2). This overall approach
is visualized in Fig. 1. As stated, our work focuses on the decoding procedure
(S2), the right block in Fig. 1.

Figure 1:  _Data collection:_ The behavior policy $\pi_{B}$ collects the
training batch $\mathcal{D}$. _Model Learning (S1, offline):_ Given the
previously collected batch, $\mathcal{D}$, we learn a sequence model.
_Decoding (S2, online):_ Decides which partial trajectories to simulate at
each time point based on some scoring function outputs: $w_{1},...,w_{N}$ (we
formally define score in 4.2)

###  4.1 Offline Model Learning (S1)

As previously mentioned, our approach to offline RL builds heavily on recent
model-based offline RL methods (Janner, Li, and Levine 2021; Chen et al. 2021)
using Transformers to learn an MDP model (S1). We describe this approach in
this subsection.

#### Dataset

To create the dataset used to train the Transformer model, each transition in
the original trajectory $\upsilon$ (Sec. 3) is augmented with a discounted
reward-to-go estimate (namely,
$R_{t}=\sum_{t^{\prime}=t}^{H}\gamma^{t^{\prime}-t}r_{t}^{\prime}$) to obtain
an _augmented trajectory_ consisting of a sequence of $H$ states, actions,
rewards, and reward-to-go. Here, it is important to note that the estimated
reward-to-go is computed using the training data and estimates the return
obtained by following the behavior policy $\pi_{B}$. In general, it does not
necessarily approximate the values of the learned policy. However, since we
use it as a heuristic estimation, it suffices for guiding the search in the
decoding stage (Janner, Li, and Levine 2021).

Using a discrete-token architecture (a Transformer) forces tokenization of the
augmented trajectory. This is done by discretizing each dimension
independently in the event of continuous inputs. Assuming $N$-dimensional
states, $M$-dimensional actions, and scalar reward and reward-to-go, the
_tokenized trajectory_ $\mathbf{y}$ is a sequence of length $T=H\cdot(N+M+2)$,
where every dimension of the trajectory is a token (subscripts on all tokens
denote timestep, and superscripts on states and actions denote dimension.):

$\mathbf{y}=(\ldots,s_{t}^{1},s_{t}^{2},\ldots,s_{t}^{N},a_{t}^{1},a_{t}^{2},\ldots,a_{t}^{M},r_{t},R_{t},s_{t+1}^{1}\ldots).$

#### Notation

The elements of the vector $\mathbf{y}$ (a tokenized trajectory) are denoted
using an ordered set of components, $y_{t}$,
$\mathbf{y}=(y_{1},y_{2},\ldots,y_{T})$. Each $y_{t}$ is an element of
$\mathcal{V}$, the set of output tokens. $\mathbf{y}_{<t}$ denotes a
trajectory from the first time step up to $t-1$, $\mathcal{Y}$ is the set of
all valid output trajectories, $x$ is the initial state, sampled from
$\rho_{0}$ (the initial state distribution of the MDP $\mathcal{M}$). Finally,
$\circ$ denotes a concatenation operator.

#### Model

A _Transformer_ (Vaswani et al. 2017) is a sequence-transduction model whose
network architecture relies solely on attention mechanisms (Bahdanau, Cho, and
Bengio 2015), which allows for high parallelization and is shown to be highly
effective in natural language processing (NLP) tasks (Devlin et al. 2019;
Brown et al. 2020). In our case, we use the Transformer to select an action
sequence to be executed in the environment.

The model (parametrized via the network’s weights $\theta$) predicts the
probability distribution over an output space of possible trajectories
$\mathbf{y}$ given an initial state $x$, which is factorized as:

$P_{\theta}(\mathbf{y}|x)=\prod_{t=1}^{T}P_{\theta}(y_{t}|\mathbf{y}_{<t},x).$

Where $P_{\theta}(y_{t}|\mathbf{y}_{<t},x)$ typically defines a multinomial
classification model. In such case, assuming a per-dimension vocabulary size
of $\mathcal{V}$, the network’s output layer consists of logits over a
vocabulary of size $\mathcal{V}$. Namely, for a discrete output variable
$y_{t}$ (which represents a dimension of a state, action, reward, or reward-
to-go estimate), the outputs correspond to a mapping: $y_{t}^{i}\mapsto p_{i}$
for $i\in\\{1,\ldots\mathcal{V}\\}$ where $p_{i}$ is the probability of the
next token to be $y_{t}^{i}$.

As we use the original network as-is, we refer the reader to the original
paper for the architecture’s details.

###  4.2 Online Decoding (S2)

To decode a sequence from a trained model, we use the model’s autoregressive
nature to predict a single token, $y_{t}$, at each step. Given an initial
state $x$, we generate tokens sequentially using some heuristic-based
algorithms. The core algorithm providing the foundation of our planning
techniques is beam search (BS) (Lowerre and Reddy 1976). Here we describe BS
as a meta-algorithm. As we will see in Sec. 5, this will significantly
simplify the presentation of our planning algorithm WSTS. Accordingly, we
present in Alg. 1 BS as an algorithm that takes in additional two functions as
inputs- a scoring function score and a filter function filter and returns a
single approximated best trajectory (we will define best shortly). At each
iteration, BS expands all currently-considered sequences and creates new
candidate sequences. These candidate sequences are evaluated using the score
function. Subsequently, BS uses these scores as inputs to the filter function
and decides which $\mathcal{B}$ sequences to keep for the next iteration.

Algorithm 1 Beam Search (BS)

Input: start state $x$, scoring function score(), filter function filter(),
sequence model $P_{\theta}(y_{t}|\mathbf{y}_{<t},x)$
Parameters: beam width $\mathcal{B}$, planning horizon $H$
Output: Approx. single best trajectory.

$C_{0}\leftarrow\\{x\\}$

for $t\leftarrow 1$ to $H$ do

$C\leftarrow\\{\\}$

for all $\mathbf{y}_{<t}\in C_{t-1}$ do

// Autoregressively simulate transitions

$(s_{t},a_{t},r_{t},R_{t})\sim P_{\theta}(y_{t}|\mathbf{y}_{<t},x)$

$C\leftarrow C\cup\\{(\mathbf{y}_{<t}\circ(s_{t},a_{t},r_{t},R_{t}))\\}$

$\mathbf{w}\leftarrow{\tt\textbf{score}}(C)$

$C_{t}\leftarrow{\tt\textbf{filter}}(C,\mathcal{B},\mathbf{w})$

return $C.max()$

As we’ll see next, different definitions of ”best” result in different
implementations of these functions and, in turn, correspond to different
search algorithms that may be used during online decoding. For example, in NLP
and imitation learning, the overarching objective is to find the most probable
sequence under the model at inference times (commonly known as maximum a
posteriori, or MAP, decoding (Meister, Cotterell, and Vieira 2020)). This
corresponds to solving the following optimization problem:

$y^{*}={\arg\max}_{\mathbf{y}\in\mathcal{Y}}\log P_{\theta}(\mathbf{y}|x).$
(1)

In this case, finding an exact solution is computationally hard. An
approximate solution is often found using top-$K$ sampling (Fan, Lewis, and
Dauphin 2018), a widely used BS variant with the following selection of
scoring and filter functions: The scoring function is the conditional
probability distribution ${\tt score}(\mathbf{y}_{<t}\circ y_{t})=\log
P_{\theta}(y_{t}|\mathbf{y}_{<t})$. The filter function randomly picks the
$\mathcal{B}$ most likely next tokens with proportion to their score.

In the context of offline RL, Janner, Li, and Levine (2021) used a different
score() function and modified the top-K sampling BS variant to decode
trajectories that achieve the maximum cumulative expected rewards. For this
case, the log probabilities of transitions are replaced by the log probability
of the predicted reward signal. Consequently, the scoring function is the
expected cumulative reward added to the reward-to-go estimate while the filter
function remains unchanged. For simplicity, here we will refer to this variant
simply as Expectation Maximizing Beam Search EM-BS.

As we will see shortly (Sec. 5), we explain WSTS using a particular choice of
these functions.

###  4.3 Portfolio Optimization

In finance, a _portfolio_ is defined as a combination of financial assets,
each typically associated with some expected reward and risk. To form the
portfolio, given such $N$ different risky assets and wealth $w$, we need to
decide how much wealth to invest in each asset, i.e., determine the vector of
weights on assets $1$ to $N$: $w=(w_{1},w_{2},...,w_{N})$. The _portfolio
optimization_ problem corresponds to selecting the best portfolio out of the
set of all portfolios being considered, according to some objective. These
objectives typically balance expected reward and risk allowing investors a
principled way to maximize return while bounding risk. Modern portfolio theory
(Markowitz 1952) or mean-variance analysis solves the portfolio selection
problem by taking only the mean and variance of the portfolio into
consideration, weighing the risk, expressed as variance, against the expected
return.

_Expected utility theory_ estimates the utility of action when the outcome is
uncertain. It takes into account that individuals may be risk-averse, meaning
they tend to prefer outcomes with low uncertainty to those with high
uncertainty (Tversky and Kahneman 1992). The expected utility maximization is
consistent with mean-variance analysis if the utility function is quadratic or
if the asset returns are normally distributed (Levy and Markowitz 1979). In
such a case, the constrained portfolio-optimization problem can be formulated
as a utility-maximization problem. Specifically, given some risk-aversion
parameter $\delta$, assets mean and variance vectors $\mu$ and $\Sigma$, the
portfolio-optimization problem is solved by computing a weight vector $w$
dictating the relative amount to invest in each asset by solving the following
optimization problem:

$\max_{w}\quad w^{T}\mu-\frac{\delta}{2}w^{T}\Sigma w.$ (2)

##  5 Method: Wall Street Tree Search (WSTS)

In this section, we present “Wall Street Tree Search” (WSTS)—our approach for
online risk-aware decoding. WSTS expands beam search with portfolio
optimization to create a new planning algorithm for sequential decision-making
under uncertainty. Like portfolio optimization, our risk-aware planning
algorithm is about budget allocation to different assets. However, in our
setting, the budget is not money— it is the computational effort, and the
assets are the candidate trajectories. In our variation of beam search, we
allocate our limited amount of node-expansion budget to determine which
$\mathcal{B}$ trajectories we keep at each time step. To this end, we use
portfolio optimization to weigh trajectories according to their “risk” (which
is uncertainty in our setting), expected return, and our risk-aversion
parameter $\delta$ (we discuss the implications of this parameter in Sec. 6).

We start by describing how we compute the mean vector and covariance matrix
that will be used in Eq. 2. We then continue to detail how WSTS uses portfolio
optimization to instantiate the generic score and filter functions described
in Sec. 4.2.

#### Mean Vector and Covariance Matrix

Recall that to solve the portfolio-optimization problem (Eq. 2), in addition
to the risk-aversion paramaeter $\delta$, we require (i) a vector of expected
returns, $\mu$, and (ii) a covariance matrix $\Sigma{}$ of the assets. To
estimate these quantities, we use the Transformer decoder sequence model
trained in the model-learning stage (Sec. 4.1). Namely, we use the variance
computed from the Transformer’s output as a measure of predictive uncertainty
and as a proxy indicating the amount of distribution shift. This approach is
based on Desai and Durrett (2020), which showed that Transformer-based models
are well-calibrated when trained with temperature scaling (Guo et al. 2017) as
is done in our setting.

To start, recall that the Transformer is a multinomial classification model
which at inference time autoregressively outputs a probability $p_{i}$ of the
subsequent trajectory token $y_{t}^{i}$ at step $t$, conditioned on the
preceding trajectory tokens $\mathbf{y}_{<t}$. Hence, by conditioning on each
of the partial trajectories, the mean $E[y_{t}]$ and variance
$\mathrm{Var}[y_{t}]$ of each output random variable $y_{t}$ are simply
$E[y_{t}]=\sum_{i=1}^{\mathcal{V}}p_{i}y_{t}^{i}$ and
$\mathrm{Var}[y_{t}]=\sum_{i=1}^{\mathcal{V}}p_{i}\cdot(y_{t}^{i}-\mu)^{2}$,
respectively.

Now, using these values to predict the vector of expected returns, $\mu$, and
the covariance matrix $\Sigma{}$ for all the candidate solutions on each time
step requires more care. The longer the effective planning horizon, the more
error the model inaccuracy introduces. Thus, in contrast to the common
approach in RL where immediate rewards are incentivized over long-term rewards
via a discount factor $\gamma\in(0,1]$ which exponentially scales down the
rewards after each step, in our case, we want to account for the potential
effect of compounding future risks. Specifically, in addition to scaling down
the mean, we also scale up the variance, informing our downstream planner that
the uncertainty increases for future predictions. Consequently, the means
vector, $\mu$, is an $N$-dimensional vector consisting of the discounted means
of the cumulative reward plus the reward-to-go estimate associated with each
one of the $N$ candidate trajectories. The covariance matrix, $\Sigma$, is an
$N\times N$ diagonal matrix consisting of the discounted variances of the
cumulative reward plus reward-to-go estimate associated with each candidate
trajectory along its main diagonal. Namely, $\mu_{j}$ and $\sigma_{j}$ which
are the $j$th entry in $\mu$ and $\Sigma$, respectively are defined as:

$\displaystyle\mu_{j}$
$\displaystyle=\sum_{t=1}^{T-1}\gamma^{t}\mathrm{E[r_{t}^{i}]}+\gamma^{T}\mathrm{E[\hat{R}_{T}^{i}]},$
$\displaystyle\sigma^{2}_{j}$
$\displaystyle=\sum_{t=1}^{T-1}\gamma^{-2t}\mathrm{Var[r_{t}]}+\gamma^{-2T}\mathrm{Var[\hat{R}_{T}]}.$

#### score

function: Given the mean vector $\mu$, and the covariance matrix $\Sigma$, we
compute a score for each of the $N$ candidate trajectories available at time
step $t$ by solving a portfolio-optimization problem via Eq. 2 with the input
risk-aversion parameter, $\delta$. The output of the score function is the
solution to the portfolio-optimization problem, namely the wealth coefficients
$w_{1},\ldots,w_{N}$, representing the proportion of the computational time
(money) to invest in exploring each trajectory (asset).

#### filter

function: Given the coefficients $w_{1},\ldots,w_{N}$, which are computed
using the above score function, our filter function samples $\mathcal{B}$
trajectories (with repetition) where the $j$th trajectory is sampled with
probability $w_{j}$. We underline that a trajectory may be sampled more than
once. In such cases, the search is biased toward more-promising trajectories.
Thus, our practical beam width is smaller than $\mathcal{B}$ (but never
bigger).

##  6 Experiments and Results

In our experiments, we aim to study the following questions:

  * Q1

How does WSTS perform compared to prior approaches?

  * Q2

How to choose the risk-aversion parameter $\delta$?

#### Experimental setup:

Our experimental evaluation focuses on continuous control tasks from the
MuJoCo physics engine, which are formulated as RL tasks and aims to facilitate
research and development in robotics. We focus on Gym-MuJoCo locomotion tasks:
Walker2d, HalfCheetah, and Hopper (displayed in Fig. 2). To this end, we use
the offline dataset from the widely used D4RL benchmark (Fu et al. 2020).

Figure 2:  The Gym-MuJoCo tasks, from left to right: HalfCheetah, Hopper, and
Walker2d.

The different dataset settings of D4RL are described below:

  1. 1.

The “medium” dataset is generated by collecting 1 million samples from a
partially-trained policy.

  2. 2.

The “medium-replay” dataset records all samples observed during training until
the policy reaches the “medium” level of performance.

  3. 3.

The “medium-expert” dataset is generated by mixing 1 million samples generated
by the medium policy concatenated with 1 million samples generated by an
expert policy.

We ran all tests on NVIDIA Tesla V100 GPUs. For both planning algorithms, EM-
BS and WSTS 111Code will be made publicly available upon paper publication, we
take an MPC approach, interleaving planning and execution (Sec. 4). For both
algorithms, we used the trained network weights and the same hyperparameters
provided by the original paper’s authors (Janner, Li, and Levine 2021).

#### (Q1) How does WSTS perform compared to prior approaches?

In this section, we compare WSTS with SOTA methods spanning other approaches
of offline RL. However, we focus our comparison on a EM-BS, which is the
decoding algorithm employed in Trajectory Transformer (Janner, Li, and Levine
2021). We concentrate our comparison on EM-BS for two reasons:

  1. 1.

When using EM-BS to decode the Trajectory Transformer, it was shown to
outperform prior offline RL methods on D4RL locomotion tasks. We will shortly
show that we have improved upon these results. Consequently, WSTS performs on
par or better than the other offline RL methods (see Table 1).

  2. 2.

Our work focuses on the decoding procedure. Therefore we want to investigate
the contribution of this individual component to the overall system
performance. To this end, we modify only the overarching planning algorithm
that decodes the Transformer’s outputs from EM-BS to WSTS.

Dataset | Environment | CQL | MOPO | MBOP | IQL | DT | TT (EM-BS) | WSTS | $\delta$
---|---|---|---|---|---|---|---|---|---
Med-Replay | HalfCheetah | $45.5$ | 53.1 | $42.3$ | $44.2$ | $36.6$ | $41.9\pm 2.5$ | $44.8\pm 0.3$ | $1.0$
Med-Replay | Hopper | 95.0 | $67.5$ | $12.4$ | $94.7$ | $82.7$ | $91.5\pm 3.6$ | $94.2\pm 2.6$ | $2.0$
Med-Replay | Walker2d | $77.2$ | $39.0$ | $9.7$ | $73.9$ | $66.6$ | $82.6\pm 6.9$ | $\textbf{86.1}\pm 3.8$ | $2.0$
Medium | HalfCheetah | $44.0$ | $42.3$ | $44.6$ | $47.4$ | $42.6$ | $46.9\pm 0.4$ | $\textbf{47.6}\pm 0.2$ | $0.1$
Medium | Hopper | $58.5$ | $28.0$ | $48.8$ | $66.3$ | 67.6 | $67.4\pm 2.9$ | $67.0\pm 4.7$ | $0.5$
Medium | Walker2d | $72.5$ | $17.8$ | $41.0$ | $68.3$ | $74.0$ | $79.0\pm 2.8$ | $\textbf{82.1}\pm 0.8$ | $1.0$
Med-Expert | HalfCheetah | $91.6$ | $63.3$ | 105.0 | $86.7$ | $86.8$ | $95.0\pm 0.2$ | $94.4\pm 0.3$ | $0.1$
Med-Expert | Hopper | $105.4$ | $23.7$ | $55.1$ | $91.5$ | $107.6$ | $110.0\pm 2.7$ |  111.0 $\pm 1.6$ | $0.1$
Med-Expert | Walker2d | $108.8$ | $44.6$ | $70.2$ | 109.6 | $108.1$ | $101.9\pm 6.8$ | $107.0\pm 1.5$ | $0.1$
Average |  | $77.6$ | $42.14$ | $47.8$ | $75.8$ | $74.7$ | $78.9$ | 81.6 |
Table 1: WSTS compared against other prior SOTA methods spanning other
approaches of offline RL: CQL (Kumar et al. 2020), MOPO (Yu et al. 2020), MBOP
(Argenson and Dulac-Arnold 2020), IQL (Kostrikov, Nair, and Levine 2021), DT
(Chen et al. 2021) and Trajectory Transformer (TT which runs EM-BS) (Janner,
Li, and Levine 2021). variants correspond to the mean and standard error over
15 random seeds. Results for the other algorithms are taken from the original
papers.

Table (1) shows that WSTS outperforms EM-BS. However, reporting point
estimates do not tell the whole story as they obscure some key aspects of the
comparison. Indeed WSTS matches or surpasses the average score of EM-BS, but
what is even more critical is that WSTS is consistently more stable, reducing
the variance considerably. To illustrate this, we first visualize the
differences between WSTS and EM-BS using a box plot and present the results in
Fig. 3.

Figure 3:  Score as a function of the dataset for WSTS and EM-BS obtained over
15 random seeds. The box corresponds to quartile Q1 through Q3, with a line at
the median (Q2) and a red dot on the mean. Whiskers are restricted to a
maximum of 1.5 times the interquartile range.

Though the box plot helps understand better the difference between EM-BS and
WSTS, it still doesn’t provide a complete picture. Since evaluating runs is
often computationally demanding in deep RL, only a small number of runs
(usually less than 10) are generally assessed. For such a low number of runs,
reporting reliable results is challenging. We address this by following the
methodology proposed by Agarwal et al. (2021), which suggests presenting
confidence intervals (CIs), which are a range of values that are likely to
include the correct value with a certain degree of confidence. Accordingly, in
Fig. 4, we show $95\%$ CI values of normalized mean, median, and interquartile
mean (IQM) scores (performance on middle $50\%$). The CIs for WSTS are
narrower and with higher averages than EM-BS, indicating that WSTS does not
only outperform EM-BS but is also more stable than EM-BS.

Figure 4:  Aggregate metrics on Gym-MuJoCo tasks with $95\%$ CI values based
on 15 independent runs per each of the nine environments. The CIs are
estimated using stratified sampling, with a total of $9\times 15$ normalized
scores per algorithm. The interval estimates are created with bootstrapping,
i.e., calculating the aggregate metrics by randomly re-sampling scores with
replacement from this data.

#### (Q2) How to choose the risk aversion parameter $\delta$?

Offline RL usually requires expensive online rollouts for hyperparameters
search (for example, Janner, Li, and Levine (2021) used six hyperparameters),
which makes hyperparameters undesirable. However, in our case, we can get a
hint on how to set the risk aversion parameter’s value properly by observing
the connection between the risk aversion parameter and the offline training
samples: scilicet, the training dataset size (relative to the environment
complexity), and the data quality.

Generally, the risk-aversion parameter $\delta$ takes the highest values for
the medium-replay dataset; the medium dataset takes intermediate values; for
the medium-expert dataset, the risk aversion parameter takes the lowest
values, as seen in the rightmost column of Table 1. This is expected as the
more expert the behavior policy that generated the training samples is, the
more we trust it and hence the less risk-averse we need to be. This
observation is concise with imitation learning. In the extreme case of expert
demonstrations with unlimited data, the optimal strategy would be to select a
single action with the maximum expected reward. Such an approach is equivalent
to having $\delta-0$, which narrows the beam width $\mathcal{V}$ to one.

The number of training samples also plays a role when setting the value of
$\delta$. More training samples reduce the subjective uncertainty, which in
turn can reduce the risk aversion parameter value. In our experiments, we can
get affirmation for this observation when examining the comparison between EM-
BS and WSTS in the medium-expert dataset (left boxplot in each image of Fig.
3). Recall that a well-trained behavior policy generated the medium-expert
dataset. In addition, this dataset contains twice as many samples as other
datasets. For such a dataset, a behavior that tries to maximize the expected
reward and does not account for the risk (i.e., low the risk-aversion
parameter) is a good option. Indeed WSTS performs approximately as EM-BS does
on this dataset.

It is also interesting to examine how the risk aversion parameter $\delta$
affects the performance. This question is relevant for scenarios where system
engineers can select the amount of risk aversion. Such a selection generally
depends on the application and system domain. Fig. 5 presents a kernel density
estimate plot, visualizing the distribution of scores in the Walker2d medium-
replay dataset, that demonstrates that: (i) Our algorithm is robust when using
diversified $\delta$ values and that (ii) Decreasing the value of $\delta$
allows achieving higher results if we are willing to take the risk of failing.

Figure 5:  A kernel density estimate plot, presenting WSTS sensitivity to the
risk aversion parameter $\delta$. Results showed for Walker2d medium-replay
dataset over ten random seeds for each risk aversion value.

##  7 Conclusion and Future Work

The bitter lesson (Sutton 2019) is based on the historical observations that
generic models tend to overtake specialized domain-specific approaches. This
observation is manifested in a recently emerging body of work that shows that
Transformer models that work well in other domains, such as natural language
processing, can also provide an effective and generic solution to the offline
RL problem. However, prediction with Transformers is currently made by
repurposing the same tools from NLP, ignoring uncertainty. In this work, we
addressed the question of how we can modify the decoding algorithms of large-
scale language models– the Transformer to tackle the uncertainty present in
offline problems.

To this end, we introduced WSTS, a planning algorithm that suggests a
constructed way to control the risk. Our planning algorithm leverages the
capacity of the learned model to generalize to states outside the static batch
support. Still, it is cautious when drifting to states where the model can’t
give a confident prediction based on the offline dataset. For such states, our
planner weighs this risk vs. the expected return by incorporating modern
portfolio theory into sequential decision-making.

In this work, we rely on temperature scaling for calibration. However, a
question remains whether different Bayesian formulations of deep learning that
can potentially improve predictive uncertainty quantification can further
improve our results. Moreover, our planning algorithm penalizes trajectories
with uncertainty. However, we are more interested in the risk of incurred loss
than the variability of returns measured by standard deviation; hence we
should only penalize for the risk of low returns (but we need not penalize the
risk of high returns). We hope to explore these limitations in future work
using better-suited portfolio optimization methods.

We demonstrated our algorithm on an offline model-based RL setting. Our
experiments show that our method for controlling the risk is advantageous in
the offline RL setting, as WSTS outperforms state-of-the-art offline RL
methods in the D4RL standard benchmark. Moreover, it is consistently more
stable, reducing the variance of the results considerably, an essential
property in a range of real-world domains where real-world failures can be
costly or dangerous.

We conclude with a quote by Warren Buffett (Goodman 2018), representing his
view on investors’ uncertainty which also captures the essence of our approach
to planning with an uncertain model applied to offline RL.

> There is nothing wrong with a ‘know nothing’ investor who realizes it. The
> problem is when you are a ‘know nothing’ investor, but you think you know
> something.

## References

  * Agarwal, Schuurmans, and Norouzi (2020) Agarwal, R.; Schuurmans, D.; and Norouzi, M. 2020.  An Optimistic Perspective on Offline Reinforcement Learning.  In _International Conference on Machine Learning (ICML)_ , 104–114.
  * Agarwal et al. (2021) Agarwal, R.; Schwarzer, M.; Castro, P. S.; Courville, A. C.; and Bellemare, M. 2021\.  Deep Reinforcement Learning at the Edge of the statistical precipice.  In _Advances in Neural Information Processing Systems (NeurIPS)_ , 29304–29320.
  * Argenson and Dulac-Arnold (2020) Argenson, A.; and Dulac-Arnold, G. 2020.  Model-Based Offline Planning.  In _International Conference on Learning Representations (ICLR)_ , 14129–14142.
  * Bahdanau, Cho, and Bengio (2015) Bahdanau, D.; Cho, K. H.; and Bengio, Y. 2015.  Neural Machine Translation by Jointly Learning to Align and Translate.  In _International Conference on Learning Representations (ICLR)_.
  * Brown et al. (2020) Brown, T.; Mann, B.; Ryder, N.; Subbiah, M.; Kaplan, J. D.; Dhariwal, P.; Neelakantan, A.; Shyam, P.; Sastry, G.; Askell, A.; Agarwal, S.; Herbert-Voss, A.; Krueger, G.; Henighan, T.; Child, R.; Ramesh, A.; Ziegler, D.; Wu, J.; Winter, C.; Hesse, C.; Chen, M.; Sigler, E.; Litwin, M.; Gray, S.; Chess, B.; Clark, J.; Berner, C.; McCandlish, S.; Radford, A.; Sutskever, I.; and Amodei, D. 2020.  Language Models are Few-Shot Learners.  In _Advances in Neural Information Processing Systems (NeurIPS)_ , 1877–1901.
  * Chen et al. (2021) Chen, L.; Lu, K.; Rajeswaran, A.; Lee, K.; Grover, A.; Laskin, M.; Abbeel, P.; Srinivas, A.; and Mordatch, I. 2021.  Decision Transformer: Reinforcement Learning via Sequence Modeling.
  * Covington, Adams, and Sargin (2016) Covington, P.; Adams, J.; and Sargin, E. 2016.  Deep Neural Networks for Youtube Recommendations.  In _ACM Conference on Recommender Systems_ , 191–198.
  * Desai and Durrett (2020) Desai, S.; and Durrett, G. 2020.  Calibration of Pre-trained Transformers.  In _Empirical Methods in Natural Language Processing (EMNLP)_ , 295–302.
  * Devlin et al. (2019) Devlin, J.; Chang, M. W.; Lee, K.; and Toutanova, K. 2019.  BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding.  In _North American Chapter of the Association for Computational Linguistics: Human Language Technologies (NAACL HLT)_ , 4171–4186.
  * Fan, Lewis, and Dauphin (2018) Fan, A.; Lewis, M.; and Dauphin, Y. 2018.  Hierarchical Neural Story Generation.  In _Association for Computational Linguistics (ACL)_ , 889–898.
  * Fu et al. (2020) Fu, J.; Kumar, A.; Nachum, O.; Tucker, G.; and Levine, S. 2020.  D4RL: Datasets for Deep Data-Driven Reinforcement Learning.
  * Fujimoto, Meger, and Precup (2019) Fujimoto, S.; Meger, D.; and Precup, D. 2019.  Off-policy Deep Reinforcement Learning Without Exploration.  In _International Conference on Machine Learning (ICML)_ , 2052–2062.
  * Goodman (2018) Goodman, J. 2018.  _The Best Quotes Book: 555 Daily Inspirational and Motivational Quotations by Famous People_.  Lulu.
  * Gottesman et al. (2018) Gottesman, O.; Johansson, F.; Meier, J.; Dent, J.; Lee, D.; Srinivasan, S.; Zhang, L.; Ding, Y.; Wihl, D.; Peng, X.; Yao, J.; Lage, I.; Mosch, C.; Lehman, L.-w. H.; Komorowski, M.; Komorowski, M.; Faisal, A.; Celi, L. A.; Sontag, D.; and Doshi-Velez, F. 2018.  Evaluating Reinforcement Learning Algorithms in Observational Health Settings.
  * Guo et al. (2017) Guo, C.; Pleiss, G.; Sun, Y.; and Weinberger, K. Q. 2017.  On Calibration of Modern Neural Networks.  In _International Conference on Machine Learning (ICML)_ , 1321–1330.
  * Janner, Li, and Levine (2021) Janner, M.; Li, Q.; and Levine, S. 2021.  Offline Reinforcement Learning as One Big Sequence Modeling Problem.
  * Jaques et al. (2019) Jaques, N.; Ghandeharioun, A.; Shen, J. H.; Ferguson, C.; Lapedriza, A.; Jones, N.; Gu, S.; and Picard, R. 2019.  Way Off-Policy Batch Deep Reinforcement Learning of Implicit Human Preferences in Dialog.
  * Kidambi et al. (2020) Kidambi, R.; Rajeswaran, A.; Netrapalli, P.; and Joachims, T. 2020.  MOReL : Model-Based Offline Reinforcement Learning.  In _Advances in Neural Information Processing Systems (NeurIPS)_ , 21810–21823.
  * Kocsis and Szepesvári (2006) Kocsis, L.; and Szepesvári, C. 2006.  Bandit Based Monte-Carlo Planning.  In _European Conference on Machine Learning (ECML)_ , 282–293.
  * Kostrikov, Nair, and Levine (2021) Kostrikov, I.; Nair, A.; and Levine, S. 2021.  Offline Reinforcement Learning with Implicit Q-Learning.
  * Kumar et al. (2019) Kumar, A.; Fu, J.; Tucker, G.; and Levine, S. 2019.  Stabilizing Off-Policy Q-Learning via Bootstrapping Error Reduction.  In _Advances in Neural Information Processing Systems (NeurIPS)_ , 11761–11771.
  * Kumar et al. (2020) Kumar, A.; Zhou, A.; Tucker, G.; and Levine, S. 2020.  Conservative Q-Learning for Offline Reinforcement Learning.  In _Advances in Neural Information Processing Systems (NeurIPS)_ , 1179–1191.
  * Levine et al. (2020) Levine, S.; Kumar, A.; Tucker, G.; and Fu, J. 2020.  Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems.
  * Levy and Markowitz (1979) Levy, H.; and Markowitz, H. M. 1979.  Approximating Expected Utility by a Function of Mean and Variance.  _American Economic Review_ , 69(3): 308–317.
  * Lowerre and Reddy (1976) Lowerre, B.; and Reddy, R. 1976.  The Harpy Speech Recognition System: Performance with Large Vocabularies.  _The Journal of the Acoustical Society of America_ , 60(S1): S10–S11.
  * Markowitz (1952) Markowitz, H. 1952.  Portfolio Selection.  _The Journal of Finance_ , 7(1): 77–91.
  * Meister, Cotterell, and Vieira (2020) Meister, C.; Cotterell, R.; and Vieira, T. 2020.  If Beam Search is the Answer, What Was the Question?  In _Empirical Methods in Natural Language Processing (EMNLP)_ , 2173–2185.
  * Peng et al. (2019) Peng, X. B.; Kumar, A.; Zhang, G.; and Levine, S. 2019.  Advantage-Weighted Regression: Simple and Scalable Off-Policy Reinforcement Learning.
  * Richalet et al. (1978) Richalet, J.; Rault, A.; Testud, J. L.; and Papon, J. 1978.  Model Predictive Heuristic Control. Applications to Industrial Processes.  _Automatica_ , 14(5): 413–428.
  * Sallab et al. (2017) Sallab, A. E.; Abdou, M.; Perot, E.; and Yogamani, S. 2017.  Deep reinforcement learning framework for autonomous driving.  _Electronic Imaging_ , 19: 70–76.
  * Siegel et al. (2020) Siegel, N. Y.; Springenberg, J. T.; Berkenkamp, F.; Neunert, M.; Lampe, T.; Hafner, R.; and Riedmiller, M. 2020.  Keep Doing What Worked : Behavioral Cloning Priors for Fully Offline Learning.  In _International Conference on Learning Representations (ICLR)_.
  * Strehl et al. (2010) Strehl, A.; Langford, J.; Li, L.; and Kakade, S. M. 2010.  Learning from logged implicit exploration data.  _Advances in Neural Information Processing Systems_ , 23.
  * Sutton (2019) Sutton, R. 2019.  The Bitter Lesson.  Incomplete Ideas (blog).
  * Sutton and Barto (1998) Sutton, R.; and Barto, A. 1998.  _Reinforcement Learning: An Introduction_.  MIT Press.
  * Tversky and Kahneman (1992) Tversky, A.; and Kahneman, D. 1992.  Advances in Prospect Theory: Cumulative Representation of Uncertainty.  _Journal of Risk and Uncertainty_ , 5(4): 297–323.
  * Vaswani et al. (2017) Vaswani, A.; Shazeer, N.; Parmar, N.; Uszkoreit, J.; Jones, L.; Gomez, A. N.; Kaiser, L.; and Polosukhin, I. 2017.  Attention is all you need.  In _Advances in Neural Information Processing Systems (NIPS)_ , 6000–6010.
  * Wang et al. (2018) Wang, L.; Zhang, W.; He, X.; and Zha, H. 2018.  Supervised reinforcement learning with recurrent neural network for dynamic treatment recommendation.  In _International Conference on Knowledge Discovery and Data Mining (SIGKDD)_ , 2447–2456.
  * Wu, Tucker, and Nachum (2019) Wu, Y.; Tucker, G.; and Nachum, O. 2019.  Behavior Regularized Offline Reinforcement Learning.
  * Yu, Ren, and Liu (2019) Yu, C.; Ren, G.; and Liu, J. 2019.  Deep inverse reinforcement learning for sepsis treatment.  In _International Conference on Healthcare Informatics (ICHI)_ , 1–3.
  * Yu et al. (2020) Yu, T.; Thomas, G.; Yu, L.; Ermon, S.; Zou, J.; Levine, S.; Finn, C.; and Ma, T. 2020.  MOPO: Model-based Offline Policy Optimization.  In _Advances in Neural Information Processing Systems (NeurIPS)_ , 14129–14142.