REINFORCE with a Baseline

Train a policy network with the REINFORCE policy gradient.

The environment is a corridor. The agent starts in an inner cell. A cue in the observation says whether the goal is at the left end or the right end. Reaching the goal gives +1, the wrong end gives -1, and each other step costs 0.01. The policy must read the cue, so it cannot learn one fixed direction.

REINFORCE raises the log probability of each action in proportion to how good the return after it was. Write three functions:

G_t  = r_t + gamma * G_(t+1)                  # returns to go, G_T = 0
adv  = G - G.mean()                           # a constant baseline
loss = -mean( log pi(a_t | s_t) * adv_t )

Use torch.distributions.Categorical(logits=...) for the log probabilities and for sampling.

Function signature

def returns_to_go(rewards, gamma):
def reinforce_loss(logits, actions, returns):
def train_policy(env):

The environment has env.obs_dim, env.n_actions, env.reset(), which returns a float observation of shape (obs_dim,), and env.step(action), which returns (obs, reward, done). Every episode ends within 20 steps.

Example

Input

returns_to_go([0.0, 0.0, 1.0], gamma=0.9)
reinforce_loss(logits=torch.tensor([[0.0, 0.0], [2.0, 0.0]]),
               actions=torch.tensor([1, 0]),
               returns=torch.tensor([1.0, 0.0]))

Output

tensor([0.8100, 0.9000, 1.0000])
tensor(0.1416)

Reasoning

The reward of 1 at the last step is worth 0.9 one step earlier and 0.81 two steps earlier. For the loss, the mean return is 0.5, so the advantages are 0.5 and -0.5. The log probabilities of the chosen actions are -0.6931 and -0.1269. The loss is -(-0.6931 * 0.5 + -0.1269 * -0.5) / 2 = 0.1416.

Constraints and notes

• returns_to_go takes a list of floats and returns a float tensor of the same length.
• reinforce_loss gets logits of shape (T, n_actions), int64 actions of shape (T,) and float returns of shape (T,). The returns are constants: no gradient flows into them.
• train_policy returns an nn.Module that maps an observation to logits. With greedy actions (argmax), it must reach the goal from at least 90% of all start cells and cues. It must use env.obs_dim, because the tests also use a longer corridor.
• Keep training under about 10 seconds on a CPU.