RL — Model-based Reinforcement Learning
In reinforcement learning RL, we maximize the rewards for our actions. By simply looking at the equation below, rewards depend on the…
Category: Machine Learning
Author: Jonathan Hui

RL — Model-based Reinforcement Learning
Photo by Jonny Caspari
In reinforcement learning RL, we maximize the rewards for our actions. By simply looking at the equation below, rewards depend on the policy and the system dynamics (model).

In Model-free RL, we ignore the model. To accomplish this, we depend on sampling and observation heavily so we don’t need to know the inner working of the system. In Model-based RL, if we can define a cost function ourselves, we can calculate the optimal actions using the model directly. So basically, RL can be roughly divided into Model-free and Model-based methods. In this article, we will discuss how to establish a model and use it to make the best decisions.
Terms
Model-based RL has a strong influence from the control theory and often explained in terms from different disciplines. To get the concepts when explained in different ways, let’s go through some of the terms that refer to similar ideas.
In reinforcement learning, we find an optimal policy to decide actions. In control theory, we optimize a controller.

Control is just another term for action in RL. An action is often written as a or u with states as s or x. The controller uses the model (the system dynamics) to decide the optimal controls for a trajectory expressed as a sequence of states and controls.

In model-based RL, we optimize the trajectory for the least cost instead of the maximum rewards.
Model-free RL v.s. Model-based RL
As mentioned before, Model-free RL ignores the model and care less about the inner working. We fall back to sampling to estimate rewards.

We use Policy Gradients, Value Learning or other Model-free RL to find a policy that takes the best actions for maximum rewards.

On the contrary, Model-based RL focuses on the model.

With a cost function, we find an optimal trajectory with the lowest cost.


Known models
In many games, like GO, the rule of the game is the model.
AlphaGO
Sometimes, it is the law of Physics which can be modeled easily. We even build simulators for it.
Source: Vehicle dynamics model & Kinematic Bicycle Model
Mathematically, the model predicts the next state.

Besides using an explicit model with rules or equations, we can represent it with a generic model, like the Gaussian Process, Gaussian Mixture Model (GMM) or Deep networks. To fit those models, we run a controller to collect sample trajectories and train them with supervised learning.
Motivation
Model-based RL has a strong advantage of being sample efficient. For example, if it is appropriate to approximate the dynamics as locally linear, it will take much fewer samples to learn the model. Once the model and the cost function is known, we can plan the optimal controls without further sampling. As shown below, On-policy Policy Gradient methods can take 10M training iterations while Model-based RL is in the range of hundreds. To train a physical robot for a simple task, a Model-based method may take about 20 minutes while a Policy Gradient method may take weeks. However, such an advantage diminishes if the sampling cost is low. In particular. the optimization methods in the Model-based methods are more complex than Model-free methods in a few orders of magnitude. If sampling can be done in a computer simulation, Model-free methods could finish faster. Also, to simplify the computation, Model-based methods have more assumptions and approximations and therefore, may limit themselves to specific types of tasks.

Learn the model
In Model-based RL, the model may be known or can be learned. We run a base policy, like a random or any educated policy, and observe the trajectory. Then, we fit a model using the sampled data.
Source
In step 2, we use supervised learning to train a model to minimize the least square error from the sampled data. In step 3, we use any trajectory optimization method, like iLQR, to calculate the optimal trajectory using the model and a cost function. The cost function can measure how far we are from the target location and the amount of effort spent.
Learn the model iteratively
However, the last solution is vulnerable to drifting. Tiny errors accumulate fast along the trajectory. The search space is too big for any base policy to cover fully. We may land in areas where the model has not been learned yet. Without a proper model around these areas, we cannot plan the optimal controls.

To address that, instead of learning the model at the beginning, we continue to sample and fit the model as we move along the path.
Source
So we repeat step 2 and step 4 and continue collecting samples and fitting the model around the searched space.
MPC (Model Predictive Control)
Nevertheless, the previous method executes all planned actions before fitting the model again. We may be off-course too far already.

In MPC, we optimize the whole trajectory but we take the first action only. We observe and replan again. The replan gives us a chance to take corrective action after observed the current state again. For a stochastic model, this is particularly helpful.
Source
Backpropagate to policy
What is the difference between a controller and a policy? Are they both serve the same purpose of selecting actions based on states?
The controls produced by a controller is calculated using a model and a cost function through the trajectory optimization methods like iLQR.

However, we can also model a policy directly using a deep network or a Gaussian Process. For example, we can train the policy by backpropagating the cost function and the model.

Source
PILCO
Here is the PILCO algorithm of training a policy directly through backpropagation.
Simplified from source
However, consecutive states in a trajectory are highly correlated. Backpropagation functions with correlated values often lead to vanishing or exploding gradients. So its promise is limited.

Global Model
What kind of models can we use to represent the system dynamics?
Gaussian Process
One possibility is the Gaussian Process. The intuition is simple. If two inputs are similar, their output should be similar too. For two data points, if one is closer to a known training data, its prediction is more certain.

Say, we sampled two data points x1 and x2 with observed values f(x1)=150 and f(x2)=200 respectively. The plot above shows the possible values for f(x) within one standard deviation. The output of other data points, like f1 and f2, can be modeled as a Gaussian Distribution in the following form.


where 175 is the mean of f. For details on how to calculate K, please refer to here. As another example, the right figure below collects 5 data points and plot the mean prediction with the shaded area represents values within one standard deviation.
Source: Wikipedia
Here is the algorithm of the policy search using a Gaussian Process (GP) model:
Source
Gaussian Mixture Model GMM
Another possibility is the GMM. Gaussian Mixture Model is a mixture of K Gaussian distributions. We assume the model has k most likely outcomes and we weight them according:


To identify those k Gaussian distributions, we use Expectation Maximization (EM) to cluster the sample data into k clusters (modes) each represented by a mean and a variance.

Deep network
Finally, we can use a deep network.
Global model
Previously, we model the dynamics with a global model. If the dynamic is complex, we need a more expressive model like the deep network. But it requires a lot of samples to train it. If we land in space which is not yet trained properly, the model can be erroneous and lead us to bad actions that destroy the progress of the training.
Local Model
Alternatively, we can adopt an on-demand approach in developing models locally when we need it.
Source
The local model is Gaussian distributed with linear dynamics.

The parameters for the linear dynamics is computed as:

Controller
Modified from source
We run the controller p on the robot and observe the trajectory. With the collected samples, we can fit the model locally by estimating the derivative using samples.

However, the potential error increases as we move away from the trajectory where the local model is developed. Hence, we add a constraint to make sure the trajectory optimization is done within a trust region. Outside this region, we will not trust the optimization result.

This trust region is measured by the KL-divergence which measures the difference between the new controller and the old controller. If the trajectory is too different from the sampled one, we may get too aggressive and the calculated value can be too far away from the real value.
Image source
So what controller should we use? It will be too boring if the controller explores the same action over and over again for the same state during training. The repeating samples provide no new information to make the model better. Hence, we use a linear Gaussian controller to explore actions better. (linear dynamic with a Gaussian distributed output actions)

and Q is the cost function:

Σ is large when the cost Q is small. So we allow the exploration to go off-course more if the estimated cost is low. So how can we solve:

Before solving that, we need to check out some of the properties of the KL-divergence between the new and the old controller. This will involve some maths but it should be easy to follow or feel free to browse through them quickly.
KL-divergence
We want to establish:

First, we expand the distribution of the trajectory using the new and the old controller:

(Note: both trajectories have same initial and final state.)
Let’s compute the log term below first which is part of the KL-divergence definition:
Source
KL-divergence is defined as:

Therefore,
Source
Dual Gradient Descent
Next, we need optimization methods that work with constraints below:

DGD is our choice. Let’s have a quick summary here. For those want more details, you can read DGD later. First, we need to find its Lagrangian 𝓛 and the Lagrange dual function g which is defined as:

Then the optimal x* of our objective can be solved by:
Source
Intuitively, we start with some random value of λ and find the corresponding optimal values of 𝓛 using an optimization method. In our context, it is often a trajectory optimization method like LQR. g is actually the lower bound of our objective. So we find the optimal value of g by updating λ iteratively using the gradient ascent and the corresponding optimal 𝓛.
Model-based RL with local model
To solve our objective, we are going to reuse LQR. Our objective is:

LQR is quite complicated but for this context, you just need to know LQR is a trajectory optimization method. You can read LQR later if you want more.
The Lagrangian for our objective is:

We divide the objective by λ and create a new surrogate cost function

Here is the objective with the Lagrangian:

As explained in a previous article, LQR solves:

With a Linear Gaussian controller,

LQR minimizes

This looks almost the same as our objective but using the original cost function. Hence, we can reuse LQR exactly but with a new surrogate cost function.
Optimize trajectory with a local model
This is the final algorithm:
Source
Planning
We can also use planning to generate simulated experience and use them to fit the value functions or the policy better.
Source
The difference between learning and planning is one from real experience generated by the environment and one from simulated experience by a model. We use the real experience to fit the value function. We build a model of the transition and sample experience from it. Later, we can fit the value function again with the sampled experience. This improves sample efficiency because sample data are reused and it produces a more accurate value for V.
Model-based RL using sample-based planning (source)
Dyna-Q Algorithm
Here is the Dyna-Q algorithm which used the sampled data and the model to fit the Q-value.
Source
Thoughts
Instead of fitting a policy or a value function, we develop a model to understand the inner workings better. Knowing how things work, we end up with fewer required samples which is the key selling point when the physical simulation is expensive. Here, we discuss the basic of the model learning. But I wish it is that simple. Model learning can be very hard and not generalize well for other tasks. Learning directly for raw images is difficult when information is all tangled. Can we deploy the solution that is economically feasible? Stay tuned for these questions for more in-depth Model-based RL methods in later articles.
Credit and references
UC Berkeley Reinforcement Learning Class
UCL Course on RL