Libraries¶
For visualizing our agent playing, we'll need to load the following library.
In [18]:
!pip3 install imageio==2.4.1
!pip3 install gym[classic_control]
Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/
Requirement already satisfied: imageio==2.4.1 in /usr/local/lib/python3.7/dist-packages (2.4.1)
Requirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from imageio==2.4.1) (1.21.6)
Requirement already satisfied: pillow in /usr/local/lib/python3.7/dist-packages (from imageio==2.4.1) (7.1.2)
Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/
Requirement already satisfied: gym[classic_control] in /usr/local/lib/python3.7/dist-packages (0.25.2)
Requirement already satisfied: gym-notices>=0.0.4 in /usr/local/lib/python3.7/dist-packages (from gym[classic_control]) (0.0.8)
Requirement already satisfied: numpy>=1.18.0 in /usr/local/lib/python3.7/dist-packages (from gym[classic_control]) (1.21.6)
Requirement already satisfied: importlib-metadata>=4.8.0 in /usr/local/lib/python3.7/dist-packages (from gym[classic_control]) (4.13.0)
Requirement already satisfied: cloudpickle>=1.2.0 in /usr/local/lib/python3.7/dist-packages (from gym[classic_control]) (1.5.0)
Collecting pygame==2.1.0
Downloading pygame-2.1.0-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (18.3 MB)
|████████████████████████████████| 18.3 MB 96 kB/s
Requirement already satisfied: typing-extensions>=3.6.4 in /usr/local/lib/python3.7/dist-packages (from importlib-metadata>=4.8.0->gym[classic_control]) (4.1.1)
Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.7/dist-packages (from importlib-metadata>=4.8.0->gym[classic_control]) (3.10.0)
Installing collected packages: pygame
Successfully installed pygame-2.1.0
In [2]:
# Libraries
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import torch.distributions as distributions
import matplotlib.pyplot as plt
import numpy as np
import gym
import moviepy.video.io.ImageSequenceClip
from moviepy.editor import *
Imageio: 'ffmpeg-linux64-v3.3.1' was not found on your computer; downloading it now. Try 1. Download from https://github.com/imageio/imageio-binaries/raw/master/ffmpeg/ffmpeg-linux64-v3.3.1 (43.8 MB) Downloading: 8192/45929032 bytes (0.0%)1327104/45929032 bytes (2.9%)3727360/45929032 bytes (8.1%)5840896/45929032 bytes (12.7%)7225344/45929032 bytes (15.7%)9052160/45929032 bytes (19.7%)10887168/45929032 bytes (23.7%)13041664/45929032 bytes (28.4%)15097856/45929032 bytes (32.9%)17514496/45929032 bytes (38.1%)19800064/45929032 bytes (43.1%)23150592/45929032 bytes (50.4%)26329088/45929032 bytes (57.3%)29859840/45929032 bytes (65.0%)32899072/45929032 bytes (71.6%)34684928/45929032 bytes (75.5%)36872192/45929032 bytes (80.3%)38600704/45929032 bytes (84.0%)41852928/45929032 bytes (91.1%)45424640/45929032 bytes (98.9%)45929032/45929032 bytes (100.0%) Done File saved as /root/.imageio/ffmpeg/ffmpeg-linux64-v3.3.1.
Check out a list of classic games here. We'll be experimenting with Cart Pole where the object is to keep a pole on top of a cart without falling.
In [3]:
env = gym.make('CartPole-v1')
WARNING:py.warnings:/usr/local/lib/python3.7/dist-packages/gym/core.py:318: DeprecationWarning: WARN: Initializing wrapper in old step API which returns one bool instead of two. It is recommended to set `new_step_api=True` to use new step API. This will be the default behaviour in future. "Initializing wrapper in old step API which returns one bool instead of two. It is recommended to set `new_step_api=True` to use new step API. This will be the default behaviour in future." WARNING:py.warnings:/usr/local/lib/python3.7/dist-packages/gym/wrappers/step_api_compatibility.py:40: DeprecationWarning: WARN: Initializing environment in old step API which returns one bool instead of two. It is recommended to set `new_step_api=True` to use new step API. This will be the default behaviour in future. "Initializing environment in old step API which returns one bool instead of two. It is recommended to set `new_step_api=True` to use new step API. This will be the default behaviour in future."
The observation space consits of four variables:
Cart position (on a line)
Cart velocity
Pole angle
Pole angular velocity
In [4]:
print(env.observation_space)
print(env.observation_space.shape)
Box([-4.8000002e+00 -3.4028235e+38 -4.1887903e-01 -3.4028235e+38], [4.8000002e+00 3.4028235e+38 4.1887903e-01 3.4028235e+38], (4,), float32) (4,)
We can take two actions: either push the cart to the left or to the right.
In [5]:
print(env.action_space)
Discrete(2)
The game is actually pretty hard! You can try your hand at it here. Now let's train a policy to beat it. We start by loading two different copies of the environment; one for training and one for testing.
In [6]:
train_env = gym.make('CartPole-v1')
test_env = gym.make('CartPole-v1')
We'll train a neural network that is actually incredibly simple: only two layers with dropout and ReLU.
In [7]:
class PolicyNet(nn.Module):
def __init__(self, input_dim, hidden_dim, output_dim, dropout=.5):
super().__init__()
self.net = nn.Sequential(
nn.Linear(input_dim, hidden_dim),
nn.Dropout(dropout),
nn.ReLU(),
nn.Linear(hidden_dim, output_dim)
)
def forward(self, x):
return self.net(x)
As we've done many times before, we initialize our architecture and optimizer. However, the loss will be slightly more tricky to compute than normal.
In [8]:
input_dim = train_env.observation_space.shape[0]
print('input dim:', input_dim)
output_dim = train_env.action_space.n
print('output dim:', output_dim)
hidden_dim = 128
policy = PolicyNet(input_dim, hidden_dim, output_dim)
print(policy)
optimizer = optim.Adam(policy.parameters(), lr=0.01)
input dim: 4
output dim: 2
PolicyNet(
(net): Sequential(
(0): Linear(in_features=4, out_features=128, bias=True)
(1): Dropout(p=0.5, inplace=False)
(2): ReLU()
(3): Linear(in_features=128, out_features=2, bias=True)
)
)
We will write a function for calculating a sequence of returns. With these in hand, we can calculate the empirical average over many trajectories as described in the lecture notes.
In [9]:
def calculate_returns(rewards, discount_factor, normalize=True):
returns = []
R = 0
for r in reversed(rewards):
R = r + R * discount_factor
returns.insert(0,R)
returns = torch.tensor(returns)
# Remember we normalize to differentiate between
# "good" and "bad" actions
if normalize:
returns = (returns - returns.mean()) / returns.std()
return returns
Once we have returns and log probabilities of the actions we take, we will compute the loss and optimize.
In [10]:
def update_policy(returns, log_prob_actions, optimizer):
returns = returns.detach()
# Because we minimize (by convention) and we actually
# want to maximize returns, we negate the expected sum
loss = - (returns * log_prob_actions).sum()
optimizer.zero_grad()
loss.backward()
optimizer.step()
return loss.item()
Now we write the training and evaluation functions. Notice how training in RL is more involved than our standard applications.
In [11]:
def train(env, policy, optimizer, discount_factor=.99):
policy.train()
log_prob_actions = []
rewards = []
done = False
episode_reward = 0
state = env.reset()
while not done:
state = torch.FloatTensor(state).unsqueeze(0)
action_pred = policy(state)
action_prob = F.softmax(action_pred, dim = -1)
dist = distributions.Categorical(action_prob)
action = dist.sample() # we sample (rather than take the best)
log_prob_action = dist.log_prob(action)
state, reward, done, _ = env.step(action.item())
log_prob_actions.append(log_prob_action)
rewards.append(reward)
episode_reward += reward
log_prob_actions = torch.cat(log_prob_actions)
returns = calculate_returns(rewards, discount_factor)
loss = update_policy(returns, log_prob_actions, optimizer)
return loss, episode_reward
In [20]:
def evaluate(env, policy, return_frames=False):
policy.eval()
done = False
episode_reward = 0
state = env.reset()
frames = []
while not done:
if return_frames:
frames += [env.render(mode='rgb_array')]
state = torch.FloatTensor(state).unsqueeze(0)
action_pred = policy(state)
action_prob = F.softmax(action_pred, dim=-1)
action = torch.argmax(action_prob, dim=-1) # we take the best
state, reward, done, _ = env.step(action.item())
episode_reward += reward
if return_frames:
return episode_reward, frames
return episode_reward
Now we're ready to train. Notice the very jagged nature of the loss history. This is because RL is different from our regular training: our policy is stochastic and the reward is more "sparse".
In [13]:
max_episodes = 500
num_trials = 25
reward_threshold = 475
print_every = 10
train_rewards = []
test_rewards = []
for episode in range(1, max_episodes+1):
loss, train_reward = train(train_env, policy, optimizer)
test_reward = evaluate(test_env, policy)
train_rewards += [train_reward]
test_rewards += [test_reward]
mean_train_rewards = np.mean(train_rewards[-num_trials:])
mean_test_rewards = np.mean(test_rewards[-num_trials:])
if episode % print_every == 0:
print(f'| Episode: {episode:3} | Mean Train Rewards: {mean_train_rewards:5.1f} | Mean Test Rewards: {mean_test_rewards:5.1f} |')
if mean_test_rewards >= reward_threshold:
print(f'Reached reward threshold in {episode} episodes')
break
| Episode: 10 | Mean Train Rewards: 21.4 | Mean Test Rewards: 9.4 | | Episode: 20 | Mean Train Rewards: 15.9 | Mean Test Rewards: 9.3 | | Episode: 30 | Mean Train Rewards: 11.2 | Mean Test Rewards: 9.2 | | Episode: 40 | Mean Train Rewards: 12.0 | Mean Test Rewards: 9.3 | | Episode: 50 | Mean Train Rewards: 20.0 | Mean Test Rewards: 23.7 | | Episode: 60 | Mean Train Rewards: 36.9 | Mean Test Rewards: 48.7 | | Episode: 70 | Mean Train Rewards: 53.6 | Mean Test Rewards: 93.6 | | Episode: 80 | Mean Train Rewards: 55.8 | Mean Test Rewards: 133.5 | | Episode: 90 | Mean Train Rewards: 78.9 | Mean Test Rewards: 144.8 | | Episode: 100 | Mean Train Rewards: 95.2 | Mean Test Rewards: 181.8 | | Episode: 110 | Mean Train Rewards: 104.7 | Mean Test Rewards: 191.0 | | Episode: 120 | Mean Train Rewards: 149.4 | Mean Test Rewards: 298.0 | | Episode: 130 | Mean Train Rewards: 233.9 | Mean Test Rewards: 404.8 | Reached reward threshold in 136 episodes
In [14]:
plt.figure(figsize=(8,8))
plt.plot(test_rewards, label='Test Reward')
plt.plot(train_rewards, label='Train Reward')
plt.xlabel('Episode', fontsize=20)
plt.ylabel('Reward', fontsize=20)
plt.hlines(reward_threshold, 0, len(test_rewards), color='r')
plt.legend(loc='lower right')
plt.grid()
In [21]:
_, frames = evaluate(test_env, policy, return_frames = True)
WARNING:py.warnings:/usr/local/lib/python3.7/dist-packages/gym/core.py:44: DeprecationWarning: WARN: The argument mode in render method is deprecated; use render_mode during environment initialization instead.
See here for more information: https://www.gymlibrary.ml/content/api/
"The argument mode in render method is deprecated; "
In [22]:
clip = moviepy.video.io.ImageSequenceClip.ImageSequenceClip(frames, fps=50)
ipython_display(clip)
100%|██████████| 422/422 [00:02<00:00, 140.77it/s]
Out[22]: