def ppo(env_fn,
actor_critic=core.mlp_actor_critic,
ac_kwargs=dict(),
seed=0,
steps_per_epoch=4000,
epochs=50,
gamma=0.99,
clip_ratio=0.2,
pi_lr=3e-4,
vf_lr=1e-3,
train_pi_iters=80,
train_v_iters=80,
lam=0.97,
max_ep_len=1000,
target_kl=0.01,
logger_kwargs=dict(),
save_freq=10):
"""
Proximal Policy Optimization (by clipping),
with early stopping based on approximate KL
Args:
env_fn : A function which creates a copy of the environment.
The environment must satisfy the OpenAI Gym API.
actor_critic: A function which takes in placeholder symbols
for state, ``x_ph``, and action, ``a_ph``, and returns the main
outputs from the agent's Tensorflow computation graph:
=========== ================ ======================================
Symbol Shape Description
=========== ================ ======================================
``pi`` (batch, act_dim) | Samples actions from policy given
| states.
``logp`` (batch,) | Gives log probability, according to
| the policy, of taking actions ``a_ph``
| in states ``x_ph``.
``logp_pi`` (batch,) | Gives log probability, according to
| the policy, of the action sampled by
| ``pi``.
``v`` (batch,) | Gives the value estimate for states
| in ``x_ph``. (Critical: make sure
| to flatten this!)
=========== ================ ======================================
ac_kwargs (dict): Any kwargs appropriate for the actor_critic
function you provided to PPO.
seed (int): Seed for random number generators.
steps_per_epoch (int): Number of steps of interaction (state-action pairs)
for the agent and the environment in each epoch.
epochs (int): Number of epochs of interaction (equivalent to
number of policy updates) to perform.
gamma (float): Discount factor. (Always between 0 and 1.)
clip_ratio (float): Hyperparameter for clipping in the policy objective.
Roughly: how far can the new policy go from the old policy while
still profiting (improving the objective function)? The new policy
can still go farther than the clip_ratio says, but it doesn't help
on the objective anymore. (Usually small, 0.1 to 0.3.) Typically
denoted by :math:`\epsilon`.
pi_lr (float): Learning rate for policy optimizer.
vf_lr (float): Learning rate for value function optimizer.
train_pi_iters (int): Maximum number of gradient descent steps to take
on policy loss per epoch. (Early stopping may cause optimizer
to take fewer than this.)
train_v_iters (int): Number of gradient descent steps to take on
value function per epoch.
lam (float): Lambda for GAE-Lambda. (Always between 0 and 1,
close to 1.)
max_ep_len (int): Maximum length of trajectory / episode / rollout.
target_kl (float): Roughly what KL divergence we think is appropriate
between new and old policies after an update. This will get used
for early stopping. (Usually small, 0.01 or 0.05.)
logger_kwargs (dict): Keyword args for EpochLogger.
save_freq (int): How often (in terms of gap between epochs) to save
the current policy and value function.
"""
logger = EpochLogger(**logger_kwargs)
logger.save_config(locals())
seed += 10000 * proc_id()
tf.set_random_seed(seed)
np.random.seed(seed)
env = env_fn()
obs_dim = env.observation_space.shape
act_dim = env.action_space.shape
# Share information about action space with policy architecture
ac_kwargs['action_space'] = env.action_space
# Inputs to computation graph
x_ph, a_ph = core.placeholders_from_spaces(env.observation_space,
env.action_space)
adv_ph, ret_ph, logp_old_ph = core.placeholders(None, None, None)
# Main outputs from computation graph
pi, logp, logp_pi, v = actor_critic(x_ph, a_ph, **ac_kwargs)
# Need all placeholders in *this* order later (to zip with data from buffer)
all_phs = [x_ph, a_ph, adv_ph, ret_ph, logp_old_ph]
# Every step, get: action, value, and logprob
get_action_ops = [pi, v, logp_pi]
# Experience buffer
local_steps_per_epoch = int(steps_per_epoch / num_procs())
buf = PPOBuffer(obs_dim, act_dim, local_steps_per_epoch, gamma, lam)
# Count variables
var_counts = tuple(core.count_vars(scope) for scope in ['pi', 'v'])
logger.log('\nNumber of parameters: \t pi: %d, \t v: %d\n' % var_counts)
# PPO objectives
ratio = tf.exp(logp - logp_old_ph) # pi(a|s) / pi_old(a|s)
min_adv = tf.where(adv_ph > 0, (1 + clip_ratio) * adv_ph,
(1 - clip_ratio) * adv_ph)
pi_loss = -tf.reduce_mean(tf.minimum(ratio * adv_ph, min_adv))
v_loss = tf.reduce_mean((ret_ph - v)**2)
# Info (useful to watch during learning)
approx_kl = tf.reduce_mean(
logp_old_ph -
logp) # a sample estimate for KL-divergence, easy to compute
approx_ent = tf.reduce_mean(
-logp) # a sample estimate for entropy, also easy to compute
clipped = tf.logical_or(ratio > (1 + clip_ratio), ratio < (1 - clip_ratio))
clipfrac = tf.reduce_mean(tf.cast(clipped, tf.float32))
# Optimizers
train_pi = MpiAdamOptimizer(learning_rate=pi_lr).minimize(pi_loss)
train_v = MpiAdamOptimizer(learning_rate=vf_lr).minimize(v_loss)
sess = tf.Session()
sess.run(tf.global_variables_initializer())
# Sync params across processes
sess.run(sync_all_params())
# Setup model saving
logger.setup_tf_saver(sess, inputs={'x': x_ph}, outputs={'pi': pi, 'v': v})
def update():
inputs = {k: v for k, v in zip(all_phs, buf.get())}
pi_l_old, v_l_old, ent = sess.run([pi_loss, v_loss, approx_ent],
feed_dict=inputs)
# Training
for i in range(train_pi_iters):
_, kl = sess.run([train_pi, approx_kl], feed_dict=inputs)
kl = mpi_avg(kl)
if kl > 1.5 * target_kl:
logger.log(
'Early stopping at step %d due to reaching max kl.' % i)
break
logger.store(StopIter=i)
for _ in range(train_v_iters):
sess.run(train_v, feed_dict=inputs)
# Log changes from update
pi_l_new, v_l_new, kl, cf = sess.run(
[pi_loss, v_loss, approx_kl, clipfrac], feed_dict=inputs)
logger.store(LossPi=pi_l_old,
LossV=v_l_old,
KL=kl,
Entropy=ent,
ClipFrac=cf,
DeltaLossPi=(pi_l_new - pi_l_old),
DeltaLossV=(v_l_new - v_l_old))
start_time = time.time()
o, r, d, ep_ret, ep_len = env.reset(), 0, False, 0, 0
# Main loop: collect experience in env and update/log each epoch
for epoch in range(epochs):
for t in range(local_steps_per_epoch):
a, v_t, logp_t = sess.run(get_action_ops,
feed_dict={x_ph: o.reshape(1, -1)})
o2, r, d, _ = env.step(a[0])
ep_ret += r
ep_len += 1
# save and log
buf.store(o, a, r, v_t, logp_t)
logger.store(VVals=v_t)
# Update obs (critical!)
o = o2
terminal = d or (ep_len == max_ep_len)
if terminal or (t == local_steps_per_epoch - 1):
if not (terminal):
print('Warning: trajectory cut off by epoch at %d steps.' %
ep_len)
# if trajectory didn't reach terminal state, bootstrap value target
last_val = 0 if d else sess.run(
v, feed_dict={x_ph: o.reshape(1, -1)})
buf.finish_path(last_val)
if terminal:
# only save EpRet / EpLen if trajectory finished
logger.store(EpRet=ep_ret, EpLen=ep_len)
o, ep_ret, ep_len = env.reset(), 0, 0
# Save model
if (epoch % save_freq == 0) or (epoch == epochs - 1):
logger.save_state({'env': env}, None)
# Perform PPO update!
update()
# Log info about epoch
logger.log_tabular('Epoch', epoch)
logger.log_tabular('EpRet', with_min_and_max=True)
logger.log_tabular('EpLen', average_only=True)
logger.log_tabular('VVals', with_min_and_max=True)
logger.log_tabular('TotalEnvInteracts', (epoch + 1) * steps_per_epoch)
logger.log_tabular('LossPi', average_only=True)
logger.log_tabular('LossV', average_only=True)
logger.log_tabular('DeltaLossPi', average_only=True)
logger.log_tabular('DeltaLossV', average_only=True)
logger.log_tabular('Entropy', average_only=True)
logger.log_tabular('KL', average_only=True)
logger.log_tabular('ClipFrac', average_only=True)
logger.log_tabular('StopIter', average_only=True)
logger.log_tabular('Time', time.time() - start_time)
logger.dump_tabular()
class EMAQ:
def __init__(self,
env_fn,
env_name=None,
actor_critic=core.MLPActorCritic,
ac_kwargs=dict(),
seed=0,
steps_per_epoch=100,
epochs=10000,
replay_size=int(2000000),
gamma=0.99,
polyak=0.995,
lr=3e-4,
p_lr=3e-5,
alpha=0.2,
batch_size=100,
start_steps=10000,
update_after=1000,
update_every=50,
num_test_episodes=10,
max_ep_len=1000,
logger_kwargs=dict(),
save_freq=1,
algo='CQL'):
"""
Soft Actor-Critic (SAC)
Args:
env_fn : A function which creates a copy of the environment.
The environment must satisfy the OpenAI Gym API.
actor_critic: The constructor method for a PyTorch Module with an ``act``
method, a ``pi`` module, a ``q1`` module, and a ``q2`` module.
The ``act`` method and ``pi`` module should accept batches of
observations as inputs, and ``q1`` and ``q2`` should accept a batch
of observations and a batch of actions as inputs. When called,
``act``, ``q1``, and ``q2`` should return:
=========== ================ ======================================
Call Output Shape Description
=========== ================ ======================================
``act`` (batch, act_dim) | Numpy array of actions for each
| observation.
``q1`` (batch,) | Tensor containing one current estimate
| of Q* for the provided observations
| and actions. (Critical: make sure to
| flatten this!)
``q2`` (batch,) | Tensor containing the other current
| estimate of Q* for the provided observations
| and actions. (Critical: make sure to
| flatten this!)
=========== ================ ======================================
Calling ``pi`` should return:
=========== ================ ======================================
Symbol Shape Description
=========== ================ ======================================
``a`` (batch, act_dim) | Tensor containing actions from policy
| given observations.
``logp_pi`` (batch,) | Tensor containing log probabilities of
| actions in ``a``. Importantly: gradients
| should be able to flow back into ``a``.
=========== ================ ======================================
ac_kwargs (dict): Any kwargs appropriate for the ActorCritic object
you provided to SAC.
seed (int): Seed for random number generators.
steps_per_epoch (int): Number of steps of interaction (state-action pairs)
for the agent and the environment in each epoch.
epochs (int): Number of epochs to run and train agent.
replay_size (int): Maximum length of replay buffer.
gamma (float): Discount factor. (Always between 0 and 1.)
polyak (float): Interpolation factor in polyak averaging for target
networks. Target networks are updated towards main networks
according to:
.. math:: \\theta_{\\text{targ}} \\leftarrow
\\rho \\theta_{\\text{targ}} + (1-\\rho) \\theta
where :math:`\\rho` is polyak. (Always between 0 and 1, usually
close to 1.)
lr (float): Learning rate (used for both policy and value learning).
alpha (float): Entropy regularization coefficient. (Equivalent to
inverse of reward scale in the original SAC paper.)
batch_size (int): Minibatch size for SGD.
start_steps (int): Number of steps for uniform-random action selection,
before running real policy. Helps exploration.
update_after (int): Number of env interactions to collect before
starting to do gradient descent updates. Ensures replay buffer
is full enough for useful updates.
update_every (int): Number of env interactions that should elapse
between gradient descent updates. Note: Regardless of how long
you wait between updates, the ratio of env steps to gradient steps
is locked to 1.
num_test_episodes (int): Number of episodes to test the deterministic
policy at the end of each epoch.
max_ep_len (int): Maximum length of trajectory / episode / rollout.
logger_kwargs (dict): Keyword args for EpochLogger.
save_freq (int): How often (in terms of gap between epochs) to save
the current policy and value function.
"""
self.logger = EpochLogger(**logger_kwargs)
self.logger.save_config(locals())
torch.manual_seed(seed)
np.random.seed(seed)
self.env, self.test_env = env_fn(), env_fn()
self.obs_dim = self.env.observation_space.shape
self.act_dim = self.env.action_space.shape[0]
# Action limit for clamping: critically, assumes all dimensions share the same bound!
self.act_limit = self.env.action_space.high[0]
# Create actor-critic module and target networks
self.ac = actor_critic(self.env.observation_space,
self.env.action_space, **ac_kwargs)
self.ac_targ = deepcopy(self.ac)
self.gamma = gamma
# Freeze target networks with respect to optimizers (only update via polyak averaging)
for p in self.ac_targ.parameters():
p.requires_grad = False
# List of parameters for both Q-networks (save this for convenience)
self.q_params = itertools.chain(self.ac.q1.parameters(),
self.ac.q2.parameters())
# Experience buffer
self.replay_buffer = ReplayBuffer(obs_dim=self.obs_dim,
act_dim=self.act_dim,
size=replay_size)
# Count variables (protip: try to get a feel for how different size networks behave!)
var_counts = tuple(
core.count_vars(module)
for module in [self.ac.pi, self.ac.q1, self.ac.q2])
self.logger.log(
'\nNumber of parameters: \t pi: %d, \t q1: %d, \t q2: %d\n' %
var_counts)
self.algo = algo
self.lagrange_threshold = 10
self.penalty_lr = 5e-2
self.lamda = Variable(torch.log(torch.exp(torch.Tensor([5])) - 1),
requires_grad=True)
self.lamda_optimizer = torch.optim.Adam([self.lamda],
lr=self.penalty_lr)
self.tune_lambda = True if 'lagrange' in self.algo else False
self.alpha = 0
self.target_update_freq = 1
self.p_lr = 3e-5
self.lr = 3e-4
self.n_samples = 100
self.env_name = env_name
# Set up optimizers for policy and q-function
self.pi_optimizer = Adam(self.ac.pi.parameters(), lr=self.p_lr)
self.q_optimizer = Adam(self.q_params, lr=self.lr)
self.num_test_episodes = num_test_episodes
self.max_ep_len = max_ep_len
self.epochs = epochs
self.steps_per_epoch = steps_per_epoch
self.update_after = update_after
self.update_every = update_every
self.batch_size = batch_size
self.save_freq = save_freq
self.polyak = polyak
# Set up model saving
self.logger.setup_pytorch_saver(self.ac)
print("Running Offline RL algorithm: {}".format(self.algo))
def populate_replay_buffer(self):
dataset = d4rl.qlearning_dataset(self.env)
self.replay_buffer.obs_buf[:dataset['observations'].
shape[0], :] = dataset['observations']
self.replay_buffer.act_buf[:dataset['actions'].
shape[0], :] = dataset['actions']
self.replay_buffer.obs2_buf[:dataset['next_observations'].
shape[0], :] = dataset['next_observations']
self.replay_buffer.rew_buf[:dataset['rewards'].
shape[0]] = dataset['rewards']
self.replay_buffer.done_buf[:dataset['terminals'].
shape[0]] = dataset['terminals']
self.replay_buffer.size = dataset['observations'].shape[0]
self.replay_buffer.ptr = (self.replay_buffer.size +
1) % (self.replay_buffer.max_size)
# Set up function for computing SAC Q-losses
def compute_loss_q(self, data):
o, a, r, o2, d = data['obs'], data['act'], data['rew'], data[
'obs2'], data['done']
sampled_actions_q1 = None
sampled_actions_q2 = None
for i in range(self.n_samples):
z = np.random.randn(a.shape[0], a.shape[1])
z = torch.FloatTensor(z)
actions, _ = self.sampling_policy.inverse(z, y=o2)
if sampled_actions_q1 is None:
sampled_actions_q1 = self.ac_targ.q1(o2, actions).view(-1, 1)
sampled_actions_q2 = self.ac_targ.q2(o2, actions).view(-1, 1)
else:
sampled_actions_q1 = torch.cat(
(sampled_actions_q1, self.ac_targ.q1(o2, actions).view(
-1, 1)),
dim=1)
sampled_actions_q2 = torch.cat(
(sampled_actions_q2, self.ac_targ.q2(o2, actions).view(
-1, 1)),
dim=1)
q1 = self.ac.q1(o, a)
q2 = self.ac.q2(o, a)
# Bellman backup for Q functions
with torch.no_grad():
# Target actions come from *current* policy
a2, logp_a2 = self.ac.pi(o2)
# Target Q-values
q1_pi_targ = torch.max(sampled_actions_q1, dim=1).values
q2_pi_targ = torch.max(sampled_actions_q2, dim=1).values
q_pi_targ = torch.min(q1_pi_targ, q2_pi_targ)
backup = r + self.gamma * (1 - d) * (q_pi_targ -
self.alpha * logp_a2)
# MSE loss against Bellman backup
loss_q1 = ((q1 - backup)**2).mean()
loss_q2 = ((q2 - backup)**2).mean()
loss_q = loss_q1 + loss_q2
# Useful info for logging
q_info = dict(Q1Vals=q1.detach().numpy(), Q2Vals=q2.detach().numpy())
return loss_q, q_info
def update(self, data, update_timestep):
# First run one gradient descent step for Q1 and Q2
self.q_optimizer.zero_grad()
loss_q, q_info = self.compute_loss_q(data)
loss_q.backward()
self.q_optimizer.step()
# Record things
self.logger.store(LossQ=loss_q.item(), **q_info)
# Finally, update target networks by polyak averaging.
if update_timestep % self.target_update_freq == 0:
with torch.no_grad():
for p, p_targ in zip(self.ac.parameters(),
self.ac_targ.parameters()):
# NB: We use an in-place operations "mul_", "add_" to update target
# params, as opposed to "mul" and "add", which would make new tensors.
p_targ.data.mul_(self.polyak)
p_targ.data.add_((1 - self.polyak) * p.data)
def get_action(self, o, deterministic=False):
sampled_actions_q1 = None
sampled_actions_q2 = None
sampled_actions = []
o = torch.FloatTensor(o).view(1, -1)
for i in range(self.n_samples):
z = np.random.randn(1, self.act_dim)
z = torch.FloatTensor(z)
actions, _ = self.sampling_policy.inverse(z, y=o)
sampled_actions.append(actions)
if sampled_actions_q1 is None:
sampled_actions_q1 = self.ac.q1(o, actions).view(-1, 1)
sampled_actions_q2 = self.ac.q2(o, actions).view(-1, 1)
else:
sampled_actions_q1 = torch.cat(
(sampled_actions_q1, self.ac.q1(o, actions).view(-1, 1)),
dim=1)
sampled_actions_q2 = torch.cat(
(sampled_actions_q2, self.ac.q2(o, actions).view(-1, 1)),
dim=1)
q_values = torch.min(sampled_actions_q1, sampled_actions_q2)
max_idx = torch.argmax(q_values.view(-1))
return sampled_actions[max_idx].detach().cpu().numpy()
def test_agent(self):
for j in range(self.num_test_episodes):
o, d, ep_ret, ep_len = self.test_env.reset(), False, 0, 0
while not (d or (ep_len == self.max_ep_len)):
# Take deterministic actions at test time
o, r, d, _ = self.test_env.step(self.get_action(o, True))
ep_ret += r
ep_len += 1
self.logger.store(TestEpRet=100 *
self.test_env.get_normalized_score(ep_ret),
TestEpLen=ep_len)
def run(self):
# Learn a generative model for data
# density_epochs = 50
# self.sampling_policy = core.MADE(self.act_dim, 256, 2 , cond_label_size = self.obs_dim[0])
# density_optimizer = torch.optim.Adam(self.sampling_policy.parameters(), lr=1e-4, weight_decay=1e-6)
# for i in range(density_epochs):
# sample_indices = np.random.choice(
# self.replay_buffer.size, self.replay_buffer.size)
# np.random.shuffle(sample_indices)
# ctr = 0
# total_loss = 0
# for j in range(0, self.replay_buffer.size, self.batch_size):
# actions = self.replay_buffer.act_buf[sample_indices[ctr * self.batch_size:(
# ctr + 1) * self.batch_size],:]
# actions = torch.FloatTensor(actions)
# obs = self.replay_buffer.obs_buf[sample_indices[ctr * self.batch_size:(
# ctr + 1) * self.batch_size],:]
# obs = torch.FloatTensor(obs)
# density_optimizer.zero_grad()
# loss = -self.sampling_policy.log_prob(actions,y=obs).mean()
# loss.backward()
# total_loss+=loss.data * self.batch_size
# density_optimizer.step()
# ctr+=1
# print("Density training loss: {}".format(total_loss/self.replay_buffer.size))
self.sampling_policy = core.MADE(self.act_dim,
256,
3,
cond_label_size=self.obs_dim[0])
self.sampling_policy.load_state_dict(
torch.load("behavior_policies/" + self.env_name + ".pt"))
# self.sampling_policy = torch.load("marginals/"+self.env_name+".pt")
# Prepare for interaction with environment
total_steps = self.epochs * self.steps_per_epoch
start_time = time.time()
o, ep_ret, ep_len = self.env.reset(), 0, 0
# Main loop: collect experience in env and update/log each epoch
for t in range(total_steps):
# # Update handling
batch = self.replay_buffer.sample_batch(self.batch_size)
self.update(data=batch, update_timestep=t)
# End of epoch handling
if (t + 1) % self.steps_per_epoch == 0:
epoch = (t + 1) // self.steps_per_epoch
# Save model
if (epoch % self.save_freq == 0) or (epoch == self.epochs):
self.logger.save_state({'env': self.env}, None)
# Test the performance of the deterministic version of the agent.
self.test_agent()
# Log info about epoch
self.logger.log_tabular('Epoch', epoch)
self.logger.log_tabular('TestEpRet', with_min_and_max=True)
self.logger.log_tabular('TestEpLen', average_only=True)
self.logger.log_tabular('TotalUpdates', t)
self.logger.log_tabular('Q1Vals', with_min_and_max=True)
self.logger.log_tabular('Q2Vals', with_min_and_max=True)
self.logger.log_tabular('LossQ', average_only=True)
self.logger.log_tabular('Time', time.time() - start_time)
self.logger.dump_tabular()
def valor(args):
if not hasattr(args, "get"):
args.get = args.__dict__.get
env_fn = args.get('env_fn', lambda: gym.make('HalfCheetah-v2'))
actor_critic = args.get('actor_critic', ActorCritic)
ac_kwargs = args.get('ac_kwargs', {})
disc = args.get('disc', Discriminator)
dc_kwargs = args.get('dc_kwargs', {})
seed = args.get('seed', 0)
episodes_per_epoch = args.get('episodes_per_epoch', 40)
epochs = args.get('epochs', 50)
gamma = args.get('gamma', 0.99)
pi_lr = args.get('pi_lr', 3e-4)
vf_lr = args.get('vf_lr', 1e-3)
dc_lr = args.get('dc_lr', 2e-3)
train_v_iters = args.get('train_v_iters', 80)
train_dc_iters = args.get('train_dc_iters', 50)
train_dc_interv = args.get('train_dc_interv', 2)
lam = args.get('lam', 0.97)
max_ep_len = args.get('max_ep_len', 1000)
logger_kwargs = args.get('logger_kwargs', {})
context_dim = args.get('context_dim', 4)
max_context_dim = args.get('max_context_dim', 64)
save_freq = args.get('save_freq', 10)
k = args.get('k', 1)
logger = EpochLogger(**logger_kwargs)
logger.save_config(locals())
# seed += 10000 * proc_id()
torch.manual_seed(seed)
np.random.seed(seed)
env = env_fn()
obs_dim = env.observation_space.shape
act_dim = env.action_space.shape
ac_kwargs['action_space'] = env.action_space
# Model
actor_critic = actor_critic(input_dim=obs_dim[0] + max_context_dim,
**ac_kwargs)
disc = disc(input_dim=obs_dim[0], context_dim=max_context_dim, **dc_kwargs)
# Buffer
local_episodes_per_epoch = episodes_per_epoch # int(episodes_per_epoch / num_procs())
buffer = Buffer(max_context_dim, obs_dim[0], act_dim[0],
local_episodes_per_epoch, max_ep_len, train_dc_interv)
# Count variables
var_counts = tuple(
count_vars(module)
for module in [actor_critic.policy, actor_critic.value_f, disc.policy])
logger.log('\nNumber of parameters: \t pi: %d, \t v: %d, \t d: %d\n' %
var_counts)
# Optimizers
#Optimizer for RL Policy
train_pi = torch.optim.Adam(actor_critic.policy.parameters(), lr=pi_lr)
#Optimizer for value function (for actor-critic)
train_v = torch.optim.Adam(actor_critic.value_f.parameters(), lr=vf_lr)
#Optimizer for decoder
train_dc = torch.optim.Adam(disc.policy.parameters(), lr=dc_lr)
#pdb.set_trace()
# Parameters Sync
#sync_all_params(actor_critic.parameters())
#sync_all_params(disc.parameters())
'''
Training function
'''
def update(e):
obs, act, adv, pos, ret, logp_old = [
torch.Tensor(x) for x in buffer.retrieve_all()
]
# Policy
#pdb.set_trace()
_, logp, _ = actor_critic.policy(obs, act, batch=False)
#pdb.set_trace()
entropy = (-logp).mean()
# Policy loss
pi_loss = -(logp * (k * adv + pos)).mean()
# Train policy (Go through policy update)
train_pi.zero_grad()
pi_loss.backward()
# average_gradients(train_pi.param_groups)
train_pi.step()
# Value function
v = actor_critic.value_f(obs)
v_l_old = F.mse_loss(v, ret)
for _ in range(train_v_iters):
v = actor_critic.value_f(obs)
v_loss = F.mse_loss(v, ret)
# Value function train
train_v.zero_grad()
v_loss.backward()
# average_gradients(train_v.param_groups)
train_v.step()
# Discriminator
if (e + 1) % train_dc_interv == 0:
print('Discriminator Update!')
con, s_diff = [torch.Tensor(x) for x in buffer.retrieve_dc_buff()]
_, logp_dc, _ = disc(s_diff, con)
d_l_old = -logp_dc.mean()
# Discriminator train
for _ in range(train_dc_iters):
_, logp_dc, _ = disc(s_diff, con)
d_loss = -logp_dc.mean()
train_dc.zero_grad()
d_loss.backward()
# average_gradients(train_dc.param_groups)
train_dc.step()
_, logp_dc, _ = disc(s_diff, con)
dc_l_new = -logp_dc.mean()
else:
d_l_old = 0
dc_l_new = 0
# Log the changes
_, logp, _, v = actor_critic(obs, act)
pi_l_new = -(logp * (k * adv + pos)).mean()
v_l_new = F.mse_loss(v, ret)
kl = (logp_old - logp).mean()
logger.store(LossPi=pi_loss,
LossV=v_l_old,
KL=kl,
Entropy=entropy,
DeltaLossPi=(pi_l_new - pi_loss),
DeltaLossV=(v_l_new - v_l_old),
LossDC=d_l_old,
DeltaLossDC=(dc_l_new - d_l_old))
# logger.store(Adv=adv.reshape(-1).numpy().tolist(), Pos=pos.reshape(-1).numpy().tolist())
start_time = time.time()
#Resets observations, rewards, done boolean
o, r, d, ep_ret, ep_len = env.reset(), 0, False, 0, 0
#Creates context distribution where each logit is equal to one (This is first place to make change)
context_dim_prob_dict = {
i: 1 / context_dim if i < context_dim else 0
for i in range(max_context_dim)
}
last_phi_dict = {i: 0 for i in range(context_dim)}
context_dist = Categorical(
probs=torch.Tensor(list(context_dim_prob_dict.values())))
total_t = 0
for epoch in range(epochs):
#Sets actor critic and decoder (discriminator) into eval mode
actor_critic.eval()
disc.eval()
#Runs the policy local_episodes_per_epoch before updating the policy
for index in range(local_episodes_per_epoch):
# Sample from context distribution and one-hot encode it (Step 2)
# Every time we run the policy we sample a new context
c = context_dist.sample()
c_onehot = F.one_hot(c, max_context_dim).squeeze().float()
for _ in range(max_ep_len):
concat_obs = torch.cat(
[torch.Tensor(o.reshape(1, -1)),
c_onehot.reshape(1, -1)], 1)
'''
Feeds in observation and context into actor_critic which spits out a distribution
Label is a sample from the observation
pi is the action sampled
logp is the log probability of some other action a
logp_pi is the log probability of pi
v_t is the value function
'''
a, _, logp_t, v_t = actor_critic(concat_obs)
#Stores context and all other info about the state in the buffer
buffer.store(c,
concat_obs.squeeze().detach().numpy(),
a.detach().numpy(), r, v_t.item(),
logp_t.detach().numpy())
logger.store(VVals=v_t)
o, r, d, _ = env.step(a.detach().numpy()[0])
ep_ret += r
ep_len += 1
total_t += 1
terminal = d or (ep_len == max_ep_len)
if terminal:
# Key stuff with discriminator
dc_diff = torch.Tensor(buffer.calc_diff()).unsqueeze(0)
#Context
con = torch.Tensor([float(c)]).unsqueeze(0)
#Feed in differences between each state in your trajectory and a specific context
#Here, this is just the log probability of the label it thinks it is
_, _, log_p = disc(dc_diff, con)
buffer.end_episode(log_p.detach().numpy())
logger.store(EpRet=ep_ret, EpLen=ep_len)
o, r, d, ep_ret, ep_len = env.reset(), 0, False, 0, 0
if (epoch % save_freq == 0) or (epoch == epochs - 1):
logger.save_state({'env': env}, [actor_critic, disc], None)
# Sets actor_critic and discriminator into training mode
actor_critic.train()
disc.train()
update(epoch)
#Need to implement curriculum learning here to update context distribution
'''
#Psuedocode:
Loop through each of d episodes taken in local_episodes_per_epoch and check log probability from discrimantor
If >= 0.86, increase k in the following manner: k = min(int(1.5*k + 1), Kmax)
Kmax = 64
'''
decoder_accs = []
stag_num = 10
stag_pct = 0.05
if (epoch + 1) % train_dc_interv == 0 and epoch > 0:
#pdb.set_trace()
con, s_diff = [torch.Tensor(x) for x in buffer.retrieve_dc_buff()]
print("Context: ", con)
print("num_contexts", len(con))
_, logp_dc, _ = disc(s_diff, con)
log_p_context_sample = logp_dc.mean().detach().numpy()
print("Log Probability context sample", log_p_context_sample)
decoder_accuracy = np.exp(log_p_context_sample)
print("Decoder Accuracy", decoder_accuracy)
logger.store(LogProbabilityContext=log_p_context_sample,
DecoderAccuracy=decoder_accuracy)
'''
Create score (phi(i)) = -log_p_context_sample.mean() for each specific context
Assign phis to each specific context
Get p(i) in the following manner: (phi(i) + epsilon)
Get Probabilities by doing p(i)/sum of all p(i)'s
'''
logp_np = logp_dc.detach().numpy()
con_np = con.detach().numpy()
phi_dict = {i: 0 for i in range(context_dim)}
count_dict = {i: 0 for i in range(context_dim)}
for i in range(len(logp_np)):
current_con = con_np[i]
phi_dict[current_con] += logp_np[i]
count_dict[current_con] += 1
print(phi_dict)
phi_dict = {
k: last_phi_dict[k] if count_dict[k] == 0 else
(-1) * v / count_dict[k]
for (k, v) in phi_dict.items()
}
sorted_dict = dict(
sorted(phi_dict.items(),
key=lambda item: item[1],
reverse=True))
sorted_dict_keys = list(sorted_dict.keys())
rank_dict = {
sorted_dict_keys[i]: 1 / (i + 1)
for i in range(len(sorted_dict_keys))
}
rank_dict_sum = sum(list(rank_dict.values()))
context_dim_prob_dict = {
k: rank_dict[k] / rank_dict_sum if k < context_dim else 0
for k in context_dim_prob_dict.keys()
}
print(context_dim_prob_dict)
decoder_accs.append(decoder_accuracy)
stagnated = (len(decoder_accs) > stag_num
and (decoder_accs[-stag_num - 1] - decoder_accuracy) /
stag_num < stag_pct)
if stagnated:
new_context_dim = max(int(0.75 * context_dim), 5)
elif decoder_accuracy >= 0.86:
new_context_dim = min(int(1.5 * context_dim + 1),
max_context_dim)
if stagnated or decoder_accuracy >= 0.86:
print("new_context_dim: ", new_context_dim)
new_context_prob_arr = np.array(
new_context_dim * [1 / new_context_dim] +
(max_context_dim - new_context_dim) * [0])
context_dist = Categorical(
probs=ptu.from_numpy(new_context_prob_arr))
context_dim = new_context_dim
for i in range(context_dim):
if i in phi_dict:
last_phi_dict[i] = phi_dict[i]
elif i not in last_phi_dict:
last_phi_dict[i] = max(phi_dict.values())
buffer.clear_dc_buff()
else:
logger.store(LogProbabilityContext=0, DecoderAccuracy=0)
# Log
logger.store(ContextDim=context_dim)
logger.log_tabular('Epoch', epoch)
logger.log_tabular('EpRet', with_min_and_max=True)
logger.log_tabular('EpLen', average_only=True)
logger.log_tabular('VVals', with_min_and_max=True)
logger.log_tabular('TotalEnvInteracts', total_t)
logger.log_tabular('LossPi', average_only=True)
logger.log_tabular('DeltaLossPi', average_only=True)
logger.log_tabular('LossV', average_only=True)
logger.log_tabular('DeltaLossV', average_only=True)
logger.log_tabular('LossDC', average_only=True)
logger.log_tabular('DeltaLossDC', average_only=True)
logger.log_tabular('Entropy', average_only=True)
logger.log_tabular('KL', average_only=True)
logger.log_tabular('Time', time.time() - start_time)
logger.log_tabular('LogProbabilityContext', average_only=True)
logger.log_tabular('DecoderAccuracy', average_only=True)
logger.log_tabular('ContextDim', average_only=True)
logger.dump_tabular()
def td3(env_fn,
actor_critic=core.mlp_actor_critic,
ac_kwargs=dict(),
seed=0,
steps_per_epoch=4000,
epochs=100,
replay_size=int(1e6),
gamma=0.99,
polyak=0.995,
pi_lr=1e-3,
q_lr=1e-3,
batch_size=100,
start_steps=10000,
update_after=1000,
update_every=50,
act_noise=0.1,
target_noise=0.2,
noise_clip=0.5,
policy_delay=2,
num_test_episodes=10,
max_ep_len=1000,
logger_kwargs=dict(),
save_freq=1):
"""
Twin Delayed Deep Deterministic Policy Gradient (TD3)
Args:
env_fn : A function which creates a copy of the environment.
The environment must satisfy the OpenAI Gym API.
actor_critic: A function which takes in placeholder symbols
for state, ``x_ph``, and action, ``a_ph``, and returns the main
outputs from the agent's Tensorflow computation graph:
=========== ================ ======================================
Symbol Shape Description
=========== ================ ======================================
``pi`` (batch, act_dim) | Deterministically computes actions
| from policy given states.
``q1`` (batch,) | Gives one estimate of Q* for
| states in ``x_ph`` and actions in
| ``a_ph``.
``q2`` (batch,) | Gives another estimate of Q* for
| states in ``x_ph`` and actions in
| ``a_ph``.
``q1_pi`` (batch,) | Gives the composition of ``q1`` and
| ``pi`` for states in ``x_ph``:
| q1(x, pi(x)).
=========== ================ ======================================
ac_kwargs (dict): Any kwargs appropriate for the actor_critic
function you provided to TD3.
seed (int): Seed for random number generators.
steps_per_epoch (int): Number of steps of interaction (state-action pairs)
for the agent and the environment in each epoch.
epochs (int): Number of epochs to run and train agent.
replay_size (int): Maximum length of replay buffer.
gamma (float): Discount factor. (Always between 0 and 1.)
polyak (float): Interpolation factor in polyak averaging for target
networks. Target networks are updated towards main networks
according to:
.. math:: \\theta_{\\text{targ}} \\leftarrow
\\rho \\theta_{\\text{targ}} + (1-\\rho) \\theta
where :math:`\\rho` is polyak. (Always between 0 and 1, usually
close to 1.)
pi_lr (float): Learning rate for policy.
q_lr (float): Learning rate for Q-networks.
batch_size (int): Minibatch size for SGD.
start_steps (int): Number of steps for uniform-random action selection,
before running real policy. Helps exploration.
update_after (int): Number of env interactions to collect before
starting to do gradient descent updates. Ensures replay buffer
is full enough for useful updates.
update_every (int): Number of env interactions that should elapse
between gradient descent updates. Note: Regardless of how long
you wait between updates, the ratio of env steps to gradient steps
is locked to 1.
act_noise (float): Stddev for Gaussian exploration noise added to
policy at training time. (At test time, no noise is added.)
target_noise (float): Stddev for smoothing noise added to target
policy.
noise_clip (float): Limit for absolute value of target policy
smoothing noise.
policy_delay (int): Policy will only be updated once every
policy_delay times for each update of the Q-networks.
num_test_episodes (int): Number of episodes to test the deterministic
policy at the end of each epoch.
max_ep_len (int): Maximum length of trajectory / episode / rollout.
logger_kwargs (dict): Keyword args for EpochLogger.
save_freq (int): How often (in terms of gap between epochs) to save
the current policy and value function.
"""
logger = EpochLogger(**logger_kwargs)
logger.save_config(locals())
tf.set_random_seed(seed)
np.random.seed(seed)
env, test_env = env_fn(), env_fn()
obs_dim = env.observation_space.shape[0]
act_dim = env.action_space.shape[0]
# Action limit for clamping: critically, assumes all dimensions share the same bound!
act_limit = env.action_space.high[0]
# Share information about action space with policy architecture
ac_kwargs['action_space'] = env.action_space
# Inputs to computation graph
x_ph, a_ph, x2_ph, r_ph, d_ph = core.placeholders(obs_dim, act_dim,
obs_dim, None, None)
# Main outputs from computation graph
with tf.variable_scope('main'):
pi, q1, q2, q1_pi = actor_critic(x_ph, a_ph, **ac_kwargs)
# Target policy network
with tf.variable_scope('target'):
pi_targ, _, _, _ = actor_critic(x2_ph, a_ph, **ac_kwargs)
# Target Q networks
with tf.variable_scope('target', reuse=True):
# Target policy smoothing, by adding clipped noise to target actions
epsilon = tf.random_normal(tf.shape(pi_targ), stddev=target_noise)
epsilon = tf.clip_by_value(epsilon, -noise_clip, noise_clip)
a2 = pi_targ + epsilon
a2 = tf.clip_by_value(a2, -act_limit, act_limit)
# Target Q-values, using action from target policy
_, q1_targ, q2_targ, _ = actor_critic(x2_ph, a2, **ac_kwargs)
# Experience buffer
replay_buffer = ReplayBuffer(obs_dim=obs_dim,
act_dim=act_dim,
size=replay_size)
# Count variables
var_counts = tuple(
core.count_vars(scope)
for scope in ['main/pi', 'main/q1', 'main/q2', 'main'])
print(
'\nNumber of parameters: \t pi: %d, \t q1: %d, \t q2: %d, \t total: %d\n'
% var_counts)
# Bellman backup for Q functions, using Clipped Double-Q targets
min_q_targ = tf.minimum(q1_targ, q2_targ)
backup = tf.stop_gradient(r_ph + gamma * (1 - d_ph) * min_q_targ)
# TD3 losses
pi_loss = -tf.reduce_mean(q1_pi)
q1_loss = tf.reduce_mean((q1 - backup)**2)
q2_loss = tf.reduce_mean((q2 - backup)**2)
q_loss = q1_loss + q2_loss
# Separate train ops for pi, q
pi_optimizer = tf.train.AdamOptimizer(learning_rate=pi_lr)
q_optimizer = tf.train.AdamOptimizer(learning_rate=q_lr)
train_pi_op = pi_optimizer.minimize(pi_loss, var_list=get_vars('main/pi'))
train_q_op = q_optimizer.minimize(q_loss, var_list=get_vars('main/q'))
# Polyak averaging for target variables
target_update = tf.group([
tf.assign(v_targ, polyak * v_targ + (1 - polyak) * v_main)
for v_main, v_targ in zip(get_vars('main'), get_vars('target'))
])
# Initializing targets to match main variables
target_init = tf.group([
tf.assign(v_targ, v_main)
for v_main, v_targ in zip(get_vars('main'), get_vars('target'))
])
sess = tf.Session()
sess.run(tf.global_variables_initializer())
sess.run(target_init)
# Setup model saving
logger.setup_tf_saver(sess,
inputs={
'x': x_ph,
'a': a_ph
},
outputs={
'pi': pi,
'q1': q1,
'q2': q2
})
def get_action(o, noise_scale):
a = sess.run(pi, feed_dict={x_ph: o.reshape(1, -1)})[0]
a += noise_scale * np.random.randn(act_dim)
return np.clip(a, -act_limit, act_limit)
def test_agent():
for j in range(num_test_episodes):
o, d, ep_ret, ep_len = test_env.reset(), False, 0, 0
while not (d or (ep_len == max_ep_len)):
# Take deterministic actions at test time (noise_scale=0)
o, r, d, _ = test_env.step(get_action(o, 0))
ep_ret += r
ep_len += 1
logger.store(TestEpRet=ep_ret, TestEpLen=ep_len)
start_time = time.time()
o, ep_ret, ep_len = env.reset(), 0, 0
total_steps = steps_per_epoch * epochs
# Main loop: collect experience in env and update/log each epoch
for t in range(total_steps):
# Until start_steps have elapsed, randomly sample actions
# from a uniform distribution for better exploration. Afterwards,
# use the learned policy (with some noise, via act_noise).
if t > start_steps:
a = get_action(o, act_noise)
else:
a = env.action_space.sample()
# Step the env
o2, r, d, _ = env.step(a)
ep_ret += r
ep_len += 1
# Ignore the "done" signal if it comes from hitting the time
# horizon (that is, when it's an artificial terminal signal
# that isn't based on the agent's state)
d = False if ep_len == max_ep_len else d
# Store experience to replay buffer
replay_buffer.store(o, a, r, o2, d)
# Super critical, easy to overlook step: make sure to update
# most recent observation!
o = o2
# End of trajectory handling
if d or (ep_len == max_ep_len):
logger.store(EpRet=ep_ret, EpLen=ep_len)
o, ep_ret, ep_len = env.reset(), 0, 0
# Update handling
if t >= update_after and t % update_every == 0:
for j in range(update_every):
batch = replay_buffer.sample_batch(batch_size)
feed_dict = {
x_ph: batch['obs1'],
x2_ph: batch['obs2'],
a_ph: batch['acts'],
r_ph: batch['rews'],
d_ph: batch['done']
}
q_step_ops = [q_loss, q1, q2, train_q_op]
outs = sess.run(q_step_ops, feed_dict)
logger.store(LossQ=outs[0], Q1Vals=outs[1], Q2Vals=outs[2])
if j % policy_delay == 0:
# Delayed policy update
outs = sess.run([pi_loss, train_pi_op, target_update],
feed_dict)
logger.store(LossPi=outs[0])
# End of epoch wrap-up
if (t + 1) % steps_per_epoch == 0:
epoch = (t + 1) // steps_per_epoch
# Save model
if (epoch % save_freq == 0) or (epoch == epochs):
logger.save_state({'env': env}, None)
# Test the performance of the deterministic version of the agent.
test_agent()
# Log info about epoch
logger.log_tabular('Epoch', epoch)
logger.log_tabular('EpRet', with_min_and_max=True)
logger.log_tabular('TestEpRet', with_min_and_max=True)
logger.log_tabular('EpLen', average_only=True)
logger.log_tabular('TestEpLen', average_only=True)
logger.log_tabular('TotalEnvInteracts', t)
logger.log_tabular('Q1Vals', with_min_and_max=True)
logger.log_tabular('Q2Vals', with_min_and_max=True)
logger.log_tabular('LossPi', average_only=True)
logger.log_tabular('LossQ', average_only=True)
logger.log_tabular('Time', time.time() - start_time)
logger.dump_tabular()
def gail(env_fn,
actor_critic=ActorCritic,
ac_kwargs=dict(),
disc=Discriminator,
dc_kwargs=dict(),
seed=0,
episodes_per_epoch=40,
epochs=500,
gamma=0.99,
lam=0.97,
pi_lr=3e-3,
vf_lr=3e-3,
dc_lr=5e-4,
train_v_iters=80,
train_dc_iters=80,
max_ep_len=1000,
logger_kwargs=dict(),
save_freq=10):
l_lam = 0 # balance two loss term
logger = EpochLogger(**logger_kwargs)
logger.save_config(locals())
seed += 10000 * proc_id()
torch.manual_seed(seed)
np.random.seed(seed)
env = env_fn()
obs_dim = env.observation_space.shape
act_dim = env.action_space.shape
ac_kwargs['action_space'] = env.action_space
# Models
ac = actor_critic(input_dim=obs_dim[0], **ac_kwargs)
disc = disc(input_dim=obs_dim[0], **dc_kwargs)
# TODO: Load expert policy here
expert = actor_critic(input_dim=obs_dim[0], **ac_kwargs)
expert_name = "expert_torch_save.pt"
expert = torch.load(osp.join(logger_kwargs['output_dir'], expert_name))
# Buffers
local_episodes_per_epoch = int(episodes_per_epoch / num_procs())
buff_s = BufferS(obs_dim[0], act_dim[0], local_episodes_per_epoch,
max_ep_len)
buff_t = BufferT(obs_dim[0], act_dim[0], local_episodes_per_epoch,
max_ep_len)
# Count variables
var_counts = tuple(
count_vars(module) for module in [ac.policy, ac.value_f, disc.policy])
logger.log('\nNumber of parameters: \t pi: %d, \t v: %d, \t d: %d\n' %
var_counts)
# Optimizers
train_pi = torch.optim.Adam(ac.policy.parameters(), lr=pi_lr)
train_v = torch.optim.Adam(ac.value_f.parameters(), lr=vf_lr)
train_dc = torch.optim.Adam(disc.policy.parameters(), lr=dc_lr)
# Parameters Sync
sync_all_params(ac.parameters())
sync_all_params(disc.parameters())
def update(e):
obs_s, act, adv, ret, lgp_old = [
torch.Tensor(x) for x in buff_s.retrieve_all()
]
obs_t, _ = [torch.Tensor(x) for x in buff_t.retrieve_all()]
# Policy
_, lgp, _ = ac.policy(obs_s, act)
entropy = (-lgp).mean()
# Policy loss
# policy gradient term + entropy term
pi_loss = -(lgp * adv).mean() - l_lam * entropy
# Train policy
if e > 10:
train_pi.zero_grad()
pi_loss.backward()
average_gradients(train_pi.param_groups)
train_pi.step()
# Value function
v = ac.value_f(obs_s)
v_l_old = F.mse_loss(v, ret)
for _ in range(train_v_iters):
v = ac.value_f(obs_s)
v_loss = F.mse_loss(v, ret)
# Value function train
train_v.zero_grad()
v_loss.backward()
average_gradients(train_v.param_groups)
train_v.step()
# Discriminator
gt1 = torch.ones(obs_s.size()[0], dtype=torch.int)
gt2 = torch.zeros(obs_t.size()[0], dtype=torch.int)
_, lgp_s, _ = disc(obs_s, gt=gt1)
_, lgp_t, _ = disc(obs_t, gt=gt2)
dc_loss_old = -lgp_s.mean() - lgp_t.mean()
for _ in range(train_dc_iters):
_, lgp_s, _ = disc(obs_s, gt=gt1)
_, lgp_t, _ = disc(obs_t, gt=gt2)
dc_loss = -lgp_s.mean() - lgp_t.mean()
# Discriminator train
train_dc.zero_grad()
dc_loss.backward()
average_gradients(train_dc.param_groups)
train_dc.step()
_, lgp_s, _ = disc(obs_s, gt=gt1)
_, lgp_t, _ = disc(obs_t, gt=gt2)
dc_loss_new = -lgp_s.mean() - lgp_t.mean()
# Log the changes
_, lgp, _, v = ac(obs, act)
entropy_new = (-lgp).mean()
pi_loss_new = -(lgp * adv).mean() - l_lam * entropy
v_loss_new = F.mse_loss(v, ret)
kl = (lgp_old - lgp).mean()
logger.store(LossPi=pi_loss,
LossV=v_l_old,
LossDC=dc_loss_old,
DeltaLossPi=(pi_loss_new - pi_loss),
DeltaLossV=(v_loss_new - v_l_old),
DeltaLossDC=(dc_loss_new - dc_loss_old),
DeltaEnt=(entropy_new - entropy),
Entropy=entropy,
KL=kl)
start_time = time.time()
o, r, sdr, d, ep_ret, ep_sdr, ep_len = env.reset(), 0, 0, False, 0, 0, 0
total_t = 0
ep_len_t = 0
for epoch in range(epochs):
ac.eval()
disc.eval()
# We recognize the probability term of index [0] correspond to the teacher's policy
# Student's policy rollout
for _ in range(local_episodes_per_epoch):
for _ in range(max_ep_len):
obs = torch.Tensor(o.reshape(1, -1))
a, _, lopg_t, v_t = ac(obs)
buff_s.store(o,
a.detach().numpy(), r, sdr, v_t.item(),
lopg_t.detach().numpy())
logger.store(VVals=v_t)
o, r, d, _ = env.step(a.detach().numpy()[0])
_, sdr, _ = disc(torch.Tensor(o.reshape(1, -1)),
gt=torch.Tensor([0]))
if sdr < -4: # Truncate rewards
sdr = -4
ep_ret += r
ep_sdr += sdr
ep_len += 1
total_t += 1
terminal = d or (ep_len == max_ep_len)
if terminal:
buff_s.end_episode()
logger.store(EpRetS=ep_ret, EpLenS=ep_len, EpSdrS=ep_sdr)
o, r, sdr, d, ep_ret, ep_sdr, ep_len = env.reset(
), 0, 0, False, 0, 0, 0
# Teacher's policy rollout
for _ in range(local_episodes_per_epoch):
for _ in range(max_ep_len):
obs = torch.Tensor(o.reshape(1, -1))
a, _, _, _ = expert(obs)
buff_t.store(o, a.detach().numpy(), r)
o, r, d, _ = env.step(a.detach().numpy()[0])
ep_ret += r
ep_len += 1
total_t += 1
terminal = d or (ep_len == max_ep_len)
if terminal:
buff_t.end_episode()
logger.store(EpRetT=ep_ret, EpLenT=ep_len)
o, r, d, ep_ret, ep_len = env.reset(), 0, False, 0, 0
if (epoch % save_freq == 0) or (epoch == epochs - 1):
logger.save_state({'env': env}, [ac, disc], None)
# Update
ac.train()
disc.train()
update(epoch)
# Log
logger.log_tabular('Epoch', epoch)
logger.log_tabular('EpRetS', with_min_and_max=True)
logger.log_tabular('EpSdrS', with_min_and_max=True)
logger.log_tabular('EpLenS', average_only=True)
logger.log_tabular('EpRetT', with_min_and_max=True)
logger.log_tabular('EpLenT', average_only=True)
logger.log_tabular('VVals', with_min_and_max=True)
logger.log_tabular('TotalEnvInteracts', total_t)
logger.log_tabular('LossPi', average_only=True)
logger.log_tabular('DeltaLossPi', average_only=True)
logger.log_tabular('LossV', average_only=True)
logger.log_tabular('DeltaLossV', average_only=True)
logger.log_tabular('LossDC', average_only=True)
logger.log_tabular('DeltaLossDC', average_only=True)
logger.log_tabular('Entropy', average_only=True)
logger.log_tabular('DeltaEnt', average_only=True)
logger.log_tabular('KL', average_only=True)
logger.log_tabular('Time', time.time() - start_time)
logger.dump_tabular()
def vpg(env_fn, actor_critic=core.mlp_actor_critic, ac_kwargs=dict(), seed=0,
steps_per_epoch=4000, epochs=50, gamma=0.99, pi_lr=3e-4,
vf_lr=1e-3, train_v_iters=80, lam=0.97, max_ep_len=1000,
logger_kwargs=dict(), save_freq=10):
"""
Vanilla Policy Gradient
(with GAE-Lambda for advantage estimation)
Args:
env_fn : A function which creates a copy of the environment.
The environment must satisfy the OpenAI Gym API.
actor_critic: A function which takes in placeholder symbols
for state, ``x_ph``, and action, ``a_ph``, and returns the main
outputs from the agent's Tensorflow computation graph:
=========== ================ ======================================
Symbol Shape Description
=========== ================ ======================================
``pi`` (batch, act_dim) | Samples actions from policy given
| states.
``logp`` (batch,) | Gives log probability, according to
| the policy, of taking actions ``a_ph``
| in states ``x_ph``.
``logp_pi`` (batch,) | Gives log probability, according to
| the policy, of the action sampled by
| ``pi``.
``v`` (batch,) | Gives the value estimate for states
| in ``x_ph``. (Critical: make sure
| to flatten this!)
=========== ================ ======================================
ac_kwargs (dict): Any kwargs appropriate for the actor_critic
function you provided to VPG.
seed (int): Seed for random number generators.
steps_per_epoch (int): Number of steps of interaction (state-action pairs)
for the agent and the environment in each epoch.
epochs (int): Number of epochs of interaction (equivalent to
number of policy updates) to perform.
gamma (float): Discount factor. (Always between 0 and 1.)
pi_lr (float): Learning rate for policy optimizer.
vf_lr (float): Learning rate for value function optimizer.
train_v_iters (int): Number of gradient descent steps to take on
value function per epoch.
lam (float): Lambda for GAE-Lambda. (Always between 0 and 1,
close to 1.)
max_ep_len (int): Maximum length of trajectory / episode / rollout.
logger_kwargs (dict): Keyword args for EpochLogger.
save_freq (int): How often (in terms of gap between epochs) to save
the current policy and value function.
"""
logger = EpochLogger(**logger_kwargs)
logger.save_config(locals())
seed += 10000 * proc_id()
tf.set_random_seed(seed)
np.random.seed(seed)
env = env_fn()
obs_dim = env.observation_space.shape
act_dim = env.action_space.shape
# Share information about action space with policy architecture
ac_kwargs['action_space'] = env.action_space
# Inputs to computation graph
x_ph, a_ph = core.placeholders_from_spaces(env.observation_space, env.action_space)
adv_ph, ret_ph, logp_old_ph = core.placeholders(None, None, None)
# Main outputs from computation graph
pi, logp, logp_pi, v = actor_critic(x_ph, a_ph, **ac_kwargs)
# Need all placeholders in *this* order later (to zip with data from buffer)
all_phs = [x_ph, a_ph, adv_ph, ret_ph, logp_old_ph]
# Every step, get: action, value, and logprob
get_action_ops = [pi, v, logp_pi]
# Experience buffer
local_steps_per_epoch = int(steps_per_epoch / num_procs())
buf = VPGBuffer(obs_dim, act_dim, local_steps_per_epoch, gamma, lam)
# Count variables
var_counts = tuple(core.count_vars(scope) for scope in ['pi', 'v'])
logger.log('\nNumber of parameters: \t pi: %d, \t v: %d\n'%var_counts)
# VPG objectives
pi_loss = -tf.reduce_mean(logp * adv_ph)
v_loss = tf.reduce_mean((ret_ph - v)**2)
# Info (useful to watch during learning)
approx_kl = tf.reduce_mean(logp_old_ph - logp) # a sample estimate for KL-divergence, easy to compute
approx_ent = tf.reduce_mean(-logp) # a sample estimate for entropy, also easy to compute
# Optimizers
train_pi = MpiAdamOptimizer(learning_rate=pi_lr).minimize(pi_loss)
train_v = MpiAdamOptimizer(learning_rate=vf_lr).minimize(v_loss)
sess = tf.Session()
sess.run(tf.global_variables_initializer())
# Sync params across processes
sess.run(sync_all_params())
# Setup model saving
logger.setup_tf_saver(sess, inputs={'x': x_ph}, outputs={'pi': pi, 'v': v})
def update():
inputs = {k:v for k,v in zip(all_phs, buf.get())}
pi_l_old, v_l_old, ent = sess.run([pi_loss, v_loss, approx_ent], feed_dict=inputs)
# Policy gradient step
sess.run(train_pi, feed_dict=inputs)
# Value function learning
for _ in range(train_v_iters):
sess.run(train_v, feed_dict=inputs)
# Log changes from update
pi_l_new, v_l_new, kl = sess.run([pi_loss, v_loss, approx_kl], feed_dict=inputs)
logger.store(LossPi=pi_l_old, LossV=v_l_old,
KL=kl, Entropy=ent,
DeltaLossPi=(pi_l_new - pi_l_old),
DeltaLossV=(v_l_new - v_l_old))
start_time = time.time()
o, ep_ret, ep_len = env.reset(), 0, 0
# Main loop: collect experience in env and update/log each epoch
for epoch in range(epochs):
for t in range(local_steps_per_epoch):
a, v_t, logp_t = sess.run(get_action_ops, feed_dict={x_ph: o.reshape(1,-1)})
o2, r, d, _ = env.step(a[0])
ep_ret += r
ep_len += 1
# save and log
buf.store(o, a, r, v_t, logp_t)
logger.store(VVals=v_t)
# Update obs (critical!)
o = o2
terminal = d or (ep_len == max_ep_len)
if terminal or (t==local_steps_per_epoch-1):
if not(terminal):
print('Warning: trajectory cut off by epoch at %d steps.'%ep_len)
# if trajectory didn't reach terminal state, bootstrap value target
last_val = 0 if d else sess.run(v, feed_dict={x_ph: o.reshape(1,-1)})
buf.finish_path(last_val)
if terminal:
# only save EpRet / EpLen if trajectory finished
logger.store(EpRet=ep_ret, EpLen=ep_len)
o, ep_ret, ep_len = env.reset(), 0, 0
# Save model
if (epoch % save_freq == 0) or (epoch == epochs-1):
logger.save_state({'env': env}, None)
# Perform VPG update!
update()
# Log info about epoch
logger.log_tabular('Epoch', epoch)
logger.log_tabular('EpRet', with_min_and_max=True)
logger.log_tabular('EpLen', average_only=True)
logger.log_tabular('VVals', with_min_and_max=True)
logger.log_tabular('TotalEnvInteracts', (epoch+1)*steps_per_epoch)
logger.log_tabular('LossPi', average_only=True)
logger.log_tabular('LossV', average_only=True)
logger.log_tabular('DeltaLossPi', average_only=True)
logger.log_tabular('DeltaLossV', average_only=True)
logger.log_tabular('Entropy', average_only=True)
logger.log_tabular('KL', average_only=True)
logger.log_tabular('Time', time.time()-start_time)
logger.dump_tabular()
def sac(env_fn,
actor_critic=core.MLPActorCritic,
ac_kwargs=dict(),
seed=0,
steps_per_epoch=4000,
epochs=100,
replay_size=int(1e6),
gamma=0.99,
polyak=0.995,
lr=1e-3,
alpha=0.2,
batch_size=100,
start_steps=10000,
update_after=1000,
update_every=50,
num_test_episodes=10,
max_ep_len=1000,
logger_kwargs=dict(),
save_freq=1):
"""
Soft Actor-Critic (SAC)
Args:
env_fn : A function which creates a copy of the environment.
The environment must satisfy the OpenAI Gym API.
actor_critic: The constructor method for a PyTorch Module with an ``act``
method, a ``pi`` module, a ``q1`` module, and a ``q2`` module.
The ``act`` method and ``pi`` module should accept batches of
observations as inputs, and ``q1`` and ``q2`` should accept a batch
of observations and a batch of actions as inputs. When called,
``act``, ``q1``, and ``q2`` should return:
=========== ================ ======================================
Call Output Shape Description
=========== ================ ======================================
``act`` (batch, act_dim) | Numpy array of actions for each
| observation.
``q1`` (batch,) | Tensor containing one current estimate
| of Q* for the provided observations
| and actions. (Critical: make sure to
| flatten this!)
``q2`` (batch,) | Tensor containing the other current
| estimate of Q* for the provided observations
| and actions. (Critical: make sure to
| flatten this!)
=========== ================ ======================================
Calling ``pi`` should return:
=========== ================ ======================================
Symbol Shape Description
=========== ================ ======================================
``a`` (batch, act_dim) | Tensor containing actions from policy
| given observations.
``logp_pi`` (batch,) | Tensor containing log probabilities of
| actions in ``a``. Importantly: gradients
| should be able to flow back into ``a``.
=========== ================ ======================================
ac_kwargs (dict): Any kwargs appropriate for the ActorCritic object
you provided to SAC.
seed (int): Seed for random number generators.
steps_per_epoch (int): Number of steps of interaction (state-action pairs)
for the agent and the environment in each epoch.
epochs (int): Number of epochs to run and train agent.
replay_size (int): Maximum length of replay buffer.
gamma (float): Discount factor. (Always between 0 and 1.)
polyak (float): Interpolation factor in polyak averaging for target
networks. Target networks are updated towards main networks
according to:
.. math:: \\theta_{\\text{targ}} \\leftarrow
\\rho \\theta_{\\text{targ}} + (1-\\rho) \\theta
where :math:`\\rho` is polyak. (Always between 0 and 1, usually
close to 1.)
lr (float): Learning rate (used for both policy and value learning).
alpha (float): Entropy regularization coefficient. (Equivalent to
inverse of reward scale in the original SAC paper.)
batch_size (int): Minibatch size for SGD.
start_steps (int): Number of steps for uniform-random action selection,
before running real policy. Helps exploration.
update_after (int): Number of env interactions to collect before
starting to do gradient descent updates. Ensures replay buffer
is full enough for useful updates.
update_every (int): Number of env interactions that should elapse
between gradient descent updates. Note: Regardless of how long
you wait between updates, the ratio of env steps to gradient steps
is locked to 1.
num_test_episodes (int): Number of episodes to test the deterministic
policy at the end of each epoch.
max_ep_len (int): Maximum length of trajectory / episode / rollout.
logger_kwargs (dict): Keyword args for EpochLogger.
save_freq (int): How often (in terms of gap between epochs) to save
the current policy and value function.
"""
logger = EpochLogger(**logger_kwargs)
logger.save_config(locals())
torch.manual_seed(seed)
np.random.seed(seed)
env, test_env = env_fn(), env_fn()
obs_dim = env.observation_space.shape
act_dim = env.action_space.shape[0]
# Action limit for clamping: critically, assumes all dimensions share the same bound!
act_limit = env.action_space.high[0]
# Create actor-critic module and target networks
ac = actor_critic(env.observation_space, env.action_space, **ac_kwargs)
ac_targ = deepcopy(ac)
# Freeze target networks with respect to optimizers (only update via polyak averaging)
for p in ac_targ.parameters():
p.requires_grad = False
# List of parameters for both Q-networks (save this for convenience)
q_params = itertools.chain(ac.q1.parameters(), ac.q2.parameters())
# Experience buffer
replay_buffer = ReplayBuffer(obs_dim=obs_dim,
act_dim=act_dim,
size=replay_size)
# Count variables (protip: try to get a feel for how different size networks behave!)
var_counts = tuple(
core.count_vars(module) for module in [ac.pi, ac.q1, ac.q2])
logger.log('\nNumber of parameters: \t pi: %d, \t q1: %d, \t q2: %d\n' %
var_counts)
# Set up function for computing SAC Q-losses
def compute_loss_q(data):
o, a, r, o2, d = data['obs'], data['act'], data['rew'], data[
'obs2'], data['done']
q1 = ac.q1(o, a)
q2 = ac.q2(o, a)
# Bellman backup for Q functions
with torch.no_grad():
# Target actions come from *current* policy
a2, logp_a2 = ac.pi(o2)
# Target Q-values
q1_pi_targ = ac_targ.q1(o2, a2)
q2_pi_targ = ac_targ.q2(o2, a2)
q_pi_targ = torch.min(q1_pi_targ, q2_pi_targ)
backup = r + gamma * (1 - d) * (q_pi_targ - alpha * logp_a2)
# MSE loss against Bellman backup
loss_q1 = ((q1 - backup)**2).mean()
loss_q2 = ((q2 - backup)**2).mean()
loss_q = loss_q1 + loss_q2
# Useful info for logging
q_info = dict(Q1Vals=q1.detach().numpy(), Q2Vals=q2.detach().numpy())
return loss_q, q_info
# Set up function for computing SAC pi loss
def compute_loss_pi(data):
o = data['obs']
pi, logp_pi = ac.pi(o)
q1_pi = ac.q1(o, pi)
q2_pi = ac.q2(o, pi)
q_pi = torch.min(q1_pi, q2_pi)
# Entropy-regularized policy loss
loss_pi = (alpha * logp_pi - q_pi).mean()
# Useful info for logging
pi_info = dict(LogPi=logp_pi.detach().numpy())
return loss_pi, pi_info
# Set up optimizers for policy and q-function
pi_optimizer = Adam(ac.pi.parameters(), lr=lr)
q_optimizer = Adam(q_params, lr=lr)
# Set up model saving
logger.setup_pytorch_saver(ac)
def update(data):
# First run one gradient descent step for Q1 and Q2
q_optimizer.zero_grad()
loss_q, q_info = compute_loss_q(data)
loss_q.backward()
q_optimizer.step()
# Record things
logger.store(LossQ=loss_q.item(), **q_info)
# Freeze Q-networks so you don't waste computational effort
# computing gradients for them during the policy learning step.
for p in q_params:
p.requires_grad = False
# Next run one gradient descent step for pi.
pi_optimizer.zero_grad()
loss_pi, pi_info = compute_loss_pi(data)
loss_pi.backward()
pi_optimizer.step()
# Unfreeze Q-networks so you can optimize it at next DDPG step.
for p in q_params:
p.requires_grad = True
# Record things
logger.store(LossPi=loss_pi.item(), **pi_info)
# Finally, update target networks by polyak averaging.
with torch.no_grad():
for p, p_targ in zip(ac.parameters(), ac_targ.parameters()):
# NB: We use an in-place operations "mul_", "add_" to update target
# params, as opposed to "mul" and "add", which would make new tensors.
p_targ.data.mul_(polyak)
p_targ.data.add_((1 - polyak) * p.data)
def get_action(o, deterministic=False):
return ac.act(torch.as_tensor(o, dtype=torch.float32), deterministic)
def test_agent():
for j in range(num_test_episodes):
o, d, ep_ret, ep_len = test_env.reset(), False, 0, 0
while not (d or (ep_len == max_ep_len)):
# Take deterministic actions at test time
o, r, d, _ = test_env.step(get_action(o, True))
ep_ret += r
ep_len += 1
logger.store(TestEpRet=ep_ret, TestEpLen=ep_len)
# Prepare for interaction with environment
total_steps = steps_per_epoch * epochs
start_time = time.time()
o, ep_ret, ep_len = env.reset(), 0, 0
# Main loop: collect experience in env and update/log each epoch
for t in range(total_steps):
# Until start_steps have elapsed, randomly sample actions
# from a uniform distribution for better exploration. Afterwards,
# use the learned policy.
if t > start_steps:
a = get_action(o)
else:
a = env.action_space.sample()
# Step the env
o2, r, d, _ = env.step(a)
ep_ret += r
ep_len += 1
# Ignore the "done" signal if it comes from hitting the time
# horizon (that is, when it's an artificial terminal signal
# that isn't based on the agent's state)
d = False if ep_len == max_ep_len else d
# Store experience to replay buffer
replay_buffer.store(o, a, r, o2, d)
# Super critical, easy to overlook step: make sure to update
# most recent observation!
o = o2
# End of trajectory handling
if d or (ep_len == max_ep_len):
logger.store(EpRet=ep_ret, EpLen=ep_len)
o, ep_ret, ep_len = env.reset(), 0, 0
# Update handling
if t >= update_after and t % update_every == 0:
for j in range(update_every):
batch = replay_buffer.sample_batch(batch_size)
update(data=batch)
# End of epoch handling
if (t + 1) % steps_per_epoch == 0:
epoch = (t + 1) // steps_per_epoch
# Save model
if (epoch % save_freq == 0) or (epoch == epochs):
logger.save_state({'env': env}, None)
# Test the performance of the deterministic version of the agent.
test_agent()
# Log info about epoch
logger.log_tabular('Epoch', epoch)
logger.log_tabular('EpRet', with_min_and_max=True)
logger.log_tabular('TestEpRet', with_min_and_max=True)
logger.log_tabular('EpLen', average_only=True)
logger.log_tabular('TestEpLen', average_only=True)
logger.log_tabular('TotalEnvInteracts', t)
logger.log_tabular('Q1Vals', with_min_and_max=True)
logger.log_tabular('Q2Vals', with_min_and_max=True)
logger.log_tabular('LogPi', with_min_and_max=True)
logger.log_tabular('LossPi', average_only=True)
logger.log_tabular('LossQ', average_only=True)
logger.log_tabular('Time', time.time() - start_time)
logger.dump_tabular()
def trpo(env_fn, actor_critic=core.mlp_actor_critic, ac_kwargs=dict(), seed=0,
steps_per_epoch=4000, epochs=50, gamma=0.99, delta=0.01, vf_lr=1e-3,
train_v_iters=80, damping_coeff=0.1, cg_iters=10, backtrack_iters=10,
backtrack_coeff=0.8, lam=0.97, max_ep_len=1000, logger_kwargs=dict(),
save_freq=10, algo='trpo'):
"""
Trust Region Policy Optimization
(with support for Natural Policy Gradient)
Args:
env_fn : A function which creates a copy of the environment.
The environment must satisfy the OpenAI Gym API.
actor_critic: A function which takes in placeholder symbols
for state, ``x_ph``, and action, ``a_ph``, and returns the main
outputs from the agent's Tensorflow computation graph:
============ ================ ========================================
Symbol Shape Description
============ ================ ========================================
``pi`` (batch, act_dim) | Samples actions from policy given
| states.
``logp`` (batch,) | Gives log probability, according to
| the policy, of taking actions ``a_ph``
| in states ``x_ph``.
``logp_pi`` (batch,) | Gives log probability, according to
| the policy, of the action sampled by
| ``pi``.
``info`` N/A | A dict of any intermediate quantities
| (from calculating the policy or log
| probabilities) which are needed for
| analytically computing KL divergence.
| (eg sufficient statistics of the
| distributions)
``info_phs`` N/A | A dict of placeholders for old values
| of the entries in ``info``.
``d_kl`` () | A symbol for computing the mean KL
| divergence between the current policy
| (``pi``) and the old policy (as
| specified by the inputs to
| ``info_phs``) over the batch of
| states given in ``x_ph``.
``v`` (batch,) | Gives the value estimate for states
| in ``x_ph``. (Critical: make sure
| to flatten this!)
============ ================ ========================================
ac_kwargs (dict): Any kwargs appropriate for the actor_critic
function you provided to TRPO.
seed (int): Seed for random number generators.
steps_per_epoch (int): Number of steps of interaction (state-action pairs)
for the agent and the environment in each epoch.
epochs (int): Number of epochs of interaction (equivalent to
number of policy updates) to perform.
gamma (float): Discount factor. (Always between 0 and 1.)
delta (float): KL-divergence limit for TRPO / NPG update.
(Should be small for stability. Values like 0.01, 0.05.)
vf_lr (float): Learning rate for value function optimizer.
train_v_iters (int): Number of gradient descent steps to take on
value function per epoch.
damping_coeff (float): Artifact for numerical stability, should be
smallish. Adjusts Hessian-vector product calculation:
.. math:: Hv \\rightarrow (\\alpha I + H)v
where :math:`\\alpha` is the damping coefficient.
Probably don't play with this hyperparameter.
cg_iters (int): Number of iterations of conjugate gradient to perform.
Increasing this will lead to a more accurate approximation
to :math:`H^{-1} g`, and possibly slightly-improved performance,
but at the cost of slowing things down.
Also probably don't play with this hyperparameter.
backtrack_iters (int): Maximum number of steps allowed in the
backtracking line search. Since the line search usually doesn't
backtrack, and usually only steps back once when it does, this
hyperparameter doesn't often matter.
backtrack_coeff (float): How far back to step during backtracking line
search. (Always between 0 and 1, usually above 0.5.)
lam (float): Lambda for GAE-Lambda. (Always between 0 and 1,
close to 1.)
max_ep_len (int): Maximum length of trajectory / episode / rollout.
logger_kwargs (dict): Keyword args for EpochLogger.
save_freq (int): How often (in terms of gap between epochs) to save
the current policy and value function.
algo: Either 'trpo' or 'npg': this code supports both, since they are
almost the same.
"""
logger = EpochLogger(**logger_kwargs)
logger.save_config(locals())
seed += 10000 * proc_id()
tf.set_random_seed(seed)
np.random.seed(seed)
env = env_fn()
obs_dim = env.observation_space.shape
act_dim = env.action_space.shape
# Share information about action space with policy architecture
ac_kwargs['action_space'] = env.action_space
# Inputs to computation graph
x_ph, a_ph = core.placeholders_from_spaces(env.observation_space, env.action_space)
adv_ph, ret_ph, logp_old_ph = core.placeholders(None, None, None)
# Main outputs from computation graph, plus placeholders for old pdist (for KL)
pi, logp, logp_pi, info, info_phs, d_kl, v = actor_critic(x_ph, a_ph, **ac_kwargs)
# Need all placeholders in *this* order later (to zip with data from buffer)
all_phs = [x_ph, a_ph, adv_ph, ret_ph, logp_old_ph] + core.values_as_sorted_list(info_phs)
# Every step, get: action, value, logprob, & info for pdist (for computing kl div)
get_action_ops = [pi, v, logp_pi] + core.values_as_sorted_list(info)
# Experience buffer
local_steps_per_epoch = int(steps_per_epoch / num_procs())
info_shapes = {k: v.shape.as_list()[1:] for k,v in info_phs.items()}
buf = GAEBuffer(obs_dim, act_dim, local_steps_per_epoch, info_shapes, gamma, lam)
# Count variables
var_counts = tuple(core.count_vars(scope) for scope in ['pi', 'v'])
logger.log('\nNumber of parameters: \t pi: %d, \t v: %d\n'%var_counts)
# TRPO losses
ratio = tf.exp(logp - logp_old_ph) # pi(a|s) / pi_old(a|s)
pi_loss = -tf.reduce_mean(ratio * adv_ph)
v_loss = tf.reduce_mean((ret_ph - v)**2)
# Optimizer for value function
train_vf = MpiAdamOptimizer(learning_rate=vf_lr).minimize(v_loss)
# Symbols needed for CG solver
pi_params = core.get_vars('pi')
gradient = core.flat_grad(pi_loss, pi_params)
v_ph, hvp = core.hessian_vector_product(d_kl, pi_params)
if damping_coeff > 0:
hvp += damping_coeff * v_ph
# Symbols for getting and setting params
get_pi_params = core.flat_concat(pi_params)
set_pi_params = core.assign_params_from_flat(v_ph, pi_params)
sess = tf.Session()
sess.run(tf.global_variables_initializer())
# Sync params across processes
sess.run(sync_all_params())
# Setup model saving
logger.setup_tf_saver(sess, inputs={'x': x_ph}, outputs={'pi': pi, 'v': v})
def cg(Ax, b):
"""
Conjugate gradient algorithm
(see https://en.wikipedia.org/wiki/Conjugate_gradient_method)
"""
x = np.zeros_like(b)
r = b.copy() # Note: should be 'b - Ax(x)', but for x=0, Ax(x)=0. Change if doing warm start.
p = r.copy()
r_dot_old = np.dot(r,r)
for _ in range(cg_iters):
z = Ax(p)
alpha = r_dot_old / (np.dot(p, z) + EPS)
x += alpha * p
r -= alpha * z
r_dot_new = np.dot(r,r)
p = r + (r_dot_new / r_dot_old) * p
r_dot_old = r_dot_new
return x
def update():
# Prepare hessian func, gradient eval
inputs = {k:v for k,v in zip(all_phs, buf.get())}
Hx = lambda x : mpi_avg(sess.run(hvp, feed_dict={**inputs, v_ph: x}))
g, pi_l_old, v_l_old = sess.run([gradient, pi_loss, v_loss], feed_dict=inputs)
g, pi_l_old = mpi_avg(g), mpi_avg(pi_l_old)
# Core calculations for TRPO or NPG
x = cg(Hx, g)
alpha = np.sqrt(2*delta/(np.dot(x, Hx(x))+EPS))
old_params = sess.run(get_pi_params)
def set_and_eval(step):
sess.run(set_pi_params, feed_dict={v_ph: old_params - alpha * x * step})
return mpi_avg(sess.run([d_kl, pi_loss], feed_dict=inputs))
if algo=='npg':
# npg has no backtracking or hard kl constraint enforcement
kl, pi_l_new = set_and_eval(step=1.)
elif algo=='trpo':
# trpo augments npg with backtracking line search, hard kl
for j in range(backtrack_iters):
kl, pi_l_new = set_and_eval(step=backtrack_coeff**j)
if kl <= delta and pi_l_new <= pi_l_old:
logger.log('Accepting new params at step %d of line search.'%j)
logger.store(BacktrackIters=j)
break
if j==backtrack_iters-1:
logger.log('Line search failed! Keeping old params.')
logger.store(BacktrackIters=j)
kl, pi_l_new = set_and_eval(step=0.)
# Value function updates
for _ in range(train_v_iters):
sess.run(train_vf, feed_dict=inputs)
v_l_new = sess.run(v_loss, feed_dict=inputs)
# Log changes from update
logger.store(LossPi=pi_l_old, LossV=v_l_old, KL=kl,
DeltaLossPi=(pi_l_new - pi_l_old),
DeltaLossV=(v_l_new - v_l_old))
start_time = time.time()
o, ep_ret, ep_len = env.reset(), 0, 0
# Main loop: collect experience in env and update/log each epoch
for epoch in range(epochs):
for t in range(local_steps_per_epoch):
agent_outs = sess.run(get_action_ops, feed_dict={x_ph: o.reshape(1,-1)})
a, v_t, logp_t, info_t = agent_outs[0][0], agent_outs[1], agent_outs[2], agent_outs[3:]
o2, r, d, _ = env.step(a)
ep_ret += r
ep_len += 1
# save and log
buf.store(o, a, r, v_t, logp_t, info_t)
logger.store(VVals=v_t)
# Update obs (critical!)
o = o2
terminal = d or (ep_len == max_ep_len)
if terminal or (t==local_steps_per_epoch-1):
if not(terminal):
print('Warning: trajectory cut off by epoch at %d steps.'%ep_len)
# if trajectory didn't reach terminal state, bootstrap value target
last_val = 0 if d else sess.run(v, feed_dict={x_ph: o.reshape(1,-1)})
buf.finish_path(last_val)
if terminal:
# only save EpRet / EpLen if trajectory finished
logger.store(EpRet=ep_ret, EpLen=ep_len)
o, ep_ret, ep_len = env.reset(), 0, 0
# Save model
if (epoch % save_freq == 0) or (epoch == epochs-1):
logger.save_state({'env': env}, None)
# Perform TRPO or NPG update!
update()
# Log info about epoch
logger.log_tabular('Epoch', epoch)
logger.log_tabular('EpRet', with_min_and_max=True)
logger.log_tabular('EpLen', average_only=True)
logger.log_tabular('VVals', with_min_and_max=True)
logger.log_tabular('TotalEnvInteracts', (epoch+1)*steps_per_epoch)
logger.log_tabular('LossPi', average_only=True)
logger.log_tabular('LossV', average_only=True)
logger.log_tabular('DeltaLossPi', average_only=True)
logger.log_tabular('DeltaLossV', average_only=True)
logger.log_tabular('KL', average_only=True)
if algo=='trpo':
logger.log_tabular('BacktrackIters', average_only=True)
logger.log_tabular('Time', time.time()-start_time)
logger.dump_tabular()
def train(env_fn,
seed=0,
ppo_epoch=10,
steps_per_epoch=2048,
mini_batch_size=64,
num_epoch=1500,
gamma=0.99,
clip_ratio=0.2,
value_clip_ratio=10,
value_loss_coef=0.5,
entropy_loss_coef=0,
use_value_clipped_loss=True,
lr=3e-4,
eps=1e-5,
lam=0.95,
max_grad_norm=0.5,
max_ep_len=1000,
save_freq=10,
device=torch.device('cpu'),
ac_kwargs=dict(),
logger_kwargs=dict()):
# Set up logger and save configuration
logger = EpochLogger(**logger_kwargs)
logger.save_config(locals())
# Random seed
torch.manual_seed(seed)
np.random.seed(seed)
# Instantiate environment
env = env_fn()
env.seed(seed)
obs_dim = env.observation_space.shape
act_dim = env.action_space.shape
actor_critic = MLPActorCritic(env.observation_space, env.action_space,
**ac_kwargs).to(device)
ppo = PPO(actor_critic, clip_ratio, value_clip_ratio, ppo_epoch,
mini_batch_size, value_loss_coef, entropy_loss_coef, lr, eps,
max_grad_norm, use_value_clipped_loss)
# Set up experience buffer
buf = PPOBuffer(obs_dim, act_dim, steps_per_epoch, gamma, lam, device)
# Set up model saving
logger.setup_pytorch_saver(ppo.actor_critic)
# Prepare for interaction with environment
start_time = time.time()
running_state = ZFilter((obs_dim[0], ), clip=10)
# running_reward = ZFilter((1,), demean=False, clip=10)
obs, ep_ret, ep_len = env.reset(), 0, 0
obs = running_state(obs)
# Main loop: collect experience in env and update/log each epoch
for epoch in range(num_epoch):
for t in range(steps_per_epoch):
action, value, logp = ppo.actor_critic.step(
torch.as_tensor(obs, dtype=torch.float32).to(device))
next_obs, rew, done, _ = env.step(action)
next_obs = running_state(next_obs)
# rew = running_reward([rew])[0]
ep_ret += rew
ep_len += 1
# save and log
buf.store(obs, action, rew, value, logp)
# Update obs (critical!)
obs = next_obs
timeout = ep_len == max_ep_len
terminal = done or timeout
epoch_ended = t == steps_per_epoch - 1
if terminal or epoch_ended:
if epoch_ended and not (terminal):
print('Warning: trajectory cut off by epoch at %d steps.' %
ep_len,
flush=True)
# if trajectory didn't reach terminal state, bootstrap value target
if timeout or epoch_ended:
_, value, _ = ppo.actor_critic.step(
torch.as_tensor(obs, dtype=torch.float32).to(device))
else:
value = 0
buf.finish_path(value)
if terminal:
# only save EpRet / EpLen if trajectory finished
logger.store(EpRet=ep_ret, EpLen=ep_len)
obs, ep_ret, ep_len = env.reset(), 0, 0
obs = running_state(obs)
# Save model
if save_freq != 0 and ((epoch % save_freq == 0) or
(epoch == num_epoch - 1)):
logger.save_state({'env': env}, None)
# perform update
data = buf.get()
policy_loss, value_loss, entropy, kl = ppo.update(data)
# Log info about epoch
logger.log_tabular('Epoch', epoch)
logger.log_tabular('EpRet', with_min_and_max=True)
logger.log_tabular('EpLen', average_only=True)
logger.log_tabular('TotalEnvInteracts', (epoch + 1) * steps_per_epoch)
logger.log_tabular('LossPi', policy_loss)
logger.log_tabular('LossV', value_loss)
logger.log_tabular('Entropy', entropy)
logger.log_tabular('KL', kl)
logger.log_tabular('Time', time.time() - start_time)
logger.dump_tabular()
def train(env_fn,
env_name,
ac_kwargs=dict(),
seed=0,
steps_per_epoch=1000,
epochs=3000,
replay_size=int(1e6),
gamma=0.99,
polyak=0.995,
lr=3e-4,
batch_size=64,
start_steps=10000,
update_after=10000,
update_every=1,
num_test_episodes=10,
value_coef=0.5,
entropy_coef=0.02,
max_ep_len=1000,
logger_kwargs=dict(),
save_freq=10,
device=torch.device('cpu')):
logger = EpochLogger(**logger_kwargs)
logger.save_config(locals())
torch.manual_seed(seed)
np.random.seed(seed)
env, test_env = env_fn(), env_fn()
env.seed(seed)
test_env.seed(seed)
obs_dim = env.observation_space.shape
act_dim = env.action_space.shape[0]
# Action limit for clamping: critically, assumes all dimensions share the same bound!
act_limit = env.action_space.high[0]
actor_critic = MLPActorCritic(env.observation_space, env.action_space,
**ac_kwargs).to(device)
sql = SQL(actor_critic, lr, batch_size, update_every, gamma, polyak,
value_coef, entropy_coef)
# Experience buffer
replay_buffer = ReplayBuffer(obs_dim=obs_dim,
act_dim=act_dim,
size=replay_size,
device=device)
rewards_log = []
episode_rewards = deque(maxlen=10)
# Set up model saving
logger.setup_pytorch_saver(sql.actor_critic)
def test_agent():
for j in range(num_test_episodes):
o, d, ep_ret, ep_len = test_env.reset(), False, 0, 0
while not (d or (ep_len == max_ep_len)):
action = sql.actor_critic.act(
torch.as_tensor(o, dtype=torch.float32).to(device))
o, r, d, _ = test_env.step(action)
ep_ret += r
ep_len += 1
logger.store(TestEpRet=ep_ret, TestEpLen=ep_len)
episode_rewards.append(ep_ret)
# Prepare for interaction with environment
total_steps = steps_per_epoch * epochs
start_time = time.time()
o, ep_ret, ep_len = env.reset(), 0, 0
# Main loop: collect experience in env and update/log each epoch
for t in range(total_steps):
# Until start_steps have elapsed, randomly sample actions
# from a uniform distribution for better exploration. Afterwards,
# use the learned policy (with some noise, via act_noise).
if t > start_steps:
a = sql.actor_critic.act(
torch.as_tensor(o, dtype=torch.float32).to(device))
else:
a = env.action_space.sample()
# Step the env
o2, r, d, _ = env.step(a)
ep_ret += r
ep_len += 1
# Ignore the "done" signal if it comes from hitting the time
# horizon (that is, when it's an artificial terminal signal
# that isn't based on the agent's state)
d = False if ep_len == max_ep_len else d
# Store experience to replay buffer
replay_buffer.store(o, a, r, o2, d)
# Super critical, easy to overlook step: make sure to update
# most recent observation!
o = o2
# End of trajectory handling
if d or (ep_len == max_ep_len):
logger.store(EpRet=ep_ret, EpLen=ep_len)
o, ep_ret, ep_len = env.reset(), 0, 0
# Update handling
if t >= update_after:
for _ in range(update_every):
batch = replay_buffer.sample_batch(batch_size)
loss = sql.update(data=batch)
logger.store(Loss=loss)
else:
logger.store(Loss=0.)
# End of epoch handling
if (t + 1) % steps_per_epoch == 0:
epoch = (t + 1) // steps_per_epoch
# Save model
if save_freq != 0 and ((epoch % save_freq == 0) or
(epoch == epochs)):
logger.save_state({'env': env}, None)
# Test the performance of the deterministic version of the agent.
test_agent()
rewards_log.append(np.mean(episode_rewards))
# Log info about epoch
logger.log_tabular('Epoch', epoch)
logger.log_tabular('EpRet', with_min_and_max=True)
logger.log_tabular('TestEpRet', with_min_and_max=True)
logger.log_tabular('EpLen', average_only=True)
logger.log_tabular('TestEpLen', average_only=True)
logger.log_tabular('TotalEnvInteracts', t)
logger.log_tabular('Time', time.time() - start_time)
logger.log_tabular('Loss', average_only=True)
logger.dump_tabular()
rewards_log = np.array(rewards_log)
save_path = '../../log/modified_sql/' + env_name + '/' + str(seed) + '.npy'
np.save(save_path, rewards_log)
def valor(env_fn, actor_critic=ActorCritic, ac_kwargs=dict(), disc=Discriminator, dc_kwargs=dict(), seed=0, episodes_per_epoch=40,
epochs=50, gamma=0.99, pi_lr=3e-4, vf_lr=1e-3, dc_lr=5e-4, train_v_iters=80, train_dc_iters=10, train_dc_interv=10,
lam=0.97, max_ep_len=1000, logger_kwargs=dict(), con_dim=5, save_freq=10, k=1):
logger = EpochLogger(**logger_kwargs)
logger.save_config(locals())
seed += 10000 * proc_id()
torch.manual_seed(seed)
np.random.seed(seed)
env = env_fn()
obs_dim = env.observation_space.shape
act_dim = env.action_space.shape
ac_kwargs['action_space'] = env.action_space
# Model
actor_critic = actor_critic(input_dim=obs_dim[0]+con_dim, **ac_kwargs)
disc = disc(input_dim=obs_dim[0], context_dim=con_dim, **dc_kwargs)
# Buffer
local_episodes_per_epoch = int(episodes_per_epoch / num_procs())
buffer = Buffer(con_dim, obs_dim[0], act_dim[0], local_episodes_per_epoch, max_ep_len, train_dc_interv)
# Count variables
var_counts = tuple(count_vars(module) for module in
[actor_critic.policy, actor_critic.value_f, disc.policy])
logger.log('\nNumber of parameters: \t pi: %d, \t v: %d, \t d: %d\n'%var_counts)
# Optimizers
train_pi = torch.optim.Adam(actor_critic.policy.parameters(), lr=pi_lr)
train_v = torch.optim.Adam(actor_critic.value_f.parameters(), lr=vf_lr)
train_dc = torch.optim.Adam(disc.policy.parameters(), lr=dc_lr)
# Parameters Sync
sync_all_params(actor_critic.parameters())
sync_all_params(disc.parameters())
def update(e):
obs, act, adv, pos, ret, logp_old = [torch.Tensor(x) for x in buffer.retrieve_all()]
# Policy
_, logp, _ = actor_critic.policy(obs, act)
entropy = (-logp).mean()
# Policy loss
pi_loss = -(logp*(k*adv+pos)).mean()
# Train policy
train_pi.zero_grad()
pi_loss.backward()
average_gradients(train_pi.param_groups)
train_pi.step()
# Value function
v = actor_critic.value_f(obs)
v_l_old = F.mse_loss(v, ret)
for _ in range(train_v_iters):
v = actor_critic.value_f(obs)
v_loss = F.mse_loss(v, ret)
# Value function train
train_v.zero_grad()
v_loss.backward()
average_gradients(train_v.param_groups)
train_v.step()
# Discriminator
if (e+1) % train_dc_interv == 0:
print('Discriminator Update!')
con, s_diff = [torch.Tensor(x) for x in buffer.retrieve_dc_buff()]
_, logp_dc, _ = disc(s_diff, con)
d_l_old = -logp_dc.mean()
# Discriminator train
for _ in range(train_dc_iters):
_, logp_dc, _ = disc(s_diff, con)
d_loss = -logp_dc.mean()
train_dc.zero_grad()
d_loss.backward()
average_gradients(train_dc.param_groups)
train_dc.step()
_, logp_dc, _ = disc(s_diff, con)
dc_l_new = -logp_dc.mean()
else:
d_l_old = 0
dc_l_new = 0
# Log the changes
_, logp, _, v = actor_critic(obs, act)
pi_l_new = -(logp*(k*adv+pos)).mean()
v_l_new = F.mse_loss(v, ret)
kl = (logp_old - logp).mean()
logger.store(LossPi=pi_loss, LossV=v_l_old, KL=kl, Entropy=entropy, DeltaLossPi=(pi_l_new-pi_loss),
DeltaLossV=(v_l_new-v_l_old), LossDC=d_l_old, DeltaLossDC=(dc_l_new-d_l_old))
# logger.store(Adv=adv.reshape(-1).numpy().tolist(), Pos=pos.reshape(-1).numpy().tolist())
start_time = time.time()
o, r, d, ep_ret, ep_len = env.reset(), 0, False, 0, 0
context_dist = Categorical(logits=torch.Tensor(np.ones(con_dim)))
total_t = 0
for epoch in range(epochs):
actor_critic.eval()
disc.eval()
for _ in range(local_episodes_per_epoch):
c = context_dist.sample()
c_onehot = F.one_hot(c, con_dim).squeeze().float()
for _ in range(max_ep_len):
concat_obs = torch.cat([torch.Tensor(o.reshape(1, -1)), c_onehot.reshape(1, -1)], 1)
a, _, logp_t, v_t = actor_critic(concat_obs)
buffer.store(c, concat_obs.squeeze().detach().numpy(), a.detach().numpy(), r, v_t.item(), logp_t.detach().numpy())
logger.store(VVals=v_t)
o, r, d, _ = env.step(a.detach().numpy()[0])
ep_ret += r
ep_len += 1
total_t += 1
terminal = d or (ep_len == max_ep_len)
if terminal:
dc_diff = torch.Tensor(buffer.calc_diff()).unsqueeze(0)
con = torch.Tensor([float(c)]).unsqueeze(0)
_, _, log_p = disc(dc_diff, con)
buffer.end_episode(log_p.detach().numpy())
logger.store(EpRet=ep_ret, EpLen=ep_len)
o, r, d, ep_ret, ep_len = env.reset(), 0, False, 0, 0
if (epoch % save_freq == 0) or (epoch == epochs - 1):
logger.save_state({'env': env}, [actor_critic, disc], None)
# Update
actor_critic.train()
disc.train()
update(epoch)
# Log
logger.log_tabular('Epoch', epoch)
logger.log_tabular('EpRet', with_min_and_max=True)
logger.log_tabular('EpLen', average_only=True)
logger.log_tabular('VVals', with_min_and_max=True)
logger.log_tabular('TotalEnvInteracts', total_t)
logger.log_tabular('LossPi', average_only=True)
logger.log_tabular('DeltaLossPi', average_only=True)
logger.log_tabular('LossV', average_only=True)
logger.log_tabular('DeltaLossV', average_only=True)
logger.log_tabular('LossDC', average_only=True)
logger.log_tabular('DeltaLossDC', average_only=True)
logger.log_tabular('Entropy', average_only=True)
logger.log_tabular('KL', average_only=True)
logger.log_tabular('Time', time.time()-start_time)
logger.dump_tabular()
def ddpg(trajectory_output_path, env_fn, actor_critic=core.MLPActorCritic, ac_kwargs=dict(), seed=0,
steps_per_epoch=4000, epochs=100, replay_size=int(1e6), gamma=0.99,
polyak=0.995, pi_lr=1e-3, q_lr=1e-3, batch_size=100, start_steps=10000,
update_after=1000, update_every=50, act_noise=0.1, num_test_episodes=10,
max_ep_len=1000, logger_kwargs=dict(), save_freq=1):
"""
Deep Deterministic Policy Gradient (DDPG)
trajectory_output_path: The path to output the trajectory results
Args:
env_fn : A function which creates a copy of the environment.
The environment must satisfy the OpenAI Gym API.
actor_critic: The constructor method for a PyTorch Module with an ``act``
method, a ``pi`` module, and a ``q`` module. The ``act`` method and
``pi`` module should accept batches of observations as inputs,
and ``q`` should accept a batch of observations and a batch of
actions as inputs. When called, these should return:
=========== ================ ======================================
Call Output Shape Description
=========== ================ ======================================
``act`` (batch, act_dim) | Numpy array of actions for each
| observation.
``pi`` (batch, act_dim) | Tensor containing actions from policy
| given observations.
``q`` (batch,) | Tensor containing the current estimate
| of Q* for the provided observations
| and actions. (Critical: make sure to
| flatten this!)
=========== ================ ======================================
ac_kwargs (dict): Any kwargs appropriate for the ActorCritic object
you provided to DDPG.
seed (int): Seed for random number generators.
steps_per_epoch (int): Number of steps of interaction (state-action pairs)
for the agent and the environment in each epoch.
epochs (int): Number of epochs to run and train agent.
replay_size (int): Maximum length of replay buffer.
gamma (float): Discount factor. (Always between 0 and 1.)
polyak (float): Interpolation factor in polyak averaging for target
networks. Target networks are updated towards main networks
according to:
.. math:: \\theta_{\\text{targ}} \\leftarrow
\\rho \\theta_{\\text{targ}} + (1-\\rho) \\theta
where :math:`\\rho` is polyak. (Always between 0 and 1, usually
close to 1.)
pi_lr (float): Learning rate for policy.
q_lr (float): Learning rate for Q-networks.
batch_size (int): Minibatch size for SGD.
start_steps (int): Number of steps for uniform-random action selection,
before running real policy. Helps exploration.
update_after (int): Number of env interactions to collect before
starting to do gradient descent updates. Ensures replay buffer
is full enough for useful updates.
update_every (int): Number of env interactions that should elapse
between gradient descent updates. Note: Regardless of how long
you wait between updates, the ratio of env steps to gradient steps
is locked to 1.
act_noise (float): Stddev for Gaussian exploration noise added to
policy at training time. (At test time, no noise is added.)
num_test_episodes (int): Number of episodes to test the deterministic
policy at the end of each epoch.
max_ep_len (int): Maximum length of trajectory / episode / rollout.
logger_kwargs (dict): Keyword args for EpochLogger.
save_freq (int): How often (in terms of gap between epochs) to save
the current policy and value function.
"""
logger = EpochLogger(**logger_kwargs)
logger.save_config(locals())
torch.manual_seed(seed)
np.random.seed(seed)
env, test_env = env_fn(), env_fn()
obs_dim = env.observation_space.shape
act_dim = env.action_space.shape[0]
# Action limit for clamping: critically, assumes all dimensions share the same bound!
act_limit = env.action_space.high[0]
# Create actor-critic module and target networks
ac = actor_critic(env.observation_space, env.action_space, **ac_kwargs)
ac_targ = deepcopy(ac)
# Freeze target networks with respect to optimizers (only update via polyak averaging)
for p in ac_targ.parameters():
p.requires_grad = False
# Experience buffer
replay_buffer = ReplayBuffer(obs_dim=obs_dim, act_dim=act_dim, size=replay_size)
# Count variables (protip: try to get a feel for how different size networks behave!)
var_counts = tuple(core.count_vars(module) for module in [ac.pi, ac.q])
logger.log('\nNumber of parameters: \t pi: %d, \t q: %d\n' % var_counts)
# Set up function for computing DDPG Q-loss
def compute_loss_q(data):
o, a, r, o2, d = data['obs'], data['act'], data['rew'], data['obs2'], data['done']
q = ac.q(o, a)
# Bellman backup for Q function
with torch.no_grad():
q_pi_targ = ac_targ.q(o2, ac_targ.pi(o2))
backup = r + gamma * (1 - d) * q_pi_targ
# MSE loss against Bellman backup
loss_q = ((q - backup) ** 2).mean()
# Useful info for logging
loss_info = dict(QVals=q.detach().numpy())
return loss_q, loss_info
# Set up function for computing DDPG pi loss
def compute_loss_pi(data):
o = data['obs']
q_pi = ac.q(o, ac.pi(o))
return -q_pi.mean()
# Set up optimizers for policy and q-function
pi_optimizer = Adam(ac.pi.parameters(), lr=pi_lr)
q_optimizer = Adam(ac.q.parameters(), lr=q_lr)
# Set up model saving
logger.setup_pytorch_saver(ac)
def update(data):
# First run one gradient descent step for Q.
q_optimizer.zero_grad()
loss_q, loss_info = compute_loss_q(data)
loss_q.backward()
q_optimizer.step()
# Freeze Q-network so you don't waste computational effort
# computing gradients for it during the policy learning step.
for p in ac.q.parameters():
p.requires_grad = False
# Next run one gradient descent step for pi.
pi_optimizer.zero_grad()
loss_pi = compute_loss_pi(data)
loss_pi.backward()
pi_optimizer.step()
# Unfreeze Q-network so you can optimize it at next DDPG step.
for p in ac.q.parameters():
p.requires_grad = True
# Record things
logger.store(LossQ=loss_q.item(), LossPi=loss_pi.item(), **loss_info)
# Finally, update target networks by polyak averaging.
with torch.no_grad():
for p, p_targ in zip(ac.parameters(), ac_targ.parameters()):
# NB: We use an in-place operations "mul_", "add_" to update target
# params, as opposed to "mul" and "add", which would make new tensors.
p_targ.data.mul_(polyak)
p_targ.data.add_((1 - polyak) * p.data)
def get_action(o, noise_scale):
a = ac.act(torch.as_tensor(o, dtype=torch.float32))
a += noise_scale * np.random.randn(act_dim)
return np.clip(a, -act_limit, act_limit)
# def test_agent():
# for j in range(num_test_episodes):
# o, d, ep_ret, ep_len = test_env.reset(), False, 0, 0
# while not (d or (ep_len == max_ep_len)):
# # Take deterministic actions at test time (noise_scale=0)
# o, r, d, _ = test_env.step(get_action(o, 0))
# ep_ret += r
# ep_len += 1
# logger.store(TestEpRet=ep_ret, TestEpLen=ep_len)
# Prepare for interaction with environment
total_steps = steps_per_epoch * epochs
start_time = time.time()
o, ep_ret, ep_len, r, d = env.reset(), 0, 0, 0, False
# init trajectory saving (obs,action, reward, done)
trajectories = []
# Main loop: collect experience in env and update/log each epoch
for t in range(total_steps):
# Until start_steps have elapsed, randomly sample actions
# from a uniform distribution for better exploration. Afterwards,
# use the learned policy (with some noise, via act_noise).
if t > start_steps:
a = get_action(o, act_noise)
else:
a = env.action_space.sample()
# save step to trajectory
trajectories.append((o, a, r, d))
# Step the env
o2, r, d, _ = env.step(a)
ep_ret += r
ep_len += 1
# Ignore the "done" signal if it comes from hitting the time
# horizon (that is, when it's an artificial terminal signal
# that isn't based on the agent's state)
d = True if ep_len == max_ep_len else d
# Store experience to replay buffer
replay_buffer.store(o, a, r, o2, d)
# Super critical, easy to overlook step: make sure to update
# most recent observation!
o = o2
# End of trajectory handling
if d or (ep_len == max_ep_len):
logger.store(EpRet=ep_ret, EpLen=ep_len)
o, ep_ret, ep_len = env.reset(), 0, 0
# Update handling
if t >= update_after and t % update_every == 0:
for _ in range(update_every):
batch = replay_buffer.sample_batch(batch_size)
update(data=batch)
# End of epoch handling
if (t + 1) % steps_per_epoch == 0:
epoch = (t + 1) // steps_per_epoch
# Save model
if (epoch % save_freq == 0) or (epoch == epochs):
logger.save_state({'env': env}, None)
# # Test the performance of the deterministic version of the agent.
# test_agent()
#
# # Log info about epoch
# logger.log_tabular('Epoch', epoch)
# logger.log_tabular('EpRet', with_min_and_max=True)
# logger.log_tabular('TestEpRet', with_min_and_max=True)
# logger.log_tabular('EpLen', average_only=True)
# logger.log_tabular('TestEpLen', average_only=True)
# logger.log_tabular('TotalEnvInteracts', t)
# logger.log_tabular('QVals', with_min_and_max=True)
# logger.log_tabular('LossPi', average_only=True)
# logger.log_tabular('LossQ', average_only=True)
# logger.log_tabular('Time', time.time() - start_time)
# logger.dump_tabular()
np.save(str(trajectory_output_path) + '.npy', np.asarray(trajectories))
class SAC(tf.keras.Model):
def __init__(self,
env_fn,
actor_critic=core.mlp_actor_critic,
ac_kwargs=dict(),
seed=0,
steps_per_epoch=4000,
epochs=100,
replay_size=int(1e6),
gamma=0.99,
polyak=0.995,
lr=1e-3,
alpha=0.2,
batch_size=100,
start_steps=10000,
update_after=1000,
update_every=50,
num_test_episodes=10,
max_ep_len=1000,
logger_kwargs=dict(),
save_freq=1):
"""
Soft Actor-Critic (SAC)
Args:
env_fn : A function which creates a copy of the environment.
The environment must satisfy the OpenAI Gym API.
actor_critic: A function which takes in placeholder symbols
for state, ``x_ph``, and action, ``a_ph``, and returns the main
outputs from the agent's Tensorflow computation graph:
=========== ================ ======================================
Symbol Shape Description
=========== ================ ======================================
``mu`` (batch, act_dim) | Computes mean actions from policy
| given states.
``pi`` (batch, act_dim) | Samples actions from policy given
| states.
``logp_pi`` (batch,) | Gives log probability, according to
| the policy, of the action sampled by
| ``pi``. Critical: must be differentiable
| with respect to policy parameters all
| the way through action sampling.
``q1`` (batch,) | Gives one estimate of Q* for
| states in ``x_ph`` and actions in
| ``a_ph``.
``q2`` (batch,) | Gives another estimate of Q* for
| states in ``x_ph`` and actions in
| ``a_ph``.
=========== ================ ======================================
ac_kwargs (dict): Any kwargs appropriate for the actor_critic
function you provided to SAC.
seed (int): Seed for random number generators.
steps_per_epoch (int): Number of steps of interaction (state-action pairs)
for the agent and the environment in each epoch.
epochs (int): Number of epochs to run and train agent.
replay_size (int): Maximum length of replay buffer.
gamma (float): Discount factor. (Always between 0 and 1.)
polyak (float): Interpolation factor in polyak averaging for target
networks. Target networks are updated towards main networks
according to:
.. math:: \\theta_{\\text{targ}} \\leftarrow
\\rho \\theta_{\\text{targ}} + (1-\\rho) \\theta
where :math:`\\rho` is polyak. (Always between 0 and 1, usually
close to 1.)
lr (float): Learning rate (used for both policy and value learning).
alpha (float): Entropy regularization coefficient. (Equivalent to
inverse of reward scale in the original SAC paper.)
batch_size (int): Minibatch size for SGD.
start_steps (int): Number of steps for uniform-random action selection,
before running real policy. Helps exploration.
update_after (int): Number of env interactions to collect before
starting to do gradient descent updates. Ensures replay buffer
is full enough for useful updates.
update_every (int): Number of env interactions that should elapse
between gradient descent updates. Note: Regardless of how long
you wait between updates, the ratio of env steps to gradient steps
is locked to 1.
num_test_episodes (int): Number of episodes to test the deterministic
policy at the end of each epoch.
max_ep_len (int): Maximum length of trajectory / episode / rollout.
logger_kwargs (dict): Keyword args for EpochLogger.
save_freq (int): How often (in terms of gap between epochs) to save
the current policy and value function.
"""
super(SAC, self).__init__()
self.alpha = alpha
self.gamma = gamma
self.steps_per_epoch = steps_per_epoch
self.epochs = epochs
self.start_steps = start_steps
self.max_ep_len = max_ep_len
self.update_after = update_after
self.update_every = update_every
self.batch_size = batch_size
self.save_freq = save_freq
self.num_test_episodes = num_test_episodes
self.polyak = polyak
self.logger = EpochLogger(**logger_kwargs)
self.logger.save_config(locals())
tf.compat.v1.set_random_seed(seed)
np.random.seed(seed)
self.env, self.test_env = env_fn(), env_fn()
obs_dim = self.env.observation_space.shape[0]
act_dim = self.env.action_space.shape[0]
# Action limit for clamping: critically, assumes all dimensions share the same bound!
act_limit = self.env.action_space.high[0]
# Share information about action space with policy architecture
ac_kwargs['action_space'] = self.env.action_space
# Main outputs from computation graph
self.policy, self.Q1, self.Q2 = actor_critic(obs_dim,
act_dim,
name='main',
**ac_kwargs)
_, self.Q1_targ, self.Q2_targ = actor_critic(obs_dim,
act_dim,
name='target',
**ac_kwargs)
# Experience buffer
self.replay_buffer = ReplayBuffer(obs_dim=obs_dim,
act_dim=act_dim,
size=replay_size)
self.policy_optimizer = tf.keras.optimizers.Adam(learning_rate=lr)
self.value_optimizer = tf.keras.optimizers.Adam(learning_rate=lr)
self.target_init()
#print([np.array_equal(self.Q1.variables[i].numpy(), self.Q1_targ.variables[i].numpy()) for i in range(len(self.Q1.variables))])
@tf.function
def compute_loss_q(self, o, a, r, o2, d):
o_a = tf.concat([o, a], axis=1)
q1 = self.Q1(o_a)
q2 = self.Q2(o_a)
# Bellman backup for Q functions
# Target actions come from *current* policy
a2_mu, a2_pi, logp_a2pi = self.policy(o2)
o2_a2pi = tf.concat([o2, a2_pi], axis=-1)
# Target Q-values
q1_pi_targ = self.Q1_targ(o2_a2pi)
q2_pi_targ = self.Q2_targ(o2_a2pi)
q_pi_targ = tf.minimum(q1_pi_targ, q2_pi_targ)
backup = tf.stop_gradient(r + self.gamma * (1 - d) *
(q_pi_targ - self.alpha * logp_a2pi))
q1_loss = 0.5 * tf.reduce_mean((backup - q1)**2)
q2_loss = 0.5 * tf.reduce_mean((backup - q2)**2)
q_loss = q1_loss + q2_loss
return q_loss, q1_loss, q2_loss, q1, q2
@tf.function
def compute_loss_pi(self, o):
_, pi, logp_pi = self.policy(o)
o_pi = tf.concat([o, pi], axis=-1)
q1_pi = self.Q1(o_pi)
q2_pi = self.Q2(o_pi)
q_pi = tf.minimum(q1_pi, q2_pi)
# Entropy-regularized policy loss
loss_pi = tf.reduce_mean(self.alpha * logp_pi - q_pi)
return loss_pi, logp_pi
@tf.function
def compute_apply_gradients(self, o, o2, a, r, d):
main_pi_vars = self.policy.trainable_variables
main_q_vars = self.Q1.trainable_variables + self.Q2.trainable_variables
with tf.GradientTape() as tape:
pi_loss, logp_pi = self.compute_loss_pi(o=o)
pi_gradients = tape.gradient(pi_loss, main_pi_vars)
self.policy_optimizer.apply_gradients(zip(pi_gradients, main_pi_vars))
with tf.GradientTape() as tape:
value_loss, q1_loss, q2_loss, q1vals, q2vals = self.compute_loss_q(
o=o, o2=o2, a=a, r=r, d=d)
v_gradients = tape.gradient(value_loss, main_q_vars)
self.value_optimizer.apply_gradients(zip(v_gradients, main_q_vars))
return pi_loss, q1_loss, q2_loss, q1vals, q2vals, logp_pi
@tf.function
def update_target(self):
# Polyak averaging for target variables
# (control flow because sess.run otherwise evaluates in nondeterministic order)
for v_main, v_targ in zip(self.Q1.variables, self.Q1_targ.variables):
v_targ.assign(self.polyak * v_targ + (1 - self.polyak) * v_main)
for v_main, v_targ in zip(self.Q2.variables, self.Q2_targ.variables):
v_targ.assign(self.polyak * v_targ + (1 - self.polyak) * v_main)
@tf.function
def target_init(self):
# Initializing targets to match main variables
for v_main, v_targ in zip(self.Q1.variables, self.Q1_targ.variables):
v_targ.assign(v_main)
for v_main, v_targ in zip(self.Q2.variables, self.Q2_targ.variables):
v_targ.assign(v_main)
def get_action(self, o, deterministic=False):
mu, pi, _ = self.policy(tf.reshape(o, (1, -1)))
if deterministic:
return mu[0].numpy()
else:
return pi[0].numpy()
def test_agent(self):
for j in range(self.num_test_episodes):
o, d, ep_ret, ep_len = self.test_env.reset(), False, 0, 0
while not (d or (ep_len == self.max_ep_len)):
a = self.get_action(o, True)
# Take deterministic actions at test time
o, r, d, _ = self.test_env.step(a)
ep_ret += r
ep_len += 1
self.logger.store(TestEpRet=ep_ret, TestEpLen=ep_len)
def train(self):
start_time = time.time()
o, ep_ret, ep_len = self.env.reset(), 0, 0
total_steps = self.steps_per_epoch * self.epochs
# Main loop: collect experience in env and update/log each epoch
for t in range(total_steps):
# Until start_steps have elapsed, randomly sample actions
# from a uniform distribution for better exploration. Afterwards,
# use the learned policy.
if t > self.start_steps:
a = self.get_action(o)
else:
a = self.env.action_space.sample()
# Step the env
o2, r, d, _ = self.env.step(a)
ep_ret += r
ep_len += 1
# Ignore the "done" signal if it comes from hitting the time
# horizon (that is, when it's an artificial terminal signal
# that isn't based on the agent's state)
d = False if ep_len == self.max_ep_len else d
# Store experience to replay buffer
self.replay_buffer.store(o, a, r, o2, d)
# Super critical, easy to overlook step: make sure to update
# most recent observation!
o = o2
# End of trajectory handling
if d or (ep_len == self.max_ep_len):
self.logger.store(EpRet=ep_ret, EpLen=ep_len)
o, ep_ret, ep_len = self.env.reset(), 0, 0
# Update handling
if t >= self.update_after and t % self.update_every == 0:
for j in range(self.update_every):
batch = self.replay_buffer.sample_batch(self.batch_size)
outs = self.compute_apply_gradients(o=batch['obs1'],
o2=batch['obs2'],
a=batch['acts'],
r=batch['rews'],
d=batch['done'])
self.update_target()
self.logger.store(LossPi=outs[0],
LossQ1=outs[1],
LossQ2=outs[2],
Q1Vals=outs[3],
Q2Vals=outs[4],
LogPi=outs[5])
# End of epoch wrap-up
if (t + 1) % self.steps_per_epoch == 0:
epoch = (t + 1) // self.steps_per_epoch
# Save model
if (epoch % self.save_freq == 0) or (epoch == epochs):
self.logger.save_state({'env': self.env}, None)
# Test the performance of the deterministic version of the agent.
self.test_agent()
# Log info about epoch
self.logger.log_tabular('Epoch', epoch)
self.logger.log_tabular('EpRet', with_min_and_max=True)
self.logger.log_tabular('TestEpRet', with_min_and_max=True)
self.logger.log_tabular('EpLen', average_only=True)
self.logger.log_tabular('TestEpLen', average_only=True)
self.logger.log_tabular('TotalEnvInteracts', t)
self.logger.log_tabular('Q1Vals', with_min_and_max=True)
self.logger.log_tabular('Q2Vals', with_min_and_max=True)
self.logger.log_tabular('LogPi', with_min_and_max=True)
self.logger.log_tabular('LossPi', average_only=True)
self.logger.log_tabular('LossQ1', average_only=True)
self.logger.log_tabular('LossQ2', average_only=True)
self.logger.log_tabular('Time', time.time() - start_time)
self.logger.dump_tabular()
def ppo(env_fn,
actor_critic=core.MLPActorCritic,
ac_kwargs=dict(),
seed=0,
steps_per_epoch=4000,
epochs=50,
gamma=0.99,
clip_ratio=0.2,
pi_lr=3e-4,
vf_lr=1e-3,
train_pi_iters=80,
train_v_iters=80,
lam=0.97,
max_ep_len=1000,
target_kl=0.01,
logger_kwargs=dict(),
save_freq=10):
"""
Proximal Policy Optimization (by clipping),
with early stopping based on approximate KL
Args:
env_fn : A function which creates a copy of the environment.
The environment must satisfy the OpenAI Gym API.
actor_critic: The constructor method for a PyTorch Module with a
``step`` method, an ``act`` method, a ``pi`` module, and a ``v``
module. The ``step`` method should accept a batch of observations
and return:
=========== ================ ======================================
Symbol Shape Description
=========== ================ ======================================
``a`` (batch, act_dim) | Numpy array of actions for each
| observation.
``v`` (batch,) | Numpy array of value estimates
| for the provided observations.
``logp_a`` (batch,) | Numpy array of log probs for the
| actions in ``a``.
=========== ================ ======================================
The ``act`` method behaves the same as ``step`` but only returns ``a``.
The ``pi`` module's forward call should accept a batch of
observations and optionally a batch of actions, and return:
=========== ================ ======================================
Symbol Shape Description
=========== ================ ======================================
``pi`` N/A | Torch Distribution object, containing
| a batch of distributions describing
| the policy for the provided observations.
``logp_a`` (batch,) | Optional (only returned if batch of
| actions is given). Tensor containing
| the log probability, according to
| the policy, of the provided actions.
| If actions not given, will contain
| ``None``.
=========== ================ ======================================
The ``v`` module's forward call should accept a batch of observations
and return:
=========== ================ ======================================
Symbol Shape Description
=========== ================ ======================================
``v`` (batch,) | Tensor containing the value estimates
| for the provided observations. (Critical:
| make sure to flatten this!)
=========== ================ ======================================
ac_kwargs (dict): Any kwargs appropriate for the ActorCritic object
you provided to PPO.
seed (int): Seed for random number generators.
steps_per_epoch (int): Number of steps of interaction (state-action pairs)
for the agent and the environment in each epoch.
epochs (int): Number of epochs of interaction (equivalent to
number of policy updates) to perform.
gamma (float): Discount factor. (Always between 0 and 1.)
clip_ratio (float): Hyperparameter for clipping in the policy objective.
Roughly: how far can the new policy go from the old policy while
still profiting (improving the objective function)? The new policy
can still go farther than the clip_ratio says, but it doesn't help
on the objective anymore. (Usually small, 0.1 to 0.3.) Typically
denoted by :math:`\epsilon`.
pi_lr (float): Learning rate for policy optimizer.
vf_lr (float): Learning rate for value function optimizer.
train_pi_iters (int): Maximum number of gradient descent steps to take
on policy loss per epoch. (Early stopping may cause optimizer
to take fewer than this.)
train_v_iters (int): Number of gradient descent steps to take on
value function per epoch.
lam (float): Lambda for GAE-Lambda. (Always between 0 and 1,
close to 1.)
max_ep_len (int): Maximum length of trajectory / episode / rollout.
target_kl (float): Roughly what KL divergence we think is appropriate
between new and old policies after an update. This will get used
for early stopping. (Usually small, 0.01 or 0.05.)
logger_kwargs (dict): Keyword args for EpochLogger.
save_freq (int): How often (in terms of gap between epochs) to save
the current policy and value function.
"""
# Special function to avoid certain slowdowns from PyTorch + MPI combo.
setup_pytorch_for_mpi()
# Set up logger and save configuration
logger = EpochLogger(**logger_kwargs)
logger.save_config(locals())
# Random seed
seed += 10000 * proc_id()
torch.manual_seed(seed)
np.random.seed(seed)
# Instantiate environment
env = env_fn()
obs_dim = env.observation_space.shape
act_dim = env.action_space.shape
# Create actor-critic module
if inspect.isclass(actor_critic):
ac = actor_critic(env.observation_space, env.action_space, **ac_kwargs)
else:
ac = actor_critic
# Sync params across processes
sync_params(ac)
# Count variables
var_counts = tuple(core.count_vars(module) for module in [ac.pi, ac.v])
logger.log('\nNumber of parameters: \t pi: %d, \t v: %d\n' % var_counts)
# Set up experience buffer
local_steps_per_epoch = int(steps_per_epoch / num_procs())
buf = PPOBuffer(obs_dim, act_dim, local_steps_per_epoch, gamma, lam)
# Set up function for computing PPO policy loss
def compute_loss_pi(data):
obs, act, adv, logp_old = data['obs'], data['act'], data['adv'], data[
'logp']
# Policy loss
pi, logp = ac.pi(obs, act)
ratio = torch.exp(logp - logp_old)
clip_adv = torch.clamp(ratio, 1 - clip_ratio, 1 + clip_ratio) * adv
loss_pi = -(torch.min(ratio * adv, clip_adv)).mean()
# Useful extra info
approx_kl = (logp_old - logp).mean().item()
ent = pi.entropy().mean().item()
clipped = ratio.gt(1 + clip_ratio) | ratio.lt(1 - clip_ratio)
clipfrac = torch.as_tensor(clipped, dtype=torch.float32).mean().item()
pi_info = dict(kl=approx_kl, ent=ent, cf=clipfrac)
return loss_pi, pi_info
# Set up function for computing value loss
def compute_loss_v(data):
obs, ret = data['obs'], data['ret']
return ((ac.v(obs) - ret)**2).mean()
# Set up optimizers for policy and value function
pi_optimizer = Adam(ac.pi.parameters(), lr=pi_lr)
vf_optimizer = Adam(ac.v.parameters(), lr=vf_lr)
# Set up model saving
logger.setup_pytorch_saver(ac)
def update():
data = buf.get()
pi_l_old, pi_info_old = compute_loss_pi(data)
pi_l_old = pi_l_old.item()
v_l_old = compute_loss_v(data).item()
# Train policy with multiple steps of gradient descent
for i in range(train_pi_iters):
pi_optimizer.zero_grad()
loss_pi, pi_info = compute_loss_pi(data)
kl = mpi_avg(pi_info['kl'])
if kl > 1.5 * target_kl:
logger.log(
'Early stopping at step %d due to reaching max kl.' % i)
break
loss_pi.backward()
mpi_avg_grads(ac.pi) # average grads across MPI processes
pi_optimizer.step()
logger.store(StopIter=i)
# Value function learning
for i in range(train_v_iters):
vf_optimizer.zero_grad()
loss_v = compute_loss_v(data)
loss_v.backward()
mpi_avg_grads(ac.v) # average grads across MPI processes
vf_optimizer.step()
# Log changes from update
kl, ent, cf = pi_info['kl'], pi_info_old['ent'], pi_info['cf']
logger.store(LossPi=pi_l_old,
LossV=v_l_old,
KL=kl,
Entropy=ent,
ClipFrac=cf,
DeltaLossPi=(loss_pi.item() - pi_l_old),
DeltaLossV=(loss_v.item() - v_l_old))
# Prepare for interaction with environment
start_time = time.time()
o, ep_ret, ep_len = env.reset(), 0, 0
# Main loop: collect experience in env and update/log each epoch
for epoch in range(epochs):
for t in range(local_steps_per_epoch):
a, v, logp = ac.step(torch.as_tensor(o, dtype=torch.float32))
next_o, r, d, _ = env.step(a)
ep_ret += r
ep_len += 1
# save and log
buf.store(o, a, r, v, logp)
logger.store(VVals=v)
# Update obs (critical!)
o = next_o
timeout = ep_len == max_ep_len
terminal = d or timeout
epoch_ended = t == local_steps_per_epoch - 1
if terminal or epoch_ended:
if epoch_ended and not (terminal):
print('Warning: trajectory cut off by epoch at %d steps.' %
ep_len,
flush=True)
# if trajectory didn't reach terminal state, bootstrap value target
if timeout or epoch_ended:
_, v, _ = ac.step(torch.as_tensor(o, dtype=torch.float32))
else:
v = 0
buf.finish_path(v)
if terminal:
# only save EpRet / EpLen if trajectory finished
logger.store(EpRet=ep_ret, EpLen=ep_len)
o, ep_ret, ep_len = env.reset(), 0, 0
# Save model
if (epoch % save_freq == 0) or (epoch == epochs - 1):
logger.save_state({'env': env}, None) # current state
logger.save_state({'env': env}, epoch) # for rendering
# Perform PPO update!
update()
# Log info about epoch
logger.log_tabular('Epoch', epoch)
logger.log_tabular('EpRet', with_min_and_max=True)
logger.log_tabular('EpLen', average_only=True)
logger.log_tabular('VVals', with_min_and_max=True)
logger.log_tabular('TotalEnvInteracts', (epoch + 1) * steps_per_epoch)
logger.log_tabular('LossPi', average_only=True)
logger.log_tabular('LossV', average_only=True)
logger.log_tabular('DeltaLossPi', average_only=True)
logger.log_tabular('DeltaLossV', average_only=True)
logger.log_tabular('Entropy', average_only=True)
logger.log_tabular('KL', average_only=True)
logger.log_tabular('ClipFrac', average_only=True)
logger.log_tabular('StopIter', average_only=True)
logger.log_tabular('Time', time.time() - start_time)
logger.dump_tabular()
logger.output_file.close()
def td3(env_fn,
actor_critic=core.MLPActorCritic,
ac_kwargs=dict(),
seed=0,
steps_per_epoch=4000,
epochs=100,
replay_size=int(1e6),
gamma=0.99,
polyak=0.995,
pi_lr=1e-3,
q_lr=1e-3,
batch_size=100,
start_steps=10000,
update_after=1000,
update_every=50,
act_noise=0.1,
target_noise=0.2,
noise_clip=0.5,
policy_delay=2,
num_test_episodes=10,
max_ep_len=1000,
logger_kwargs=dict(),
save_freq=1):
"""
Twin Delayed Deep Deterministic Policy Gradient (TD3)
Args:
env_fn : A function which creates a copy of the environment.
The environment must satisfy the OpenAI Gym API.
actor_critic: The constructor method for a PyTorch Module with an ``act``
method, a ``pi`` module, a ``q1`` module, and a ``q2`` module.
The ``act`` method and ``pi`` module should accept batches of
observations as inputs, and ``q1`` and ``q2`` should accept a batch
of observations and a batch of actions as inputs. When called,
these should return:
=========== ================ ======================================
Call Output Shape Description
=========== ================ ======================================
``act`` (batch, act_dim) | Numpy array of actions for each
| observation.
``pi`` (batch, act_dim) | Tensor containing actions from policy
| given observations.
``q1`` (batch,) | Tensor containing one current estimate
| of Q* for the provided observations
| and actions. (Critical: make sure to
| flatten this!)
``q2`` (batch,) | Tensor containing the other current
| estimate of Q* for the provided observations
| and actions. (Critical: make sure to
| flatten this!)
=========== ================ ======================================
ac_kwargs (dict): Any kwargs appropriate for the ActorCritic object
you provided to TD3.
seed (int): Seed for random number generators.
steps_per_epoch (int): Number of steps of interaction (state-action pairs)
for the agent and the environment in each epoch.
epochs (int): Number of epochs to run and train agent.
replay_size (int): Maximum length of replay buffer.
gamma (float): Discount factor. (Always between 0 and 1.)
polyak (float): Interpolation factor in polyak averaging for target
networks. Target networks are updated towards main networks
according to:
.. math:: \\theta_{\\text{targ}} \\leftarrow
\\rho \\theta_{\\text{targ}} + (1-\\rho) \\theta
where :math:`\\rho` is polyak. (Always between 0 and 1, usually
close to 1.)
pi_lr (float): Learning rate for policy.
q_lr (float): Learning rate for Q-networks.
batch_size (int): Minibatch size for SGD.
start_steps (int): Number of steps for uniform-random action selection,
before running real policy. Helps exploration.
update_after (int): Number of env interactions to collect before
starting to do gradient descent updates. Ensures replay buffer
is full enough for useful updates.
update_every (int): Number of env interactions that should elapse
between gradient descent updates. Note: Regardless of how long
you wait between updates, the ratio of env steps to gradient steps
is locked to 1.
act_noise (float): Stddev for Gaussian exploration noise added to
policy at training time. (At test time, no noise is added.)
target_noise (float): Stddev for smoothing noise added to target
policy.
noise_clip (float): Limit for absolute value of target policy
smoothing noise.
policy_delay (int): Policy will only be updated once every
policy_delay times for each update of the Q-networks.
num_test_episodes (int): Number of episodes to test the deterministic
policy at the end of each epoch.
max_ep_len (int): Maximum length of trajectory / episode / rollout.
logger_kwargs (dict): Keyword args for EpochLogger.
save_freq (int): How often (in terms of gap between epochs) to save
the current policy and value function.
"""
logger = EpochLogger(**logger_kwargs)
torch.manual_seed(seed)
np.random.seed(seed)
env, test_env = env_fn, env_fn
obs_dim = 0
for key in list(env.observation_spec().keys()):
obs_dim += env.observation_spec()[key].shape[0]
act_dim = env.action_spec().shape[0]
# Action limit for clamping: critically, assumes all dimensions share the same bound!
act_limit = env.action_spec().maximum[0]
# Create actor-critic module and target networks
ac = actor_critic(obs_dim, act_dim, act_limit, **ac_kwargs)
ac.to(device)
ac_targ = deepcopy(ac)
# Freeze target networks with respect to optimizers (only update via polyak averaging)
for p in ac_targ.parameters():
p.requires_grad = False
# List of parameters for both Q-networks (save this for convenience)
q_params = itertools.chain(ac.q1.parameters(), ac.q2.parameters())
# Experience buffer
replay_buffer = ReplayBuffer(obs_dim=obs_dim,
act_dim=act_dim,
size=replay_size)
# Count variables (protip: try to get a feel for how different size networks behave!)
var_counts = tuple(
core.count_vars(module) for module in [ac.pi, ac.q1, ac.q2])
logger.log('\nNumber of parameters: \t pi: %d, \t q1: %d, \t q2: %d\n' %
var_counts)
# Set up function for computing TD3 Q-losses
def compute_loss_q(data):
o, a, r, o2, d = data['obs'], data['act'], data['rew'], data[
'obs2'], data['done']
q1 = ac.q1(o, a)
q2 = ac.q2(o, a)
# Bellman backup for Q functions
with torch.no_grad():
pi_targ = ac_targ.pi(o2)
# Target policy smoothing
epsilon = torch.randn_like(pi_targ) * target_noise
epsilon = torch.clamp(epsilon, -noise_clip, noise_clip)
a2 = pi_targ + epsilon
a2 = torch.clamp(a2, -act_limit, act_limit)
# Target Q-values
q1_pi_targ = ac_targ.q1(o2, a2)
q2_pi_targ = ac_targ.q2(o2, a2)
q_pi_targ = torch.min(q1_pi_targ, q2_pi_targ)
backup = r + gamma * (1 - d) * q_pi_targ
# MSE loss against Bellman backup
loss_q1 = ((q1 - backup)**2).mean()
loss_q2 = ((q2 - backup)**2).mean()
loss_q = loss_q1 + loss_q2
# Useful info for logging
loss_info = dict(Q1Vals=q1.cpu().detach().numpy(),
Q2Vals=q2.cpu().detach().numpy())
return loss_q, loss_info
# Set up function for computing TD3 pi loss
def compute_loss_pi(data):
o = data['obs']
q1_pi = ac.q1(o, ac.pi(o))
return -q1_pi.mean()
# Set up optimizers for policy and q-function
pi_optimizer = Adam(ac.pi.parameters(), lr=pi_lr)
q_optimizer = Adam(q_params, lr=q_lr)
# Set up model saving
logger.setup_pytorch_saver(ac)
def update(data, timer):
# First run one gradient descent step for Q1 and Q2
q_optimizer.zero_grad()
loss_q, loss_info = compute_loss_q(data)
loss_q.backward()
q_optimizer.step()
# Record things
logger.store(LossQ=loss_q.cpu().item(), **loss_info)
# Possibly update pi and target networks
if timer % policy_delay == 0:
# Freeze Q-networks so you don't waste computational effort
# computing gradients for them during the policy learning step.
for p in q_params:
p.requires_grad = False
# Next run one gradient descent step for pi.
pi_optimizer.zero_grad()
loss_pi = compute_loss_pi(data)
loss_pi.backward()
pi_optimizer.step()
# Unfreeze Q-networks so you can optimize it at next DDPG step.
for p in q_params:
p.requires_grad = True
# Record things
logger.store(LossPi=loss_pi.cpu().item())
# Finally, update target networks by polyak averaging.
with torch.no_grad():
for p, p_targ in zip(ac.parameters(), ac_targ.parameters()):
# NB: We use an in-place operations "mul_", "add_" to update target
# params, as opposed to "mul" and "add", which would make new tensors.
p_targ.data.mul_(polyak)
p_targ.data.add_((1 - polyak) * p.data)
def get_action(o, noise_scale):
a = ac.act(torch.as_tensor(o, device=device, dtype=torch.float32))
a += noise_scale * np.random.randn(act_dim)
return np.clip(a, -act_limit, act_limit)
def get_obs(env, a=None):
if a is None:
time_step = env.reset()
else:
time_step = env.step(a)
y = np.concatenate([
time_step.observation[key]
for key in list(time_step.observation.keys())
])
return y
def test_agent():
for j in range(num_test_episodes):
o, d, ep_ret, ep_len = get_obs(test_env), False, 0, 0
while not (d or (ep_len == max_ep_len)):
# Take deterministic actions at test time (noise_scale=0)
o, d = get_obs(test_env, get_action(o, 0)), time_step.last()
r = 0 if time_step.reward is None else time_step.reward
ep_ret += r
ep_len += 1
logger.store(TestEpRet=ep_ret, TestEpLen=ep_len)
# Prepare for interaction with environment
total_steps = steps_per_epoch * epochs
start_time = time.time()
o, ep_ret, ep_len = get_obs(env), 0, 0
# Main loop: collect experience in env and update/log each epoch
for t in trange(total_steps):
# Until start_steps have elapsed, randomly sample actions
# from a uniform distribution for better exploration. Afterwards,
# use the learned policy (with some noise, via act_noise).
if t > start_steps:
a = get_action(o, act_noise)
else:
action_spec = env.action_spec()
a = np.random.uniform(action_spec.minimum,
action_spec.maximum,
size=action_spec.shape)
# Step the env
time_step = env.step(a)
o2, d = get_obs(env, a), time_step.last()
r = 0 if time_step.reward is None else time_step.reward
ep_ret += r
ep_len += 1
# Ignore the "done" signal if it comes from hitting the time
# horizon (that is, when it's an artificial terminal signal
# that isn't based on the agent's state)
d = False if ep_len == max_ep_len else d
# Store experience to replay buffer
replay_buffer.store(o, a, r, o2, d)
# Super critical, easy to overlook step: make sure to update
# most recent observation!
o = o2
# End of trajectory handling
if d or (ep_len == max_ep_len):
logger.store(EpRet=ep_ret, EpLen=ep_len)
o, ep_ret, ep_len = get_obs(env), 0, 0
# Update handling
if t >= update_after and t % update_every == 0:
for j in range(update_every):
batch = replay_buffer.sample_batch(batch_size)
update(data=batch, timer=j)
# End of epoch handling
if (t + 1) % steps_per_epoch == 0:
epoch = (t + 1) // steps_per_epoch
# Save model
if (epoch % save_freq == 0) or (epoch == epochs):
logger.save_state({'env': env}, None)
# Test the performance of the deterministic version of the agent.
test_agent()
# Log info about epoch
logger.log_tabular('Epoch', epoch)
logger.log_tabular('EpRet', with_min_and_max=True)
logger.log_tabular('TestEpRet', with_min_and_max=True)
logger.log_tabular('EpLen', average_only=True)
logger.log_tabular('TestEpLen', average_only=True)
logger.log_tabular('TotalEnvInteracts', t)
logger.log_tabular('Q1Vals', with_min_and_max=True)
logger.log_tabular('Q2Vals', with_min_and_max=True)
logger.log_tabular('LossPi', average_only=True)
logger.log_tabular('LossQ', average_only=True)
logger.log_tabular('Time', time.time() - start_time)
logger.dump_tabular()
class AWAC:
def __init__(self, env_fn, actor_critic=core.MLPActorCritic,
ac_kwargs=dict(),
seed=0,
steps_per_epoch=100,
epochs=10000,
replay_size=int(2000000),
gamma=0.99,
polyak=0.995,
lr=3e-4,
p_lr=3e-4,
alpha=0.0,
batch_size=1024,
start_steps=10000,
update_after=0,
update_every=50,
num_test_episodes=10,
max_ep_len=1000,
logger_kwargs=dict(),
save_freq=1,
algo='SAC'):
"""
Soft Actor-Critic (SAC)
Args:
env_fn : A function which creates a copy of the environment.
The environment must satisfy the OpenAI Gym API.
actor_critic: The constructor method for a PyTorch Module with an ``act``
method, a ``pi`` module, a ``q1`` module, and a ``q2`` module.
The ``act`` method and ``pi`` module should accept batches of
observations as inputs, and ``q1`` and ``q2`` should accept a batch
of observations and a batch of actions as inputs. When called,
``act``, ``q1``, and ``q2`` should return:
=========== ================ ======================================
Call Output Shape Description
=========== ================ ======================================
``act`` (batch, act_dim) | Numpy array of actions for each
| observation.
``q1`` (batch,) | Tensor containing one current estimate
| of Q* for the provided observations
| and actions. (Critical: make sure to
| flatten this!)
``q2`` (batch,) | Tensor containing the other current
| estimate of Q* for the provided observations
| and actions. (Critical: make sure to
| flatten this!)
=========== ================ ======================================
Calling ``pi`` should return:
=========== ================ ======================================
Symbol Shape Description
=========== ================ ======================================
``a`` (batch, act_dim) | Tensor containing actions from policy
| given observations.
``logp_pi`` (batch,) | Tensor containing log probabilities of
| actions in ``a``. Importantly: gradients
| should be able to flow back into ``a``.
=========== ================ ======================================
ac_kwargs (dict): Any kwargs appropriate for the ActorCritic object
you provided to SAC.
seed (int): Seed for random number generators.
steps_per_epoch (int): Number of steps of interaction (state-action pairs)
for the agent and the environment in each epoch.
epochs (int): Number of epochs to run and train agent.
replay_size (int): Maximum length of replay buffer.
gamma (float): Discount factor. (Always between 0 and 1.)
polyak (float): Interpolation factor in polyak averaging for target
networks. Target networks are updated towards main networks
according to:
.. math:: \\theta_{\\text{targ}} \\leftarrow
\\rho \\theta_{\\text{targ}} + (1-\\rho) \\theta
where :math:`\\rho` is polyak. (Always between 0 and 1, usually
close to 1.)
lr (float): Learning rate (used for both policy and value learning).
alpha (float): Entropy regularization coefficient. (Equivalent to
inverse of reward scale in the original SAC paper.)
batch_size (int): Minibatch size for SGD.
start_steps (int): Number of steps for uniform-random action selection,
before running real policy. Helps exploration.
update_after (int): Number of env interactions to collect before
starting to do gradient descent updates. Ensures replay buffer
is full enough for useful updates.
update_every (int): Number of env interactions that should elapse
between gradient descent updates. Note: Regardless of how long
you wait between updates, the ratio of env steps to gradient steps
is locked to 1.
num_test_episodes (int): Number of episodes to test the deterministic
policy at the end of each epoch.
max_ep_len (int): Maximum length of trajectory / episode / rollout.
logger_kwargs (dict): Keyword args for EpochLogger.
save_freq (int): How often (in terms of gap between epochs) to save
the current policy and value function.
"""
self.logger = EpochLogger(**logger_kwargs)
self.logger.save_config(locals())
torch.manual_seed(seed)
np.random.seed(seed)
self.env, self.test_env = env_fn(), env_fn()
self.obs_dim = self.env.observation_space.shape
self.act_dim = self.env.action_space.shape[0]
# Action limit for clamping: critically, assumes all dimensions share the same bound!
self.act_limit = self.env.action_space.high[0]
# Create actor-critic module and target networks
self.ac = actor_critic(self.env.observation_space, self.env.action_space,
special_policy='awac', **ac_kwargs)
self.ac_targ = actor_critic(self.env.observation_space, self.env.action_space,
special_policy='awac', **ac_kwargs)
self.ac_targ.load_state_dict(self.ac.state_dict())
self.gamma = gamma
# Freeze target networks with respect to optimizers (only update via polyak averaging)
for p in self.ac_targ.parameters():
p.requires_grad = False
# List of parameters for both Q-networks (save this for convenience)
self.q_params = itertools.chain(self.ac.q1.parameters(), self.ac.q2.parameters())
# Experience buffer
self.replay_buffer = ReplayBuffer(obs_dim=self.obs_dim, act_dim=self.act_dim,
size=replay_size)
# Count variables (protip: try to get a feel for how different size networks behave!)
var_counts = tuple(
core.count_vars(module) for module in [self.ac.pi, self.ac.q1, self.ac.q2])
self.logger.log('\nNumber of parameters: \t pi: %d, \t q1: %d, \t q2: %d\n' % var_counts)
self.algo = algo
self.p_lr = p_lr
self.lr = lr
self.alpha = 0
# # Algorithm specific hyperparams
# Set up optimizers for policy and q-function
self.pi_optimizer = Adam(self.ac.pi.parameters(), lr=self.p_lr, weight_decay=1e-4)
self.q_optimizer = Adam(self.q_params, lr=self.lr)
self.num_test_episodes = num_test_episodes
self.max_ep_len = max_ep_len
self.epochs = epochs
self.steps_per_epoch = steps_per_epoch
self.update_after = update_after
self.update_every = update_every
self.batch_size = batch_size
self.save_freq = save_freq
self.polyak = polyak
# Set up model saving
self.logger.setup_pytorch_saver(self.ac)
print("Running Offline RL algorithm: {}".format(self.algo))
def populate_replay_buffer(self, env_name):
data_envs = {
'HalfCheetah-v2': (
"awac_data/hc_action_noise_15.npy",
"awac_data/hc_off_policy_15_demos_100.npy"),
'Ant-v2': (
"awac_data/ant_action_noise_15.npy",
"awac_data/ant_off_policy_15_demos_100.npy"),
'Walker2d-v2': (
"awac_data/walker_action_noise_15.npy",
"awac_data/walker_off_policy_15_demos_100.npy"),
}
if env_name in data_envs:
print('Loading saved data')
for file in data_envs[env_name]:
if not os.path.exists(file):
warnings.warn(colored('Offline data not found. Follow awac_data/instructions.txt to download. Running without offline data.', 'red'))
break
data = np.load(file, allow_pickle=True)
for demo in data:
for transition in list(zip(demo['observations'], demo['actions'], demo['rewards'],
demo['next_observations'], demo['terminals'])):
self.replay_buffer.store(*transition)
else:
dataset = d4rl.qlearning_dataset(self.env)
N = dataset['rewards'].shape[0]
for i in range(N):
self.replay_buffer.store(dataset['observations'][i], dataset['actions'][i],
dataset['rewards'][i], dataset['next_observations'][i],
float(dataset['terminals'][i]))
print("Loaded dataset")
# Set up function for computing SAC Q-losses
def compute_loss_q(self, data):
o, a, r, o2, d = data['obs'], data['act'], data['rew'], data['obs2'], data['done']
q1 = self.ac.q1(o, a)
q2 = self.ac.q2(o, a)
# Bellman backup for Q functions
with torch.no_grad():
# Target actions come from *current* policy
a2, logp_a2 = self.ac.pi(o2)
# Target Q-values
q1_pi_targ = self.ac_targ.q1(o2, a2)
q2_pi_targ = self.ac_targ.q2(o2, a2)
q_pi_targ = torch.min(q1_pi_targ, q2_pi_targ)
backup = r + self.gamma * (1 - d) * (q_pi_targ - self.alpha * logp_a2)
# MSE loss against Bellman backup
loss_q1 = ((q1 - backup) ** 2).mean()
loss_q2 = ((q2 - backup) ** 2).mean()
loss_q = loss_q1 + loss_q2
# Useful info for logging
q_info = dict(Q1Vals=q1.detach().numpy(),
Q2Vals=q2.detach().numpy())
return loss_q, q_info
# Set up function for computing SAC pi loss
def compute_loss_pi(self, data):
o = data['obs']
pi, logp_pi = self.ac.pi(o)
q1_pi = self.ac.q1(o, pi)
q2_pi = self.ac.q2(o, pi)
v_pi = torch.min(q1_pi, q2_pi)
beta = 2
q1_old_actions = self.ac.q1(o, data['act'])
q2_old_actions = self.ac.q2(o, data['act'])
q_old_actions = torch.min(q1_old_actions, q2_old_actions)
adv_pi = q_old_actions - v_pi
weights = F.softmax(adv_pi / beta, dim=0)
policy_logpp = self.ac.pi.get_logprob(o, data['act'])
loss_pi = (-policy_logpp * len(weights) * weights.detach()).mean()
# Useful info for logging
pi_info = dict(LogPi=policy_logpp.detach().numpy())
return loss_pi, pi_info
def update(self, data, update_timestep):
# First run one gradient descent step for Q1 and Q2
self.q_optimizer.zero_grad()
loss_q, q_info = self.compute_loss_q(data)
loss_q.backward()
self.q_optimizer.step()
# Record things
self.logger.store(LossQ=loss_q.item(), **q_info)
# Freeze Q-networks so you don't waste computational effort
# computing gradients for them during the policy learning step.
for p in self.q_params:
p.requires_grad = False
# Next run one gradient descent step for pi.
self.pi_optimizer.zero_grad()
loss_pi, pi_info = self.compute_loss_pi(data)
loss_pi.backward()
self.pi_optimizer.step()
# Unfreeze Q-networks so you can optimize it at next DDPG step.
for p in self.q_params:
p.requires_grad = True
# Record things
self.logger.store(LossPi=loss_pi.item(), **pi_info)
# Finally, update target networks by polyak averaging.
with torch.no_grad():
for p, p_targ in zip(self.ac.parameters(), self.ac_targ.parameters()):
# NB: We use an in-place operations "mul_", "add_" to update target
# params, as opposed to "mul" and "add", which would make new tensors.
p_targ.data.mul_(self.polyak)
p_targ.data.add_((1 - self.polyak) * p.data)
def get_action(self, o, deterministic=False):
return self.ac.act(torch.as_tensor(o, dtype=torch.float32), deterministic)
def test_agent(self):
for j in range(self.num_test_episodes):
o, d, ep_ret, ep_len = self.test_env.reset(), False, 0, 0
while not (d or (ep_len == self.max_ep_len)):
# Take deterministic actions at test time
o, r, d, _ = self.test_env.step(self.get_action(o, True))
ep_ret += r
ep_len += 1
self.logger.store(TestEpRet=ep_ret, TestEpLen=ep_len) # Get unnormalized score
# self.logger.store(TestEpRet=100*self.test_env.get_normalized_score(ep_ret), TestEpLen=ep_len) # Get normalized score
def run(self):
# Prepare for interaction with environment
total_steps = self.epochs * self.steps_per_epoch
start_time = time.time()
obs, ep_ret, ep_len = self.env.reset(), 0, 0
done = True
num_train_episodes = 0
# Main loop: collect experience in env and update/log each epoch
for t in range(total_steps):
# Reset stuff if necessary
if done and t > 0:
self.logger.store(ExplEpRet=ep_ret, ExplEpLen=ep_len)
obs, ep_ret, ep_len = self.env.reset(), 0, 0
num_train_episodes += 1
# Collect experience
act = self.get_action(obs, deterministic=False)
next_obs, rew, done, info = self.env.step(act)
self.replay_buffer.store(obs, act, rew, next_obs, done)
obs = next_obs
# Update handling
if t > self.update_after and t % self.update_every == 0:
for _ in range(self.update_every):
batch = self.replay_buffer.sample_batch(self.batch_size)
self.update(data=batch, update_timestep=t)
# End of epoch handling
if (t + 1) % self.steps_per_epoch == 0:
epoch = (t + 1) // self.steps_per_epoch
# Save model
if (epoch % self.save_freq == 0) or (epoch == self.epochs):
self.logger.save_state({'env': self.env}, None)
# Test the performance of the deterministic version of the agent.
self.test_agent()
# Log info about epoch
self.logger.log_tabular('Epoch', epoch)
self.logger.log_tabular('TestEpRet', with_min_and_max=True)
self.logger.log_tabular('TestEpLen', average_only=True)
self.logger.log_tabular('TotalUpdates', t)
self.logger.log_tabular('Q1Vals', with_min_and_max=True)
self.logger.log_tabular('Q2Vals', with_min_and_max=True)
self.logger.log_tabular('LogPi', with_min_and_max=True)
self.logger.log_tabular('LossPi', average_only=True)
self.logger.log_tabular('LossQ', average_only=True)
self.logger.log_tabular('Time', time.time() - start_time)
self.logger.dump_tabular()
def sac(env_fn,
actor_critic=core.mlp_actor_critic,
ac_kwargs=dict(),
seed=0,
steps_per_epoch=4000,
epochs=100,
replay_size=int(1e6),
gamma=0.99,
polyak=0.995,
lr=1e-3,
alpha=0.2,
batch_size=100,
start_steps=10000,
update_after=1000,
update_every=50,
num_test_episodes=10,
max_ep_len=1000,
logger_kwargs=dict(),
save_freq=1):
"""
Soft Actor-Critic (SAC)
Args:
env_fn : A function which creates a copy of the environment.
The environment must satisfy the OpenAI Gym API.
actor_critic: A function which takes in placeholder symbols
for state, ``x_ph``, and action, ``a_ph``, and returns the main
outputs from the agent's Tensorflow computation graph:
=========== ================ ======================================
Symbol Shape Description
=========== ================ ======================================
``mu`` (batch, act_dim) | Computes mean actions from policy
| given states.
``pi`` (batch, act_dim) | Samples actions from policy given
| states.
``logp_pi`` (batch,) | Gives log probability, according to
| the policy, of the action sampled by
| ``pi``. Critical: must be differentiable
| with respect to policy parameters all
| the way through action sampling.
``q1`` (batch,) | Gives one estimate of Q* for
| states in ``x_ph`` and actions in
| ``a_ph``.
``q2`` (batch,) | Gives another estimate of Q* for
| states in ``x_ph`` and actions in
| ``a_ph``.
=========== ================ ======================================
ac_kwargs (dict): Any kwargs appropriate for the actor_critic
function you provided to SAC.
seed (int): Seed for random number generators.
steps_per_epoch (int): Number of steps of interaction (state-action pairs)
for the agent and the environment in each epoch.
epochs (int): Number of epochs to run and train agent.
replay_size (int): Maximum length of replay buffer.
gamma (float): Discount factor. (Always between 0 and 1.)
polyak (float): Interpolation factor in polyak averaging for target
networks. Target networks are updated towards main networks
according to:
.. math:: \\theta_{\\text{targ}} \\leftarrow
\\rho \\theta_{\\text{targ}} + (1-\\rho) \\theta
where :math:`\\rho` is polyak. (Always between 0 and 1, usually
close to 1.)
lr (float): Learning rate (used for both policy and value learning).
alpha (float): Entropy regularization coefficient. (Equivalent to
inverse of reward scale in the original SAC paper.)
batch_size (int): Minibatch size for SGD.
start_steps (int): Number of steps for uniform-random action selection,
before running real policy. Helps exploration.
update_after (int): Number of env interactions to collect before
starting to do gradient descent updates. Ensures replay buffer
is full enough for useful updates.
update_every (int): Number of env interactions that should elapse
between gradient descent updates. Note: Regardless of how long
you wait between updates, the ratio of env steps to gradient steps
is locked to 1.
num_test_episodes (int): Number of episodes to test the deterministic
policy at the end of each epoch.
max_ep_len (int): Maximum length of trajectory / episode / rollout.
logger_kwargs (dict): Keyword args for EpochLogger.
save_freq (int): How often (in terms of gap between epochs) to save
the current policy and value function.
"""
logger = EpochLogger(**logger_kwargs)
logger.save_config(locals())
tf.set_random_seed(seed)
np.random.seed(seed)
env, test_env = env_fn(), env_fn()
obs_dim = env.observation_space.shape[0]
act_dim = env.action_space.shape[0]
# Action limit for clamping: critically, assumes all dimensions share the same bound!
act_limit = env.action_space.high[0]
# Share information about action space with policy architecture
ac_kwargs['action_space'] = env.action_space
# Inputs to computation graph
x_ph, a_ph, x2_ph, r_ph, d_ph = core.placeholders(obs_dim, act_dim,
obs_dim, None, None)
# Main outputs from computation graph
with tf.variable_scope('main'):
mu, pi, logp_pi, q1, q2 = actor_critic(x_ph, a_ph, **ac_kwargs)
with tf.variable_scope('main', reuse=True):
# compose q with pi, for pi-learning
_, _, _, q1_pi, q2_pi = actor_critic(x_ph, pi, **ac_kwargs)
# get actions and log probs of actions for next states, for Q-learning
_, pi_next, logp_pi_next, _, _ = actor_critic(x2_ph, a_ph, **ac_kwargs)
# Target value network
with tf.variable_scope('target'):
# target q values, using actions from *current* policy
_, _, _, q1_targ, q2_targ = actor_critic(x2_ph, pi_next, **ac_kwargs)
# Experience buffer
replay_buffer = ReplayBuffer(obs_dim=obs_dim,
act_dim=act_dim,
size=replay_size)
# Count variables
var_counts = tuple(
core.count_vars(scope)
for scope in ['main/pi', 'main/q1', 'main/q2', 'main'])
print(
'\nNumber of parameters: \t pi: %d, \t q1: %d, \t q2: %d, \t total: %d\n'
% var_counts)
# Min Double-Q:
min_q_pi = tf.minimum(q1_pi, q2_pi)
min_q_targ = tf.minimum(q1_targ, q2_targ)
# Entropy-regularized Bellman backup for Q functions, using Clipped Double-Q targets
q_backup = tf.stop_gradient(r_ph + gamma * (1 - d_ph) *
(min_q_targ - alpha * logp_pi_next))
# Soft actor-critic losses
pi_loss = tf.reduce_mean(alpha * logp_pi - min_q_pi)
q1_loss = 0.5 * tf.reduce_mean((q_backup - q1)**2)
q2_loss = 0.5 * tf.reduce_mean((q_backup - q2)**2)
value_loss = q1_loss + q2_loss
# Policy train op
# (has to be separate from value train op, because q1_pi appears in pi_loss)
pi_optimizer = tf.train.AdamOptimizer(learning_rate=lr)
train_pi_op = pi_optimizer.minimize(pi_loss, var_list=get_vars('main/pi'))
# Value train op
# (control dep of train_pi_op because sess.run otherwise evaluates in nondeterministic order)
value_optimizer = tf.train.AdamOptimizer(learning_rate=lr)
value_params = get_vars('main/q')
with tf.control_dependencies([train_pi_op]):
train_value_op = value_optimizer.minimize(value_loss,
var_list=value_params)
# Polyak averaging for target variables
# (control flow because sess.run otherwise evaluates in nondeterministic order)
with tf.control_dependencies([train_value_op]):
target_update = tf.group([
tf.assign(v_targ, polyak * v_targ + (1 - polyak) * v_main)
for v_main, v_targ in zip(get_vars('main'), get_vars('target'))
])
# All ops to call during one training step
step_ops = [
pi_loss, q1_loss, q2_loss, q1, q2, logp_pi, train_pi_op,
train_value_op, target_update
]
# Initializing targets to match main variables
target_init = tf.group([
tf.assign(v_targ, v_main)
for v_main, v_targ in zip(get_vars('main'), get_vars('target'))
])
sess = tf.Session()
sess.run(tf.global_variables_initializer())
sess.run(target_init)
# Setup model saving
logger.setup_tf_saver(sess,
inputs={
'x': x_ph,
'a': a_ph
},
outputs={
'mu': mu,
'pi': pi,
'q1': q1,
'q2': q2
})
def get_action(o, deterministic=False):
act_op = mu if deterministic else pi
return sess.run(act_op, feed_dict={x_ph: o.reshape(1, -1)})[0]
def test_agent():
for j in range(num_test_episodes):
o, d, ep_ret, ep_len = test_env.reset(), False, 0, 0
while not (d or (ep_len == max_ep_len)):
# Take deterministic actions at test time
o, r, d, _ = test_env.step(get_action(o, True))
ep_ret += r
ep_len += 1
logger.store(TestEpRet=ep_ret, TestEpLen=ep_len)
start_time = time.time()
o, ep_ret, ep_len = env.reset(), 0, 0
total_steps = steps_per_epoch * epochs
# Main loop: collect experience in env and update/log each epoch
for t in range(total_steps):
# Until start_steps have elapsed, randomly sample actions
# from a uniform distribution for better exploration. Afterwards,
# use the learned policy.
if t > start_steps:
a = get_action(o)
else:
a = env.action_space.sample()
# Step the env
o2, r, d, _ = env.step(a)
ep_ret += r
ep_len += 1
# Ignore the "done" signal if it comes from hitting the time
# horizon (that is, when it's an artificial terminal signal
# that isn't based on the agent's state)
d = False if ep_len == max_ep_len else d
# Store experience to replay buffer
replay_buffer.store(o, a, r, o2, d)
# Super critical, easy to overlook step: make sure to update
# most recent observation!
o = o2
# End of trajectory handling
if d or (ep_len == max_ep_len):
logger.store(EpRet=ep_ret, EpLen=ep_len)
o, ep_ret, ep_len = env.reset(), 0, 0
# Update handling
if t >= update_after and t % update_every == 0:
for j in range(update_every):
batch = replay_buffer.sample_batch(batch_size)
feed_dict = {
x_ph: batch['obs1'],
x2_ph: batch['obs2'],
a_ph: batch['acts'],
r_ph: batch['rews'],
d_ph: batch['done'],
}
outs = sess.run(step_ops, feed_dict)
logger.store(LossPi=outs[0],
LossQ1=outs[1],
LossQ2=outs[2],
Q1Vals=outs[3],
Q2Vals=outs[4],
LogPi=outs[5])
# End of epoch wrap-up
if (t + 1) % steps_per_epoch == 0:
epoch = (t + 1) // steps_per_epoch
# Save model
if (epoch % save_freq == 0) or (epoch == epochs):
logger.save_state({'env': env}, None)
# Test the performance of the deterministic version of the agent.
test_agent()
# Log info about epoch
logger.log_tabular('Epoch', epoch)
logger.log_tabular('EpRet', with_min_and_max=True)
logger.log_tabular('TestEpRet', with_min_and_max=True)
logger.log_tabular('EpLen', average_only=True)
logger.log_tabular('TestEpLen', average_only=True)
logger.log_tabular('TotalEnvInteracts', t)
logger.log_tabular('Q1Vals', with_min_and_max=True)
logger.log_tabular('Q2Vals', with_min_and_max=True)
logger.log_tabular('LogPi', with_min_and_max=True)
logger.log_tabular('LossPi', average_only=True)
logger.log_tabular('LossQ1', average_only=True)
logger.log_tabular('LossQ2', average_only=True)
logger.log_tabular('Time', time.time() - start_time)
logger.dump_tabular()