Acting upon Imagination: When to Trust Imagined Trajectories in Model Based Reinforcement Learning (2024)

1 Introduction

Reinforcement learning can be successfully applied in continuous control of complex and highly non-linear systems. Algorithms for reinforcement learning can be categorized as model free (MFRL) or model based (MBRL).Using deep learning in MFRL has achieved success in learning complex policies from raw input, such as solving problems with high-dimensional state spaces (Sutton and Barto, 1998)(Mnih etal., 2015) and continuous action space optimization problems with algorithms like deep deterministic policy gradient (DDPG)(Lillicrap etal., 2015), Proximal Policy Optimization (PPO) (Schulman etal., 2017) or Soft Actor-Critic (SAC) (Haarnoja etal., 2018).While significant progress has been made, the high sample complexity of MFRL limits its real-world applicability.Collecting millions of transitions needed to converge is not always feasible in real-world problems, as excessive data collection is expensive, tedious or can lead to physical damage (Williams etal., 2017; Chua etal., 2018). Instead, model based methods are comparably sample efficient. MBRL techniques build a predictive model of the world to imagine trajectories of future states and plan the best actions to maximize a reward.

MBRL uses a dynamics model to predict the outcome of taking an action in an environment with given states. It can bootstrap from existing experiences and is versatile to changing objectives on the fly. Nevertheless, its performance degrades with growing model error. In the general case of nonlinear dynamics guarantees of local optimality do not hold(Janner etal., 2019).Sampling-based MPC (model predictive control) algorithms can be used to address this issue. A Sampling-based MPC generates a large number of trajectories with the goal of maximizing the expected accumulated reward. But, complex environments are typically partially observable and the problem is formulated as a receding horizon planning task. Hence, after executing a single step, trajectories are generated again from scratch to deal with the receding horizon and reduce the impact of model compound error (Rao, 2010).An additional challenge lies in the high cost incurred by frequently generating imagined trajectories from scratch during planning. For each trajectory generated, a sequence of actions has to be evaluated with the dynamics model. Each successive step in a trajectory depends on previous states. So the evaluation of a single trajectory is a recurrent process that cannot be parallelized. Acting upon imagination, the method here proposed, seeks to estimate online the uncertainty of the imagined trajectory and to reduce the planning cost by continuing to act upon an it unless it cannot be trusted.

Therefore, in this work, we present methods for uncertainty estimation, designed to address and improve the computational limitations of shooting MPC MBRL methods. Our main objectiveis the online estimation of uncertainty of the model plan. A second objective is the application of uncertainty estimation to avoid frequent replanning. Here, we propose using the degree to which computations are reduced as a practical and quantifiable proxy for the effectiveness of uncertainty estimation methods. If our methods are successful, we expect to observe a significant decrease in computations without a substantial decay in performance. The balance between computational efficiency and performance demonstrates the reliability of our uncertainty estimation methods. A robust estimation of uncertainty facilitates efficient decision-making and optimizes the use of computational resources.Our contributions are as follows:

Our experimental results on challenging benchmark control tasks demonstrate that the proposed methods effectively leverage accurate predictions as well as dynamically decide when to replan. This approach leads to substantial reductions in training time and promotes more efficient use of computational resources by eliminating unnecessary replanning steps.

2 Related work

Model-based reinforcement learning (MBRL) has been applied in various real-world control tasks, such as robotics. Compared to model-free approaches, MBRL tends to be more sample-efficient(Deisenroth etal., 2013).MBRL can be grouped into four main categories(Zhu etal., 2020):
1) Dyna-style algorithms optimize policies using samples from a learned world model (Sutton, 1990).
2) Model-augmented value expansion methods, such as MVE (Oh etal., 2017), use model-based rollouts to enhance targets for model-free Temporal Difference updates.
3) Analytic-gradient methods can be used when a differentiable world model is available, which adjust the policy through gradients that flow through the model. When compared to traditional planning algorithms that create numerous rollouts to choose the optimal action sequence, analytic-gradient methods are more computationally efficient. Stochastic Value Gradients (SVG) (Heess etal., 2015) provide a new way to calculate analytic value gradients using a generic differentiable world model. Dreamer (Hafner etal., 2020), a milestone in the realm of analytic-gradient model-based RL, demonstrates superior performance in visual control tasks. Dreamer expands upon SVG by facilitating the generation of imaginary rollouts within the latent space
4) Model Predictive Control (MPC) and sampling-based shooting methods employ planning to select actions. They are notably effective for addressing real-world scenarios since excessive data collection is not only costly and tedious, but it can also result in physical damage. Additionally, sampling-based MPC methods have the capacity to bootstrap from existing experiences and rapidly adapt to changing objectives on the fly. However, a significant drawback to these approaches is their computationally intensive nature (Rao, 2010; Chua etal., 2018). The present work belongs into the latter category.

Recently, it was demonstrated that parametric function approximators, neural networks (NN), efficiently reduce sample complexity in problems with high-dimensional non-linear dynamics(Nagabandi etal., 2018).Random shooting methods artificially generate large number of actions(Rao, 2010) and MPC is used to select candidate actions(Camacho etal., 2004). I.e.Williams etal. (2017) and Drews etal. (2017) introduced a sampling-based MPC with dynamics model to sample a large number of trajectories in parallel. A two-layer NN trained from maneuvers performed by a human pilot was superior compared with a physics model built using vehicle dynamics equations from bicycle models. One disadvantage is that NNs cannot quantify predictive uncertainty.

Lakshminarayanan etal. (2016) utilized the ensembles of probabilistic NNs to determine predictive uncertainty. Kalweit and Boedecker (2017) used a notion of uncertainty within a model-free RL (MFRL) agent to switch executing imagined trajectories from a dynamics model when MFRL agent has a high uncertainty. Conversely, Buckman etal. (2018) used imagined trajectories to improve the sample complexity of a MFRL agent. They improved the Q function by using ensembles to estimate uncertainty and prioritize trajectories thereupon.

Measuring the reliability of a learned dynamics model when generating imagined trajectories has been proposed in several works.Chua etal. (2018) identified two types of uncertainty: aleatoric (inherent to the process) and epistemic (resulting from datasets with too few data points). The former is the uncertainty inherent to the process, and the latter results from datasets with too few data points.They combined uncertainty aware probabilistic ensembles in the trajectory sampling of the MPC with a cross entropy controller and achieved asymptotic performance comparable to Proximal Policy Optimization (PPO) Schulman etal. (2017) or Soft Actor-Critic (SAC) Haarnoja etal. (2018) with more sample efficient convergence. Janner etal. (2019) generated (truncated) short trajectories with a probabilistic ensemble to train the policy of a MFRL agent, thus improving significantly its sampling efficiency. Yu etal. (2020) also exploits the uncertainty of the dynamics model to improve policy learning on an offline RL setting. They learn policies entirely from a large batch of previously collected data with rewards artificially penalized by the uncertainty of the dynamics. These works focus on sample efficiency and improving the performance, our work proposes novel methods to estimate the uncertainty of the dynamics model to determine when to replan.

The authors in (Hafez etal., 2020) propose a analytic-gradient-based method that considers the reliability of the learned dynamics model used for imagining the future. They evaluate their approach in the context of enhancing vision-based robotic grasping, aiming to improve sample efficiency in sparse reward environments. In contrast to their method, ours does not require the use of numerous local dynamics models or a self-organizing map. Instead, we introduce a technique that exploits the uncertainty of the dynamics model to estimate the uncertainty of plan during execution, primarily aimed at minimizing replanning within an MPC framework. Close to our work, Zhu etal. (2020) studied the discrepancy between imagination. Their method allows for policy generalization to real-world interactions by optimizing the mutual information between imagined and real trajectories, while simultaneously refining the policy based on the imagined trajectories. However, their focus is on analytic gradients MBRL only, our method can be applied to any MBRL which yields a notion of uncertainty and we focus on shooting methods, which are still the first choice in domains like self driving cars (Williams etal., 2017).

Hansen etal. (2022) obtained state of the art performance in terms of reward and training time on diverse continuous control tasks by significantly improving Model-Augmented Value Expansion methods. Their approach effectively combines the strengths of both MFRL and MBRL. They adopt a learned task-oriented latent dynamics model for localized trajectory optimization over a short horizon. Furthermore, they utilize a learned terminal value function to estimate the long-term returns. However, their method still necessitates the learning of the value function. Depending on the context, this could present challenges when compared to shooting methods.

Nevertheless, shooting MPC methods still suffer from expensive computationChua etal. (2018); Zhu etal. (2020).Thus, our research seeks to reduce the amount of computation continuing to act upon trajectories that seem trustworthy. Our solution builds upon results ofChua etal. (2018), using probabilistic ensembles and cross entropy in the MPC.

3 Preliminaries

RL aims to learn a policy that maximizes the accumulated reward obtained from the environment. At each time $t$, the agent is at a state $s_{t}\in S$, executes an action $a_{t}\in A$ and receives from the environment a reward $r_{t}=r(s_{t},a_{t})$ and a state $s_{t+1}$ according to some unknown dynamics function $f:S\times A\to S$. The goal is then to maximize the sum of discounted rewards $\sum_{i=t}^{\infty}\gamma^{(i-t)}r(s_{i},a_{i})$, where $\gamma\ \in[0,1]$.MBRL uses a discrete time dynamics model $\hat{f}=(s_{t},a_{t})$ to predict the future state $\hat{s}_{t+\Delta_{t}}$ after executing action $a_{t}$ at state $s_{t}$. To reach a state into the future, the dynamics model evaluates sequences of actions, $a_{t:t+H}=(a_{t},\ldots,a_{t+H-1})$ over a longer horizon $H$, to maximize their discounted reward $\sum_{i=t}^{t+H-1}\gamma^{(i-t)}r(s_{i},a_{i})$.Due to partial observability of the environment and the error of the dynamics model $\hat{f}$ in predicting the real physics $f$, the controller typically executes only one action $a_{t}$ in the trajectory and the optimization is solved again with the updated state $s_{t+1}$.Algorithm1 outlines the general steps.When training from scratch, the dynamics model $\hat{f}_{\theta}$ is learned with data, $\mathcal{D}_{env}$, collected on the fly. With $\hat{f}_{\theta}$, the simulator starts and the controller is called to plan the best trajectory resulting in $a^{*}_{t:t+H}$. Only the first action of the trajectory $a_{t}^{*}$ is executed in the environment and the rest is discarded. The data collected from the environment is added to $\mathcal{D}_{env}$ and $\hat{f}_{\theta}$ is trained further. MBRL requires a strategy to generate an action $a_{t}$, given a state $s_{t}$, a discrete time dynamics model $\hat{f}=(s_{t},a_{t})$ to predict the state $s_{t+1}$, and a reward function $r_{t}=r(s_{t},a_{t})$.

4 The Promise of Imagination

Generating trajectories is an essential part of the entire process. The predicted states within those trajectories may have a high variance and their quality will depend on the complexity of the environment as well as the number of steps in the future, H. Estimation and online update of uncertainty is needed to determine if a trajectory is reliable. We contend that when predicted trajectories are reliable it is not necessary to frequently replan them.

Starting from a state at time $t$, a run of the planner yields the optimal set of actions $a^{*}$ of H steps. Using $a^{*}$, the dynamics model yields the probability of the next state from which we can sample the next state and confidence $\hat{s}_{t+1},\sigma_{t+1}\sim p^{*}_{s(t+1)}$ and future reward $\hat{r}_{t}\sim p^{*}_{r(t+1)}$.Thus, one step sampled from the imagined trajectory is composed of $\langle s_{t},a_{t},\hat{s}_{t+1},\sigma_{t+1},\hat{r}_{t+1}\rangle$, where $s_{t}$ represents the current state, $a_{t}$ the action to be taken, $\hat{r}_{t+1}$ is the predicted reward if $a_{t}$ is executed, $\hat{s}_{t+1}$ the predicted next state and $\sigma_{t+1}$ the confidence bound issued by the dynamics model for the prediction. Then, the planner generates iteratively the entire trajectory $\tau_{t}$ of H steps, where each step $i$ is composed of the probability distribution $p^{*}_{s(t+i)}$.

Our methods stem from the following information obtained after executing each step in a trajectory: instantaneous deviation between predicted and real outcome, and impact on the projected plan, see Fig.1.Executing the action $a_{t}$ from the imagined trajectory $\tau_{t}$, yields a real-world state $s_{t+1}$ and reward $r_{t+1}$. The state $s_{t+1}$ is expected to fall within the uncertainty of the model $\hat{s}_{t+1},\sigma_{t+1}\sim p^{*}_{s(t+1)}$. Given $f(s_{t},a_{t})=s_{t+1}$ and $\hat{f}(s_{t},a_{t})=\hat{s}_{t+1}$, the instantaneous error $\epsilon_{t}=|s_{t+1}-\hat{s}_{t+1}|\geq 0$ can be measured. The error at step $i$, $\epsilon_{t+i}$ refers to the error after trajectory calculation and executing $a_{t}$. We can model the error distribution and observe if new errors fall within.This refers to immediate effect on the last state.As regards impact on future actions, forward application of the model to the remainder of the actions starting at the new state yields a new trajectory with projected state $p^{*}_{s(t+1:t+H)}$ and projected reward $p^{*}_{r(t+1:t+H)}$, which can be compared with the planned expected outcomes.

Trajectory Quality Analysis:

As a preliminary experiment, we wanted to analyze the quality of imagined trajectories with a trained dynamics, to determine boundaries of how many actions can be executed without deviating from the plan. We analyzed imagined trajectories on agents in the MuJoCoTodorov etal. (2012) physics engine, with the Cartpole (CP) environment $S\in\mathbb{R}^{4},A\in\mathbb{R}^{1}$, $TaskH$ 200, $H$ 25, with $TaskH$ being the task horizon and $H$ the trajectory horizon. The additional material shows similar findings in other environments.A dynamics model was pre-trained in conventional MBRL-MPC, running Algorithm1 five times from scratch, with trajectory (re-)planning after executing each single action. The best performing model was selected for the analysis.The procedure consisted in collecting the errors of predicted and actual state when avoiding the re-planning for $n$ steps.For each $n\in\{0,...19\}$ the algorithm was run for $TaskH$ steps, $10$ times. Therefore, the error at $n=0$ represents the average error of 10 runs executing the first action in the trajectory.Figure2 illustrates the error of predicted trajectories as a function of $n$ steps used (i.e., avoiding re-planning). While the error and its variation increases with $n$, the minimum error at each step is still at the same level (and often lower) than the average error at step 0, where re-planning can not be avoided at all.Generally, re-planning earlier results in a lower error. However, the chart also shows that some trajectories are so reliable that 19 steps can be executed with error lower than the average error of first point in trajectory.

Reward Analysis:Compared to the vector of state errors, the reward has the advantage of being a more compact representation (single scalar). It also provides substantial information. Figure2 right shows the reward of successfully solved task in CP.After 50 steps, the reward does not change significantly and the system is at a local equilibrium. We contend that when the system is at equilibrium the dynamics model can reliably anticipate the outcomes of the agent’s actions, consequentially rewards are expected to remain similar $\Delta_{r}(t)=r_{t}-r_{t-1}\simeq\Delta_{r}(t+1)=r_{t+1}-r_{t}$.

5 Acting Upon Imagination

From the above discussion, the following information is available after executing each step (shown in Fig2):(i) immediate error ($\epsilon_{t+1}$), (ii) the model uncertainty or confidence bounds for an action imagined against its execution ($\hat{s}_{t+1},\sigma_{t+1}\sim p^{*}_{s(t+1)}$), iii) deviation in projected future states($p^{*}_{s(t+1:t+H)}$) and iv) the deviation in projected future rewards ($p^{*}_{r(t+1:t+H)}$). The last two pieces (iii), (iv) are obtained by forward applying the dynamics model with the remainder of the actions starting at the new state.We leverage each piece of available information to develop methods for uncertainty estimation and evaluate their performance in avoiding replanning events.

Algorithm3 presents the core logic of our proposed methods to continue acting upon imagined trajectories and reduce computation. The variable $skip$ is updated with the result of one of four proposed methods in Algorithms4, 5, 6 or 7. Depending on the outcome, replanning can be avoided and computations reduced.If $skip$ is False, only the first predicted action is executed in the environment. Otherwise, subsequent actions from $a^{*}_{t:t+H}$ are executed until the $skip$ flag is set back to False or the $TaskHorizon$ number of steps in the environment is reached.

6 Experiments

Our ultimate goal is to reduce computations while controlling performance decay. Intuitively, we expect that a trained dynamics model can anticipate the outcome of the agent’s actions and, if its predictions are reliable, it can do so for a number of consecutive imagined steps. So, the first experiment uses pre-trained dynamics to assess the potential gains of acting upon imagined trajectories with the proposed methods: N-Skip (NS), FSA, CB, FUT and BICHO. We aim to investigate the amount of re-planning that can be avoided in terms of number of trajectory planning skipped, the impact of our approach on the reward, and the average and variance of steps executed prior to initiating replanning.

We recognize that acting upon imagination has the potential to afford significant time savings while training the dynamics models and can potentially obtain a better result in terms of percentage of re-planning. Therefore, our second experiment evaluates the selected methods online while training the dynamics to assess the savings in overall training time and its effects on performance. In this experiment, we select the best performing methods from the previous experiment to test them while training the dynamics model from scratch.We evaluate the methods on agents in the MuJoCoTodorov etal. (2012) physics engine using two workstations with a last generation GPU. To establish a valid comparison with the baseline PETs Chua etal. (2018) (denoted as NS1), we use four environments with corresponding task length ($TaskH$) and trajectory horizon ($H$).
We use the following environments: Cartpole(CP): $S\in\mathbb{R}^{4},A\in\mathbb{R}^{1}$, $TaskH$ 200, $H$ 25. Reacher (RE): $S\in\mathbb{R}^{17},A\in\mathbb{R}^{7}$, $TaskH$ 150, $H$ 25. Pusher (PU): $S\in\mathbb{R}^{20},A\in\mathbb{R}^{7}$, $TaskH$ 150, $H$ 25. HalfCheetah (HC): $S\in\mathbb{R}^{18},A\in\mathbb{R}^{6}$, $TaskH$ 1000, $H$ 30.This means that each iteration will run for $TaskH$, task horizon, steps, and that imagined trajectories include $H$ trajectory horizon steps.$S\in\mathbb{R}^{i},A\in\mathbb{R}^{j}$ refers to the dimensions of the environment state consisting in a vector of $i$ components and the action consisting in a vector of $j$ components. We assess performance in terms of reward per episode and evaluate wall time and avoided re-planing. All experiments use random seeds and randomize initial conditions per task.