In this post, we’ll get into the weeds with some of the fundamentals of reinforcement learning. Hopefully, this will serve as a thorough overview of the basics for someone who is curious and doesn’t want to invest a significant amount of time into learning all of the math and theory behind the basics of reinforcement learning.
In reinforcement learning (RL), we want to solve a Markov Decision Process (MDP) by figuring out how to take actions that maximize the reward received. The actor in an MDP is called an agent. In an MDP, actions taken influence both current and future rewards received and actions influence future states the agent finds itself in. Because of this, solving an MDP requires that an agent is able to handle delayed reward and must be able to balance a trade-off between obtaining reward immediately and delaying reward collection. Below we have a diagram of the classic MDP formulation of agent-environment interaction.
At initialization, the environment outputs some state . The agent observes this state and in response takes an action, . This action is then applied to the environment, the environment is stepped forward in response to the action taken, and it yields a new state, and reward signal to the agent. This loop continues until the episode (period of agent-environment interaction) terminates.
Some examples of a state might be a frame of an Atari game or the current layout of pieces in a board game such as chess. The reward is a scalar signal calculated following a reward function, and the action can be something like which piece to move where on a chess board, or which direction to go in an Atari game.
The model of the environment consists of a state transition function and a reward function. Here, the transition function will be discussed and later we’ll talk about reward functions. In finite MDPs (the kind of MDP we are concerned with), the number of states, actions, and rewards are all finite values. Because of this property, they also all have well-defined probability distributions. The distribution over the next state and next reward depends only on the previous state and action. This is called the Markov property. For the random variables and , there is a probability at time that they’ll have particular values. This probability is conditional on the previous state and action, and can be written out like this:
This 4-argument function, , fully captures the dynamics of the MDP and tells us, formally, that our new state () and reward are random variables whose distribution is conditioned on the previous state and action. A complete model of an MDP can be used to calculate anything we want about the environment. State-transition probabilities can be found, expected rewards for state-action pairs, and even expected rewards for state-action-next state triplets.
An episodic MDP (we can also refer to an MDP as a task) is one with a clear stopping, or terminating, state. An example of this might be something like a game of chess, where the natural stopping state is when one player wins and the other loses.
A continuing task is one that doesn’t have a clear stopping state, and can be allowed to continue indefinitely. Consider trying to train a robot to walk; there is not necessarily a clear stopping point. In practice, when we are training a simulated robot to walk, we terminate the episode when it falls over. Over time, however, the agent learns to walk successfully. Once it is successfully walking, there isn’t an easily defined stopping point. Since we don’t want to let a simulation run forever, we typically enforce some maximum number of interactions allowed per episode. Once this number is reached, the episode terminates.
We use reward to define the goal in the problem we’d like the RL agent to solve. At every step , the agent receives a single, scalar reward signal from the environment. The agent’s aim is to maximize the reward it receives over all future actions. This is called the return.
Return is the cumulative sum of all reward earned in an episode. Let’s denote return with . Then, the return of an episode is defined with:
is the reward earned at the final step of an episode. means “defined by”. We can rewrite the above expression more concisely:
The reward formulation above doesn’t work for continuing tasks, because as the number of time steps goes to infinity, so does the reward. We need a formulation of reward that will converge to a finite value as the number of time steps goes to infinity. To do this, we can assign a discount factor to future rewards. We define a discount rate parameter, (gamma) and set . We incorporate into our reward formulation like so:
Rewrite for brevity:
Since , this is a telescoping sum. As , the reward approaches zero. This is what we want, because under this formulation, the expected future reward converges to a finite value, instead of going to infinity.
In practice, we normally discount future rewards with roughly between 0.95 and 0.99, even in episodic tasks. Intuitively, this is because reward now is normally better than reward later. There are cases where a value outside of that range will yield the best performance, but most papers and libraries seem to use a in the above range.
The MDP formulation is a clear way to frame the problem of learning from interaction to achieve a certain goal. Interaction between an agent and its environment occur in discrete time steps, i.e. . As we know, at every time step our agent takes an action and receives a new observation and reward . Writing this out directly, the interaction between agent and MDP produces a trajectory that progresses like this:
We can denote interaction trajectories by , the Greek letter Tau.
Interesting note: Theoretically, the current state and reward in an MDP should only depend on the previous state and action. In practice, however, this condition (called the Markov property) is violated very regularly. When an RL agent is learning how to control a robot to move forward, the current position of the robot is not only dependent on the previous state and action, but all of the states and actions before it. Somewhat surprisingly, deep RL algorithms are able to achieve excellent performance on these domains despite some problems technically being non-Markovian.
A policy is a function that maps from states to actions (or to a probability distribution over actions) and is the decision-making part of the agent. A value function estimates how good a particular state, or state-action pair, is. “Good” is defined in terms of expected future reward following that state or state-action pair.
The policy is a mapping from states to actions (or to a probability distribution over actions), and can be written like so:
Where represents the policy, represents the current state, and is the chosen action to take while in state .
Policies can be stochastic or deterministic. In the case of a stochastic policy, the output would be a probability distribution over actions. In a deterministic policy, the output is directly what action to take. These can be written mathematically:
Let’s break this down. In the stochastic policy, is telling us that the output of the policy is conditioned on the state . says that the probability of the action being equal to depends on equaling . The deterministic policy simply tells us that the policy takes in state and maps to an action .
The value of a state is determined by how much reward is expected to follow it. If there’s lots of reward expected after state , then that must be a good state. But, if there is very little reward expected to follow state , then that’s a bad state to be in. The state-value function is a function that, given a state, outputs an estimate of the expected future return following that state. An action-value function will take in a state and an action and will output an estimate of the expected future return following that state-action pair. Value functions are defined with respect to policies. The value of a state under policy is written with ; this is the state-value function. Mathematically:
Recall that is the return and is the formula for discounted return. The action-value function is written a bit differently and makes a slightly different assumption than the state-value function. Whereas the state-value function assumes that the policy starts in state and afterward takes all actions according to , the action-value function assumes that the policy starts in state , takes an action (which may or may not be on policy), and thereafter acts following .
The main distinction to note is that is dependent only on the current state, while is dependent on both the current state and action.
The advantage function is found by subtracting the state-value from the state-action value, to get the action value (advantage function). This is often used in practice when training deep RL algorithms.
See GAE-lambda advantage estimation for state-of-the-art with regards to advantage functions.
Now that we know about policies and value functions, it’s time to see what optimal policies and value functions are.
An optimal policy always takes the action in a state that will yield the maximum expected future reward. This can be written like so:
where is the optimal policy. The means “take the action on policy that yields the highest future return”. Finding the optimal policy is the central problem in RL, because once the optimal policy is found, the agent can then always take the best action in any state. See here for more.
The optimal state-value function (written with ) is the function that always gives you the expected return when starting in state and afterwards acting according to the optimal policy.
The optimal action-value function () always gives the expected future return if you start in state , take action (that may or may not be on-policy) and then afterwards act following the optimal policy.
When the optimal policy is in state , it chooses the action which maximizes the expected return when starting from the current state. Therefore, when we have the optimal function, the optimal policy is easily found by:
Where is the optimal action, and means “take the action that maximizes ”
In RL, we need to trade off between exploiting what an agent has already learned and exploring the environment and actions to find potentially better actions. When we deal with a greedy policy (one that always chooses the action with the highest expected return), we often must sometimes force such a policy to explore non-greedy actions during training. This is done by picking a random action instead of the on-policy action with some probability .
Bellman equations demonstrate a relationship between the value of a current state and the values of following states. It looks from a current state into the future, averages all future states and the possible actions in those states, and weights each state-action pair by the probability that it will occur. The Bellman equations are a set of equations that break the value function down into the reward in the current state plus the discounted future values.
The set of equations for state-value is this:
And for action-value:
The means “action is sampled from policy ”.
If we don’t care about computing the expected future reward when following a policy, then we can use the optimal Bellman equations:
The difference between the on-policy and optimal Bellman equations is whether or not the operation is present. When the is there, it means that for the agent to act optimally, it has to take the action that has the maximum expected return (aka highest value).
An observation is the state, or portion of the state, as it is observed by the agent. In fully-observable MDPs, the state and the observation can be identical, but sometimes are not (for example, in the Atari DQN paper, the observations to the agent are the most recent few frames of the game, concatenated together), and in partially-observable MDPs, the agent cannot see the entirety of the environment’s state. Therefore, the distinction between the state of an environment and the observations an agent receives is an important one.
Observation spaces can be discretely or continuously valued. An example of a discrete observation space is the layout of a tic-tac-toe board, where empty space is represented by a zero, X pieces are represented by a 1, and O pieces are represented by a 2. A continuous observation space example might be a vector of joint torques and velocities in a robot, like the one in the gif below.
In this gif, the goal is to move the ant robot forward. The observations are a 28 dimensional vector of joint torques and velocities. The agent must learn to use these observations to take actions to achieve the goal.
Similarly to observation spaces, actions can be continuously or discretely valued. A simple example of a discrete action is one in Atari, where you move a joystick in one of four distinct directions to control where your character moves. Continuous actions can be more complex. For example, in the gif above, the actions are a real valued vector of continuous numbers representing amounts of force to apply to the robot’s joints. Continuous actions are always a vector of values.
Reward functions yield a scalar signal at every step of the environment telling the agent how good the previous action was. The agent uses these signals to learn to take good actions in every state. To the agent, goodness is defined by how much reward was earned.
Reward functions can either be simple or complex. For a simple example, in the classic gridworld environment (see diagram below), the agent starts in one corner of a grid and must navigate an end state in the other corner of the grid. The reward at every step, no matter the action, is . This incentivizes the agent to navigate from start to end as quickly as possible.
In the above image, the greyed out squares are the starting and ending points. Either one can be the start or end point. A more complicated example of a reward function would be the one for the ant robot above. The reward function for that agent is as follows:
Let’s unpack this. tells us that the reward is a function of the last state and action. and are the position of the robot before and after the action, respectively. is the change in time before the last action and current action. is the action vector, in this case the actions are a vector because we’re applying forces to 8 different joints on the robot, so each entry in the vector tells us how much force to apply to the corresponding joint. is a vector of external forces on the body of the robot. tells us to clip the vector to have all values fall within the range .
Here, we’ll briefly touch on some classical methods in reinforcement learning. These methods aren’t in use on cutting-edge problems today, but do lay an important theoretical foundation for modern algorithms.
Dynamic programming (DP) algorithms require a perfect model of the environment. Practically, this is often unrealistic to expect and therefore DP algorithms often are not actually used in RL. However, they are still theoretically important. We can use DP to find optimal value functions, and from there, optimal policies.
We can compute the state-value function for a policy by using policy evaluation. Mathematically:
where means sum over actions, and other symbols should be known from earlier sections.
Policy improvement finds a new and improved policy by greedily selecting the action with highest value in each state.
Once we’ve used to improve and get , we can again perform a policy evaluation step (compute the state value function for ) and a policy improvement step (compute the improved poicy ). By continuing this cycle we can get consistently improving policies and value functions. The trajectory of policy iteration looks like this ( is for policy evaluation and is for policy improvement):
In the previous section, we discussed policy iteration for deterministic policies. The math and theory described there extends to stochastic policies too. Monte Carlo (MC) methods do not require a model of the environment and instead can learn entirely from experience. The core idea of MC methods is to use average observed return as an approximation for the value of a state. To get an empirical return, MC methods need complete episodes and the episodes have to end. We can estimate under first visit or every visit MC. First visit MC estimates the value of as an average of the return after the first visit to , while every visit MC estimates the value by averaging the return of every time the state is visited. Both first and every visit MC converge to the true value of the state as the number of visits goes to infinity.
Image from Lilian Weng’s blog. Showing that learning an optimal policy via MC is done by following a similar idea to policy iteration.
Similarly to policy iteration, we improve the policy greedily with respect to the current value function.
Then, we use the updated policy to generate a new episode to train on. Finally, we estimate the function using information gathered from the episode.
Temporal Difference (TD) learning combines ideas from Monte Carlo and Dynamic Programming methods. Like MC methods, TD doesn’t require a model of the environment and instead learns only from experience. TD methods update value estimates based partially on other estimates that have already been learned and they can learn from incomplete episodes.
TD methods use existing estimates to update values instead of only relying on empirical and complete returns like MC methods do. This is known as bootstrapping.
Similarly to MC methods, TD methods use experience to estimate values. Both follow a policy and collect experience over episodes. Both TD and MC methods update their value estimates for every nonterminal state in the set of gathered experience. A difference is that MC methods wait until the end of an episode, when empirical return is known, to update their value estimates, whereas TD methods can update their estimates with respect to other estimates (they don’t rely on empirical return) during an episode. We can write a simple version of an MC method that works well in nonstationary environments:
It is simple to turn this MC update into a TD update by switching out the empirical return for the current value estimate of the next state :
is a step-size parameter where .
This update can also be written for the action-value function:
Learning an optimal policy using TD learning is called TD control. We’ll next look at two algorithms for TD control.
We again follow the pattern of policy iteration, except now we use TD methods for the evaluation (value estimation) steps. In SARSA, we need to learn a function and then define a greedy (or -greedy) policy with respect to that function. This can be done using the update rule from above:
This update is performed after every nonterminal state. If is terminal, then is 0. This rule uses each element in the tuple: . This tuple also gives the SARSA algorithm its name. SARSA algorithm has these steps:
Q-learning was an early breakthrough in reinforcement learning by Watkins and Dyan in 1989. The algorithms update rule is:
Under this rule, directly approximates the optimal action-value function , independent of the current policy.
The Q-learning algorithm has these steps:
Thank you for sticking with me through this long blog post. I really hope it was worth your while and should you find any errors, please email me. Please feel free to have a discussion or raise questions in the comments. In my next post, I’m planning to write about policy gradient methods and dive into the theory behind them. See you then!