Policy-Gradient and Actor-Critic Methods
그로킹 심층 강화학습 중 11장 내용인 "정책 경사법과 액터-크리틱 학습 방법들"에 대한 내용입니다.
- Different types of Cart Pole environments
- Monte-Carlo REINFORCE
- REINFORCE Agent progression
- Fully-trained REINFORCE Agent
- Monte-Carlo VPG
- VPG Agent progression
- Fully-trained VPG Agent
- Asynchronous Advantage Actor-Critic (A3C)
- A3C Agent progression
- Fully-trained A3C Agent
- Generalized Advantage Estimation (GAE)
- GAE Agent progression
- Fully-trained GAE Agent
- Advantage Actor-Critic (A2C)
- A2C Agent progression
- Fully-trained A2C Agent
Note: 실행을 위해 아래의 패키지들을 설치해주기 바랍니다.
!pip install tqdm numpy scikit-learn pyglet setuptools && \
!pip install gym asciinema pandas tabulate tornado==5.* PyBullet && \
!pip install git+https://github.com/pybox2d/pybox2d#egg=Box2D && \
!pip install git+https://github.com/mimoralea/gym-bandits#egg=gym-bandits && \
!pip install git+https://github.com/mimoralea/gym-walk#egg=gym-walk && \
!pip install git+https://github.com/mimoralea/gym-aima#egg=gym-aima && \
!pip install gym[atari]
!pip install torch torchvision
import warnings ; warnings.filterwarnings('ignore')
import os
os.environ['CUDA_DEVICE_ORDER']='PCI_BUS_ID'
os.environ['CUDA_VISIBLE_DEVICES']=''
os.environ['OMP_NUM_THREADS'] = '1'
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torch.multiprocessing as mp
import threading
import numpy as np
from IPython.display import display
from collections import namedtuple, deque
import matplotlib.pyplot as plt
import matplotlib.pylab as pylab
from itertools import cycle, count
from textwrap import wrap
import matplotlib
import subprocess
import os.path
import tempfile
import random
import base64
import pprint
import glob
import time
import json
import sys
import gym
import io
import os
import gc
import platform
from gym import wrappers
from subprocess import check_output
from IPython.display import HTML
LEAVE_PRINT_EVERY_N_SECS = 30
ERASE_LINE = '\x1b[2K'
EPS = 1e-6
RESULTS_DIR = os.path.join('.', 'gym-results')
SEEDS = (12, 34, 56, 78, 90)
%matplotlib inline
plt.style.use('fivethirtyeight')
params = {
'figure.figsize': (15, 8),
'font.size': 24,
'legend.fontsize': 20,
'axes.titlesize': 28,
'axes.labelsize': 24,
'xtick.labelsize': 20,
'ytick.labelsize': 20
}
pylab.rcParams.update(params)
np.set_printoptions(suppress=True)
torch.cuda.is_available()
def get_make_env_fn(**kargs):
def make_env_fn(env_name, seed=None, render=None, record=False,
unwrapped=False, monitor_mode=None,
inner_wrappers=None, outer_wrappers=None):
mdir = tempfile.mkdtemp()
env = None
if render:
try:
env = gym.make(env_name, render=render)
except:
pass
if env is None:
env = gym.make(env_name)
if seed is not None: env.seed(seed)
env = env.unwrapped if unwrapped else env
if inner_wrappers:
for wrapper in inner_wrappers:
env = wrapper(env)
env = wrappers.Monitor(
env, mdir, force=True,
mode=monitor_mode,
video_callable=lambda e_idx: record) if monitor_mode else env
if outer_wrappers:
for wrapper in outer_wrappers:
env = wrapper(env)
return env
return make_env_fn, kargs
def get_videos_html(env_videos, title, max_n_videos=5):
videos = np.array(env_videos)
if len(videos) == 0:
return
n_videos = max(1, min(max_n_videos, len(videos)))
idxs = np.linspace(0, len(videos) - 1, n_videos).astype(int) if n_videos > 1 else [-1,]
videos = videos[idxs,...]
strm = '<h2>{}</h2>'.format(title)
for video_path, meta_path in videos:
video = io.open(video_path, 'r+b').read()
encoded = base64.b64encode(video)
with open(meta_path) as data_file:
meta = json.load(data_file)
html_tag = """
<h3>{0}</h3>
<video width="960" height="540" controls>
<source src="data:video/mp4;base64,{1}" type="video/mp4" />
</video>"""
strm += html_tag.format('Episode ' + str(meta['episode_id']), encoded.decode('ascii'))
return strm
platform.system()
def get_gif_html(env_videos, title, subtitle_eps=None, max_n_videos=4):
videos = np.array(env_videos)
if len(videos) == 0:
return
n_videos = max(1, min(max_n_videos, len(videos)))
idxs = np.linspace(0, len(videos) - 1, n_videos).astype(int) if n_videos > 1 else [-1,]
videos = videos[idxs,...]
strm = '<h2>{}</h2>'.format(title)
for video_path, meta_path in videos:
basename = os.path.splitext(video_path)[0]
gif_path = basename + '.gif'
if not os.path.exists(gif_path):
if platform.system() == 'Linux':
ps = subprocess.Popen(
('ffmpeg',
'-i', video_path,
'-r', '7',
'-f', 'image2pipe',
'-vcodec', 'ppm',
'-crf', '20',
'-vf', 'scale=512:-1',
'-'),
stdout=subprocess.PIPE,
universal_newlines=True)
output = subprocess.check_output(
('convert',
'-coalesce',
'-delay', '7',
'-loop', '0',
'-fuzz', '2%',
'+dither',
'-deconstruct',
'-layers', 'Optimize',
'-', gif_path),
stdin=ps.stdout)
ps.wait()
else:
ps = subprocess.Popen('ffmpeg -i {} -r 7 -f image2pipe \
-vcodec ppm -crf 20 -vf scale=512:-1 - | \
convert -coalesce -delay 7 -loop 0 -fuzz 2% \
+dither -deconstruct -layers Optimize \
- {}'.format(video_path, gif_path),
stdin=subprocess.PIPE,
shell=True)
ps.wait()
gif = io.open(gif_path, 'r+b').read()
encoded = base64.b64encode(gif)
with open(meta_path) as data_file:
meta = json.load(data_file)
html_tag = """
<h3>{0}</h3>
<img src="data:image/gif;base64,{1}" />"""
prefix = 'Trial ' if subtitle_eps is None else 'Episode '
sufix = str(meta['episode_id'] if subtitle_eps is None \
else subtitle_eps[meta['episode_id']])
strm += html_tag.format(prefix + sufix, encoded.decode('ascii'))
return strm
class DiscountedCartPole(gym.Wrapper):
def __init__(self, env):
gym.Wrapper.__init__(self, env)
def reset(self, **kwargs):
return self.env.reset(**kwargs)
def step(self, a):
o, r, d, _ = self.env.step(a)
(x, x_dot, theta, theta_dot) = o
pole_fell = x < -self.env.unwrapped.x_threshold \
or x > self.env.unwrapped.x_threshold \
or theta < -self.env.unwrapped.theta_threshold_radians \
or theta > self.env.unwrapped.theta_threshold_radians
r = -1 if pole_fell else 0
return o, r, d, _
class MCCartPole(gym.Wrapper):
def __init__(self, env):
gym.Wrapper.__init__(self, env)
def reset(self, **kwargs):
return self.env.reset(**kwargs)
def step(self, a):
o, r, d, _ = self.env.step(a)
(x, x_dot, theta, theta_dot) = o
pole_fell = x < -self.env.unwrapped.x_threshold \
or x > self.env.unwrapped.x_threshold \
or theta < -self.env.unwrapped.theta_threshold_radians \
or theta > self.env.unwrapped.theta_threshold_radians
if d:
if pole_fell:
r = 0 # done, in failure
else:
r = self.env._max_episode_steps # done, but successfully
return o, r, d, _
class FCDAP(nn.Module):
def __init__(self,
input_dim,
output_dim,
hidden_dims=(32,32),
activation_fc=F.relu):
super(FCDAP, self).__init__()
self.activation_fc = activation_fc
self.input_layer = nn.Linear(input_dim, hidden_dims[0])
self.hidden_layers = nn.ModuleList()
for i in range(len(hidden_dims)-1):
hidden_layer = nn.Linear(hidden_dims[i], hidden_dims[i+1])
self.hidden_layers.append(hidden_layer)
self.output_layer = nn.Linear(hidden_dims[-1], output_dim)
def _format(self, state):
x = state
if not isinstance(x, torch.Tensor):
x = torch.tensor(x,
dtype=torch.float32)
x = x.unsqueeze(0)
return x
def forward(self, state):
x = self._format(state)
x = self.activation_fc(self.input_layer(x))
for hidden_layer in self.hidden_layers:
x = self.activation_fc(hidden_layer(x))
return self.output_layer(x)
def full_pass(self, state):
logits = self.forward(state)
dist = torch.distributions.Categorical(logits=logits)
action = dist.sample()
logpa = dist.log_prob(action).unsqueeze(-1)
entropy = dist.entropy().unsqueeze(-1)
is_exploratory = action != np.argmax(logits.detach().numpy())
return action.item(), is_exploratory.item(), logpa, entropy
def select_action(self, state):
logits = self.forward(state)
dist = torch.distributions.Categorical(logits=logits)
action = dist.sample()
return action.item()
def select_greedy_action(self, state):
logits = self.forward(state)
return np.argmax(logits.detach().numpy())
class REINFORCE():
def __init__(self, policy_model_fn, policy_optimizer_fn, policy_optimizer_lr):
self.policy_model_fn = policy_model_fn
self.policy_optimizer_fn = policy_optimizer_fn
self.policy_optimizer_lr = policy_optimizer_lr
def optimize_model(self):
T = len(self.rewards)
discounts = np.logspace(0, T, num=T, base=self.gamma, endpoint=False)
returns = np.array([np.sum(discounts[:T-t] * self.rewards[t:]) for t in range(T)])
discounts = torch.FloatTensor(discounts).unsqueeze(1)
returns = torch.FloatTensor(returns).unsqueeze(1)
self.logpas = torch.cat(self.logpas)
policy_loss = -(discounts * returns * self.logpas).mean()
self.policy_optimizer.zero_grad()
policy_loss.backward()
self.policy_optimizer.step()
def interaction_step(self, state, env):
action, is_exploratory, logpa, _ = self.policy_model.full_pass(state)
new_state, reward, is_terminal, _ = env.step(action)
self.logpas.append(logpa)
self.rewards.append(reward)
self.episode_reward[-1] += reward
self.episode_timestep[-1] += 1
self.episode_exploration[-1] += int(is_exploratory)
return new_state, is_terminal
def train(self, make_env_fn, make_env_kargs, seed, gamma,
max_minutes, max_episodes, goal_mean_100_reward):
training_start, last_debug_time = time.time(), float('-inf')
self.checkpoint_dir = tempfile.mkdtemp()
self.make_env_fn = make_env_fn
self.make_env_kargs = make_env_kargs
self.seed = seed
self.gamma = gamma
env = self.make_env_fn(**self.make_env_kargs, seed=self.seed)
torch.manual_seed(self.seed) ; np.random.seed(self.seed) ; random.seed(self.seed)
nS, nA = env.observation_space.shape[0], env.action_space.n
self.episode_timestep = []
self.episode_reward = []
self.episode_seconds = []
self.episode_exploration = []
self.evaluation_scores = []
self.policy_model = self.policy_model_fn(nS, nA)
self.policy_optimizer = self.policy_optimizer_fn(self.policy_model,
self.policy_optimizer_lr)
result = np.empty((max_episodes, 5))
result[:] = np.nan
training_time = 0
for episode in range(1, max_episodes + 1):
episode_start = time.time()
state, is_terminal = env.reset(), False
self.episode_reward.append(0.0)
self.episode_timestep.append(0.0)
self.episode_exploration.append(0.0)
# collect rollout
self.logpas, self.rewards = [], []
for step in count():
state, is_terminal = self.interaction_step(state, env)
if is_terminal:
gc.collect()
break
self.optimize_model()
# stats
episode_elapsed = time.time() - episode_start
self.episode_seconds.append(episode_elapsed)
training_time += episode_elapsed
evaluation_score, _ = self.evaluate(self.policy_model, env)
self.save_checkpoint(episode-1, self.policy_model)
total_step = int(np.sum(self.episode_timestep))
self.evaluation_scores.append(evaluation_score)
mean_10_reward = np.mean(self.episode_reward[-10:])
std_10_reward = np.std(self.episode_reward[-10:])
mean_100_reward = np.mean(self.episode_reward[-100:])
std_100_reward = np.std(self.episode_reward[-100:])
mean_100_eval_score = np.mean(self.evaluation_scores[-100:])
std_100_eval_score = np.std(self.evaluation_scores[-100:])
lst_100_exp_rat = np.array(
self.episode_exploration[-100:])/np.array(self.episode_timestep[-100:])
mean_100_exp_rat = np.mean(lst_100_exp_rat)
std_100_exp_rat = np.std(lst_100_exp_rat)
wallclock_elapsed = time.time() - training_start
result[episode-1] = total_step, mean_100_reward, \
mean_100_eval_score, training_time, wallclock_elapsed
reached_debug_time = time.time() - last_debug_time >= LEAVE_PRINT_EVERY_N_SECS
reached_max_minutes = wallclock_elapsed >= max_minutes * 60
reached_max_episodes = episode >= max_episodes
reached_goal_mean_reward = mean_100_eval_score >= goal_mean_100_reward
training_is_over = reached_max_minutes or \
reached_max_episodes or \
reached_goal_mean_reward
elapsed_str = time.strftime("%H:%M:%S", time.gmtime(time.time() - training_start))
debug_message = 'el {}, ep {:04}, ts {:06}, '
debug_message += 'ar 10 {:05.1f}\u00B1{:05.1f}, '
debug_message += '100 {:05.1f}\u00B1{:05.1f}, '
debug_message += 'ex 100 {:02.1f}\u00B1{:02.1f}, '
debug_message += 'ev {:05.1f}\u00B1{:05.1f}'
debug_message = debug_message.format(
elapsed_str, episode-1, total_step, mean_10_reward, std_10_reward,
mean_100_reward, std_100_reward, mean_100_exp_rat, std_100_exp_rat,
mean_100_eval_score, std_100_eval_score)
print(debug_message, end='\r', flush=True)
if reached_debug_time or training_is_over:
print(ERASE_LINE + debug_message, flush=True)
last_debug_time = time.time()
if training_is_over:
if reached_max_minutes: print(u'--> reached_max_minutes \u2715')
if reached_max_episodes: print(u'--> reached_max_episodes \u2715')
if reached_goal_mean_reward: print(u'--> reached_goal_mean_reward \u2713')
break
final_eval_score, score_std = self.evaluate(self.policy_model, env, n_episodes=100)
wallclock_time = time.time() - training_start
print('Training complete.')
print('Final evaluation score {:.2f}\u00B1{:.2f} in {:.2f}s training time,'
' {:.2f}s wall-clock time.\n'.format(
final_eval_score, score_std, training_time, wallclock_time))
env.close() ; del env
self.get_cleaned_checkpoints()
return result, final_eval_score, training_time, wallclock_time
def evaluate(self, eval_policy_model, eval_env, n_episodes=1, greedy=True):
rs = []
for _ in range(n_episodes):
s, d = eval_env.reset(), False
rs.append(0)
for _ in count():
if greedy:
a = eval_policy_model.select_greedy_action(s)
else:
a = eval_policy_model.select_action(s)
s, r, d, _ = eval_env.step(a)
rs[-1] += r
if d: break
return np.mean(rs), np.std(rs)
def get_cleaned_checkpoints(self, n_checkpoints=5):
try:
return self.checkpoint_paths
except AttributeError:
self.checkpoint_paths = {}
paths = glob.glob(os.path.join(self.checkpoint_dir, '*.tar'))
paths_dic = {int(path.split('.')[-2]):path for path in paths}
last_ep = max(paths_dic.keys())
# checkpoint_idxs = np.geomspace(1, last_ep+1, n_checkpoints, endpoint=True, dtype=np.int)-1
checkpoint_idxs = np.linspace(1, last_ep+1, n_checkpoints, endpoint=True, dtype=np.int)-1
for idx, path in paths_dic.items():
if idx in checkpoint_idxs:
self.checkpoint_paths[idx] = path
else:
os.unlink(path)
return self.checkpoint_paths
def demo_last(self, title='Fully-trained {} Agent', n_episodes=3, max_n_videos=3):
env = self.make_env_fn(**self.make_env_kargs, monitor_mode='evaluation', render=True, record=True)
checkpoint_paths = self.get_cleaned_checkpoints()
last_ep = max(checkpoint_paths.keys())
self.policy_model.load_state_dict(torch.load(checkpoint_paths[last_ep]))
self.evaluate(self.policy_model, env, n_episodes=n_episodes)
env.close()
data = get_gif_html(env_videos=env.videos,
title=title.format(self.__class__.__name__),
max_n_videos=max_n_videos)
del env
return HTML(data=data)
def demo_progression(self, title='{} Agent progression', max_n_videos=5):
env = self.make_env_fn(**self.make_env_kargs, monitor_mode='evaluation', render=True, record=True)
checkpoint_paths = self.get_cleaned_checkpoints()
for i in sorted(checkpoint_paths.keys()):
self.policy_model.load_state_dict(torch.load(checkpoint_paths[i]))
self.evaluate(self.policy_model, env, n_episodes=1)
env.close()
data = get_gif_html(env_videos=env.videos,
title=title.format(self.__class__.__name__),
subtitle_eps=sorted(checkpoint_paths.keys()),
max_n_videos=max_n_videos)
del env
return HTML(data=data)
def save_checkpoint(self, episode_idx, model):
torch.save(model.state_dict(),
os.path.join(self.checkpoint_dir, 'model.{}.tar'.format(episode_idx)))
reinforce_results = []
best_agent, best_eval_score = None, float('-inf')
for seed in SEEDS:
environment_settings = {
'env_name': 'CartPole-v1',
'gamma': 1.00,
'max_minutes': 10,
'max_episodes': 10000,
'goal_mean_100_reward': 475
}
policy_model_fn = lambda nS, nA: FCDAP(nS, nA, hidden_dims=(128,64))
policy_optimizer_fn = lambda net, lr: optim.Adam(net.parameters(), lr=lr)
policy_optimizer_lr = 0.0005
env_name, gamma, max_minutes, \
max_episodes, goal_mean_100_reward = environment_settings.values()
agent = REINFORCE(policy_model_fn, policy_optimizer_fn, policy_optimizer_lr)
make_env_fn, make_env_kargs = get_make_env_fn(env_name=env_name)
# make_env_fn, make_env_kargs = get_make_env_fn(env_name=env_name, unwrapped=True)
# make_env_fn, make_env_kargs = get_make_env_fn(
# env_name=env_name, addon_wrappers=[MCCartPole,])
result, final_eval_score, training_time, wallclock_time = agent.train(
make_env_fn, make_env_kargs, seed, gamma, max_minutes, max_episodes, goal_mean_100_reward)
reinforce_results.append(result)
if final_eval_score > best_eval_score:
best_eval_score = final_eval_score
best_agent = agent
reinforce_results = np.array(reinforce_results)
best_agent.demo_progression()
best_agent.demo_last()
reinforce_max_t, reinforce_max_r, reinforce_max_s, \
reinforce_max_sec, reinforce_max_rt = np.max(reinforce_results, axis=0).T
reinforce_min_t, reinforce_min_r, reinforce_min_s, \
reinforce_min_sec, reinforce_min_rt = np.min(reinforce_results, axis=0).T
reinforce_mean_t, reinforce_mean_r, reinforce_mean_s, \
reinforce_mean_sec, reinforce_mean_rt = np.mean(reinforce_results, axis=0).T
reinforce_x = np.arange(len(reinforce_mean_s))
# reinforce_max_t, reinforce_max_r, reinforce_max_s, \
# reinforce_max_sec, reinforce_max_rt = np.nanmax(reinforce_results, axis=0).T
# reinforce_min_t, reinforce_min_r, reinforce_min_s, \
# reinforce_min_sec, reinforce_min_rt = np.nanmin(reinforce_results, axis=0).T
# reinforce_mean_t, reinforce_mean_r, reinforce_mean_s, \
# reinforce_mean_sec, reinforce_mean_rt = np.nanmean(reinforce_results, axis=0).T
# reinforce_x = np.arange(len(reinforce_mean_s))
# change convergence checks to episode only (not minutes, not mean reward 'float('inf')' can help)
fig, axs = plt.subplots(5, 1, figsize=(20,30), sharey=False, sharex=True)
# REINFORCE
axs[0].plot(reinforce_max_r, 'y', linewidth=1)
axs[0].plot(reinforce_min_r, 'y', linewidth=1)
axs[0].plot(reinforce_mean_r, 'y', label='REINFORCE', linewidth=2)
axs[0].fill_between(reinforce_x, reinforce_min_r, reinforce_max_r, facecolor='y', alpha=0.3)
axs[1].plot(reinforce_max_s, 'y', linewidth=1)
axs[1].plot(reinforce_min_s, 'y', linewidth=1)
axs[1].plot(reinforce_mean_s, 'y', label='REINFORCE', linewidth=2)
axs[1].fill_between(reinforce_x, reinforce_min_s, reinforce_max_s, facecolor='y', alpha=0.3)
axs[2].plot(reinforce_max_t, 'y', linewidth=1)
axs[2].plot(reinforce_min_t, 'y', linewidth=1)
axs[2].plot(reinforce_mean_t, 'y', label='REINFORCE', linewidth=2)
axs[2].fill_between(reinforce_x, reinforce_min_t, reinforce_max_t, facecolor='y', alpha=0.3)
axs[3].plot(reinforce_max_sec, 'y', linewidth=1)
axs[3].plot(reinforce_min_sec, 'y', linewidth=1)
axs[3].plot(reinforce_mean_sec, 'y', label='REINFORCE', linewidth=2)
axs[3].fill_between(reinforce_x, reinforce_min_sec, reinforce_max_sec, facecolor='y', alpha=0.3)
axs[4].plot(reinforce_max_rt, 'y', linewidth=1)
axs[4].plot(reinforce_min_rt, 'y', linewidth=1)
axs[4].plot(reinforce_mean_rt, 'y', label='REINFORCE', linewidth=2)
axs[4].fill_between(reinforce_x, reinforce_min_rt, reinforce_max_rt, facecolor='y', alpha=0.3)
# ALL
axs[0].set_title('Moving Avg Reward (Training)')
axs[1].set_title('Moving Avg Reward (Evaluation)')
axs[2].set_title('Total Steps')
axs[3].set_title('Training Time')
axs[4].set_title('Wall-clock Time')
plt.xlabel('Episodes')
axs[0].legend(loc='upper left')
plt.show()
reinforce_root_dir = os.path.join(RESULTS_DIR, 'reinforce')
not os.path.exists(reinforce_root_dir) and os.makedirs(reinforce_root_dir)
np.save(os.path.join(reinforce_root_dir, 'x'), reinforce_x)
np.save(os.path.join(reinforce_root_dir, 'max_r'), reinforce_max_r)
np.save(os.path.join(reinforce_root_dir, 'min_r'), reinforce_min_r)
np.save(os.path.join(reinforce_root_dir, 'mean_r'), reinforce_mean_r)
np.save(os.path.join(reinforce_root_dir, 'max_s'), reinforce_max_s)
np.save(os.path.join(reinforce_root_dir, 'min_s'), reinforce_min_s )
np.save(os.path.join(reinforce_root_dir, 'mean_s'), reinforce_mean_s)
np.save(os.path.join(reinforce_root_dir, 'max_t'), reinforce_max_t)
np.save(os.path.join(reinforce_root_dir, 'min_t'), reinforce_min_t)
np.save(os.path.join(reinforce_root_dir, 'mean_t'), reinforce_mean_t)
np.save(os.path.join(reinforce_root_dir, 'max_sec'), reinforce_max_sec)
np.save(os.path.join(reinforce_root_dir, 'min_sec'), reinforce_min_sec)
np.save(os.path.join(reinforce_root_dir, 'mean_sec'), reinforce_mean_sec)
np.save(os.path.join(reinforce_root_dir, 'max_rt'), reinforce_max_rt)
np.save(os.path.join(reinforce_root_dir, 'min_rt'), reinforce_min_rt)
np.save(os.path.join(reinforce_root_dir, 'mean_rt'), reinforce_mean_rt)
weight, probs, entropies = -0.001, [], []
for p in np.arange(0, 1.01, 0.01):
probs.append(p)
p = torch.FloatTensor([p, 1-p])
d = torch.distributions.Categorical(probs=p)
entropies.append(weight * d.entropy().item())
plt.plot(probs, entropies)
plt.xlabel('Probability of action A\np(B)=1-p(A)', labelpad=20)
plt.ylabel('Negative\nweighted\nentropy', labelpad=80, rotation=0)
plt.title('Entropy contribution to the loss function\n{}*entropy(π)'.format(weight), pad=30)
plt.show()
class FCV(nn.Module):
def __init__(self,
input_dim,
hidden_dims=(32,32),
activation_fc=F.relu):
super(FCV, self).__init__()
self.activation_fc = activation_fc
self.input_layer = nn.Linear(input_dim, hidden_dims[0])
self.hidden_layers = nn.ModuleList()
for i in range(len(hidden_dims)-1):
hidden_layer = nn.Linear(hidden_dims[i], hidden_dims[i+1])
self.hidden_layers.append(hidden_layer)
self.output_layer = nn.Linear(hidden_dims[-1], 1)
def _format(self, state):
x = state
if not isinstance(x, torch.Tensor):
x = torch.tensor(x,
dtype=torch.float32)
x = x.unsqueeze(0)
return x
def forward(self, state):
x = self._format(state)
x = self.activation_fc(self.input_layer(x))
for hidden_layer in self.hidden_layers:
x = self.activation_fc(hidden_layer(x))
return self.output_layer(x)
class VPG():
def __init__(self,
policy_model_fn,
policy_model_max_grad_norm,
policy_optimizer_fn,
policy_optimizer_lr,
value_model_fn,
value_model_max_grad_norm,
value_optimizer_fn,
value_optimizer_lr,
entropy_loss_weight):
self.policy_model_fn = policy_model_fn
self.policy_model_max_grad_norm = policy_model_max_grad_norm
self.policy_optimizer_fn = policy_optimizer_fn
self.policy_optimizer_lr = policy_optimizer_lr
self.value_model_fn = value_model_fn
self.value_model_max_grad_norm = value_model_max_grad_norm
self.value_optimizer_fn = value_optimizer_fn
self.value_optimizer_lr = value_optimizer_lr
self.entropy_loss_weight = entropy_loss_weight
def optimize_model(self):
T = len(self.rewards)
discounts = np.logspace(0, T, num=T, base=self.gamma, endpoint=False)
returns = np.array([np.sum(discounts[:T-t] * self.rewards[t:]) for t in range(T)])
discounts = torch.FloatTensor(discounts[:-1]).unsqueeze(1)
returns = torch.FloatTensor(returns[:-1]).unsqueeze(1)
self.logpas = torch.cat(self.logpas)
self.entropies = torch.cat(self.entropies)
self.values = torch.cat(self.values)
value_error = returns - self.values
policy_loss = -(discounts * value_error.detach() * self.logpas).mean()
entropy_loss = -self.entropies.mean()
loss = policy_loss + self.entropy_loss_weight * entropy_loss
self.policy_optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(self.policy_model.parameters(),
self.policy_model_max_grad_norm)
self.policy_optimizer.step()
value_loss = value_error.pow(2).mul(0.5).mean()
self.value_optimizer.zero_grad()
value_loss.backward()
torch.nn.utils.clip_grad_norm_(self.value_model.parameters(),
self.value_model_max_grad_norm)
self.value_optimizer.step()
def interaction_step(self, state, env):
action, is_exploratory, logpa, entropy = self.policy_model.full_pass(state)
new_state, reward, is_terminal, info = env.step(action)
is_truncated = 'TimeLimit.truncated' in info and info['TimeLimit.truncated']
self.logpas.append(logpa)
self.entropies.append(entropy)
self.rewards.append(reward)
self.values.append(self.value_model(state))
self.episode_reward[-1] += reward
self.episode_timestep[-1] += 1
self.episode_exploration[-1] += int(is_exploratory)
return new_state, is_terminal, is_truncated
def train(self, make_env_fn, make_env_kargs, seed, gamma,
max_minutes, max_episodes, goal_mean_100_reward):
training_start, last_debug_time = time.time(), float('-inf')
self.checkpoint_dir = tempfile.mkdtemp()
self.make_env_fn = make_env_fn
self.make_env_kargs = make_env_kargs
self.seed = seed
self.gamma = gamma
env = self.make_env_fn(**self.make_env_kargs, seed=self.seed)
torch.manual_seed(self.seed) ; np.random.seed(self.seed) ; random.seed(self.seed)
nS, nA = env.observation_space.shape[0], env.action_space.n
self.episode_timestep = []
self.episode_reward = []
self.episode_seconds = []
self.episode_exploration = []
self.evaluation_scores = []
self.policy_model = self.policy_model_fn(nS, nA)
self.policy_optimizer = self.policy_optimizer_fn(self.policy_model,
self.policy_optimizer_lr)
self.value_model = self.value_model_fn(nS)
self.value_optimizer = self.value_optimizer_fn(self.value_model,
self.value_optimizer_lr)
result = np.empty((max_episodes, 5))
result[:] = np.nan
training_time = 0
for episode in range(1, max_episodes + 1):
episode_start = time.time()
state, is_terminal = env.reset(), False
self.episode_reward.append(0.0)
self.episode_timestep.append(0.0)
self.episode_exploration.append(0.0)
# collect rollout
self.logpas, self.entropies, self.rewards, self.values = [], [], [], []
for step in count():
state, is_terminal, is_truncated = self.interaction_step(state, env)
if is_terminal:
gc.collect()
break
is_failure = is_terminal and not is_truncated
next_value = 0 if is_failure else self.value_model(state).detach().item()
self.rewards.append(next_value)
self.optimize_model()
# stats
episode_elapsed = time.time() - episode_start
self.episode_seconds.append(episode_elapsed)
training_time += episode_elapsed
evaluation_score, _ = self.evaluate(self.policy_model, env)
self.save_checkpoint(episode-1, self.policy_model)
total_step = int(np.sum(self.episode_timestep))
self.evaluation_scores.append(evaluation_score)
mean_10_reward = np.mean(self.episode_reward[-10:])
std_10_reward = np.std(self.episode_reward[-10:])
mean_100_reward = np.mean(self.episode_reward[-100:])
std_100_reward = np.std(self.episode_reward[-100:])
mean_100_eval_score = np.mean(self.evaluation_scores[-100:])
std_100_eval_score = np.std(self.evaluation_scores[-100:])
lst_100_exp_rat = np.array(
self.episode_exploration[-100:])/np.array(self.episode_timestep[-100:])
mean_100_exp_rat = np.mean(lst_100_exp_rat)
std_100_exp_rat = np.std(lst_100_exp_rat)
wallclock_elapsed = time.time() - training_start
result[episode-1] = total_step, mean_100_reward, \
mean_100_eval_score, training_time, wallclock_elapsed
reached_debug_time = time.time() - last_debug_time >= LEAVE_PRINT_EVERY_N_SECS
reached_max_minutes = wallclock_elapsed >= max_minutes * 60
reached_max_episodes = episode >= max_episodes
reached_goal_mean_reward = mean_100_eval_score >= goal_mean_100_reward
training_is_over = reached_max_minutes or \
reached_max_episodes or \
reached_goal_mean_reward
elapsed_str = time.strftime("%H:%M:%S", time.gmtime(time.time() - training_start))
debug_message = 'el {}, ep {:04}, ts {:06}, '
debug_message += 'ar 10 {:05.1f}\u00B1{:05.1f}, '
debug_message += '100 {:05.1f}\u00B1{:05.1f}, '
debug_message += 'ex 100 {:02.1f}\u00B1{:02.1f}, '
debug_message += 'ev {:05.1f}\u00B1{:05.1f}'
debug_message = debug_message.format(
elapsed_str, episode-1, total_step, mean_10_reward, std_10_reward,
mean_100_reward, std_100_reward, mean_100_exp_rat, std_100_exp_rat,
mean_100_eval_score, std_100_eval_score)
print(debug_message, end='\r', flush=True)
if reached_debug_time or training_is_over:
print(ERASE_LINE + debug_message, flush=True)
last_debug_time = time.time()
if training_is_over:
if reached_max_minutes: print(u'--> reached_max_minutes \u2715')
if reached_max_episodes: print(u'--> reached_max_episodes \u2715')
if reached_goal_mean_reward: print(u'--> reached_goal_mean_reward \u2713')
break
final_eval_score, score_std = self.evaluate(self.policy_model, env, n_episodes=100)
wallclock_time = time.time() - training_start
print('Training complete.')
print('Final evaluation score {:.2f}\u00B1{:.2f} in {:.2f}s training time,'
' {:.2f}s wall-clock time.\n'.format(
final_eval_score, score_std, training_time, wallclock_time))
env.close() ; del env
self.get_cleaned_checkpoints()
return result, final_eval_score, training_time, wallclock_time
def evaluate(self, eval_policy_model, eval_env, n_episodes=1, greedy=True):
rs = []
for _ in range(n_episodes):
s, d = eval_env.reset(), False
rs.append(0)
for _ in count():
if greedy:
a = eval_policy_model.select_greedy_action(s)
else:
a = eval_policy_model.select_action(s)
s, r, d, _ = eval_env.step(a)
rs[-1] += r
if d: break
return np.mean(rs), np.std(rs)
def get_cleaned_checkpoints(self, n_checkpoints=5):
try:
return self.checkpoint_paths
except AttributeError:
self.checkpoint_paths = {}
paths = glob.glob(os.path.join(self.checkpoint_dir, '*.tar'))
paths_dic = {int(path.split('.')[-2]):path for path in paths}
last_ep = max(paths_dic.keys())
# checkpoint_idxs = np.geomspace(1, last_ep+1, n_checkpoints, endpoint=True, dtype=np.int)-1
checkpoint_idxs = np.linspace(1, last_ep+1, n_checkpoints, endpoint=True, dtype=np.int)-1
for idx, path in paths_dic.items():
if idx in checkpoint_idxs:
self.checkpoint_paths[idx] = path
else:
os.unlink(path)
return self.checkpoint_paths
def demo_last(self, title='Fully-trained {} Agent', n_episodes=3, max_n_videos=3):
env = self.make_env_fn(**self.make_env_kargs, monitor_mode='evaluation', render=True, record=True)
checkpoint_paths = self.get_cleaned_checkpoints()
last_ep = max(checkpoint_paths.keys())
self.policy_model.load_state_dict(torch.load(checkpoint_paths[last_ep]))
self.evaluate(self.policy_model, env, n_episodes=n_episodes)
env.close()
data = get_gif_html(env_videos=env.videos,
title=title.format(self.__class__.__name__),
max_n_videos=max_n_videos)
del env
return HTML(data=data)
def demo_progression(self, title='{} Agent progression', max_n_videos=5):
env = self.make_env_fn(**self.make_env_kargs, monitor_mode='evaluation', render=True, record=True)
checkpoint_paths = self.get_cleaned_checkpoints()
for i in sorted(checkpoint_paths.keys()):
self.policy_model.load_state_dict(torch.load(checkpoint_paths[i]))
self.evaluate(self.policy_model, env, n_episodes=1)
env.close()
data = get_gif_html(env_videos=env.videos,
title=title.format(self.__class__.__name__),
subtitle_eps=sorted(checkpoint_paths.keys()),
max_n_videos=max_n_videos)
del env
return HTML(data=data)
def save_checkpoint(self, episode_idx, model):
torch.save(model.state_dict(),
os.path.join(self.checkpoint_dir, 'model.{}.tar'.format(episode_idx)))
vpg_results = []
best_agent, best_eval_score = None, float('-inf')
for seed in SEEDS:
environment_settings = {
'env_name': 'CartPole-v1',
'gamma': 1.00,
'max_minutes': 10,
'max_episodes': 10000,
'goal_mean_100_reward': 475
}
policy_model_fn = lambda nS, nA: FCDAP(nS, nA, hidden_dims=(128,64))
policy_model_max_grad_norm = 1
policy_optimizer_fn = lambda net, lr: optim.Adam(net.parameters(), lr=lr)
policy_optimizer_lr = 0.0005
value_model_fn = lambda nS: FCV(nS, hidden_dims=(256,128))
value_model_max_grad_norm = float('inf')
value_optimizer_fn = lambda net, lr: optim.RMSprop(net.parameters(), lr=lr)
value_optimizer_lr = 0.0007
entropy_loss_weight = 0.001
env_name, gamma, max_minutes, \
max_episodes, goal_mean_100_reward = environment_settings.values()
agent = VPG(policy_model_fn,
policy_model_max_grad_norm,
policy_optimizer_fn,
policy_optimizer_lr,
value_model_fn,
value_model_max_grad_norm,
value_optimizer_fn,
value_optimizer_lr,
entropy_loss_weight)
make_env_fn, make_env_kargs = get_make_env_fn(env_name=env_name)
result, final_eval_score, training_time, wallclock_time = agent.train(
make_env_fn, make_env_kargs, seed, gamma, max_minutes, max_episodes, goal_mean_100_reward)
vpg_results.append(result)
if final_eval_score > best_eval_score:
best_eval_score = final_eval_score
best_agent = agent
vpg_results = np.array(vpg_results)
best_agent.demo_progression()