-
-
-
-
-import numpy as np
-import tensorflow as tf
-import gym
-import time
-from spinup.algos.ddpg import core
-from spinup.algos.ddpg.core import get_vars
-from spinup.utils.logx import EpochLogger
-
-
-class ReplayBuffer:
- """
- A simple FIFO experience replay buffer for DDPG agents.
- """
-
- def __init__(self, obs_dim, act_dim, size):
- self.obs1_buf = np.zeros([size, obs_dim], dtype=np.float32)
- self.obs2_buf = np.zeros([size, obs_dim], dtype=np.float32)
- self.acts_buf = np.zeros([size, act_dim], dtype=np.float32)
- self.rews_buf = np.zeros(size, dtype=np.float32)
- self.done_buf = np.zeros(size, dtype=np.float32)
- self.ptr, self.size, self.max_size = 0, 0, size
-
- def store(self, obs, act, rew, next_obs, done):
- self.obs1_buf[self.ptr] = obs
- self.obs2_buf[self.ptr] = next_obs
- self.acts_buf[self.ptr] = act
- self.rews_buf[self.ptr] = rew
- self.done_buf[self.ptr] = done
- self.ptr = (self.ptr+1) % self.max_size
- self.size = min(self.size+1, self.max_size)
-
- def sample_batch(self, batch_size=32):
- idxs = np.random.randint(0, self.size, size=batch_size)
- return dict(obs1=self.obs1_buf[idxs],
- obs2=self.obs2_buf[idxs],
- acts=self.acts_buf[idxs],
- rews=self.rews_buf[idxs],
- done=self.done_buf[idxs])
-
-"""
-
-Deep Deterministic Policy Gradient (DDPG)
-
-"""
-[docs]def ddpg(env_fn, actor_critic=core.mlp_actor_critic, ac_kwargs=dict(), seed=0,
- steps_per_epoch=5000, 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,
- act_noise=0.1, max_ep_len=1000, logger_kwargs=dict(), save_freq=1):
- """
-
- 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.
- ``q`` (batch,) | Gives the current estimate of Q* for
- | states in ``x_ph`` and actions in
- | ``a_ph``.
- ``q_pi`` (batch,) | Gives the composition of ``q`` and
- | ``pi`` for states in ``x_ph``:
- | q(x, pi(x)).
- =========== ================ ======================================
-
- ac_kwargs (dict): Any kwargs appropriate for the actor_critic
- function 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.
-
- act_noise (float): Stddev for Gaussian exploration noise added to
- policy at training time. (At test time, no noise is added.)
-
- 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, q, q_pi = actor_critic(x_ph, a_ph, **ac_kwargs)
-
- # Target networks
- with tf.variable_scope('target'):
- # Note that the action placeholder going to actor_critic here is
- # irrelevant, because we only need q_targ(s, pi_targ(s)).
- pi_targ, _, q_pi_targ = actor_critic(x2_ph, a_ph, **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/q', 'main'])
- print('\nNumber of parameters: \t pi: %d, \t q: %d, \t total: %d\n'%var_counts)
-
- # Bellman backup for Q function
- backup = tf.stop_gradient(r_ph + gamma*(1-d_ph)*q_pi_targ)
-
- # DDPG losses
- pi_loss = -tf.reduce_mean(q_pi)
- q_loss = tf.reduce_mean((q-backup)**2)
-
- # 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, 'q': q})
-
- 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(n=10):
- for j in range(n):
- o, r, d, ep_ret, ep_len = test_env.reset(), 0, 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, r, d, ep_ret, ep_len = env.reset(), 0, False, 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
-
- if d or (ep_len == max_ep_len):
- """
- Perform all DDPG updates at the end of the trajectory,
- in accordance with tuning done by TD3 paper authors.
- """
- for _ in range(ep_len):
- 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-learning update
- outs = sess.run([q_loss, q, train_q_op], feed_dict)
- logger.store(LossQ=outs[0], QVals=outs[1])
-
- # Policy update
- outs = sess.run([pi_loss, train_pi_op, target_update], feed_dict)
- logger.store(LossPi=outs[0])
-
- logger.store(EpRet=ep_ret, EpLen=ep_len)
- o, r, d, ep_ret, ep_len = env.reset(), 0, False, 0, 0
-
- # End of epoch wrap-up
- if t > 0 and t % steps_per_epoch == 0:
- epoch = t // steps_per_epoch
-
- # Save model
- if (epoch % save_freq == 0) or (epoch == epochs-1):
- 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()
-
-if __name__ == '__main__':
- import argparse
- parser = argparse.ArgumentParser()
- parser.add_argument('--env', type=str, default='HalfCheetah-v2')
- parser.add_argument('--hid', type=int, default=300)
- parser.add_argument('--l', type=int, default=1)
- parser.add_argument('--gamma', type=float, default=0.99)
- parser.add_argument('--seed', '-s', type=int, default=0)
- parser.add_argument('--epochs', type=int, default=50)
- parser.add_argument('--exp_name', type=str, default='ddpg')
- args = parser.parse_args()
-
- from spinup.utils.run_utils import setup_logger_kwargs
- logger_kwargs = setup_logger_kwargs(args.exp_name, args.seed)
-
- ddpg(lambda : gym.make(args.env), actor_critic=core.mlp_actor_critic,
- ac_kwargs=dict(hidden_sizes=[args.hid]*args.l),
- gamma=args.gamma, seed=args.seed, epochs=args.epochs,
- logger_kwargs=logger_kwargs)
-
-import numpy as np
-import tensorflow as tf
-import gym
-import time
-import spinup.algos.ppo.core as core
-from spinup.utils.logx import EpochLogger
-from spinup.utils.mpi_tf import MpiAdamOptimizer, sync_all_params
-from spinup.utils.mpi_tools import mpi_fork, mpi_avg, proc_id, mpi_statistics_scalar, num_procs
-
-
-class PPOBuffer:
- """
- A buffer for storing trajectories experienced by a PPO agent interacting
- with the environment, and using Generalized Advantage Estimation (GAE-Lambda)
- for calculating the advantages of state-action pairs.
- """
-
- def __init__(self, obs_dim, act_dim, size, gamma=0.99, lam=0.95):
- self.obs_buf = np.zeros(core.combined_shape(size, obs_dim), dtype=np.float32)
- self.act_buf = np.zeros(core.combined_shape(size, act_dim), dtype=np.float32)
- self.adv_buf = np.zeros(size, dtype=np.float32)
- self.rew_buf = np.zeros(size, dtype=np.float32)
- self.ret_buf = np.zeros(size, dtype=np.float32)
- self.val_buf = np.zeros(size, dtype=np.float32)
- self.logp_buf = np.zeros(size, dtype=np.float32)
- self.gamma, self.lam = gamma, lam
- self.ptr, self.path_start_idx, self.max_size = 0, 0, size
-
- def store(self, obs, act, rew, val, logp):
- """
- Append one timestep of agent-environment interaction to the buffer.
- """
- assert self.ptr < self.max_size # buffer has to have room so you can store
- self.obs_buf[self.ptr] = obs
- self.act_buf[self.ptr] = act
- self.rew_buf[self.ptr] = rew
- self.val_buf[self.ptr] = val
- self.logp_buf[self.ptr] = logp
- self.ptr += 1
-
- def finish_path(self, last_val=0):
- """
- Call this at the end of a trajectory, or when one gets cut off
- by an epoch ending. This looks back in the buffer to where the
- trajectory started, and uses rewards and value estimates from
- the whole trajectory to compute advantage estimates with GAE-Lambda,
- as well as compute the rewards-to-go for each state, to use as
- the targets for the value function.
-
- The "last_val" argument should be 0 if the trajectory ended
- because the agent reached a terminal state (died), and otherwise
- should be V(s_T), the value function estimated for the last state.
- This allows us to bootstrap the reward-to-go calculation to account
- for timesteps beyond the arbitrary episode horizon (or epoch cutoff).
- """
-
- path_slice = slice(self.path_start_idx, self.ptr)
- rews = np.append(self.rew_buf[path_slice], last_val)
- vals = np.append(self.val_buf[path_slice], last_val)
-
- # the next two lines implement GAE-Lambda advantage calculation
- deltas = rews[:-1] + self.gamma * vals[1:] - vals[:-1]
- self.adv_buf[path_slice] = core.discount_cumsum(deltas, self.gamma * self.lam)
-
- # the next line computes rewards-to-go, to be targets for the value function
- self.ret_buf[path_slice] = core.discount_cumsum(rews, self.gamma)[:-1]
-
- self.path_start_idx = self.ptr
-
- def get(self):
- """
- Call this at the end of an epoch to get all of the data from
- the buffer, with advantages appropriately normalized (shifted to have
- mean zero and std one). Also, resets some pointers in the buffer.
- """
- assert self.ptr == self.max_size # buffer has to be full before you can get
- self.ptr, self.path_start_idx = 0, 0
- # the next two lines implement the advantage normalization trick
- adv_mean, adv_std = mpi_statistics_scalar(self.adv_buf)
- self.adv_buf = (self.adv_buf - adv_mean) / adv_std
- return [self.obs_buf, self.act_buf, self.adv_buf,
- self.ret_buf, self.logp_buf]
-
-
-"""
-
-Proximal Policy Optimization (by clipping),
-
-with early stopping based on approximate KL
-
-"""
-[docs]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):
- """
-
- 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.)
-
- 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)})
-
- # save and log
- buf.store(o, a, r, v_t, logp_t)
- logger.store(VVals=v_t)
-
- o, r, d, _ = env.step(a[0])
- ep_ret += r
- ep_len += 1
-
- 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 = r 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, r, d, ep_ret, ep_len = env.reset(), 0, False, 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()
-
-if __name__ == '__main__':
- import argparse
- parser = argparse.ArgumentParser()
- parser.add_argument('--env', type=str, default='HalfCheetah-v2')
- parser.add_argument('--hid', type=int, default=64)
- parser.add_argument('--l', type=int, default=2)
- parser.add_argument('--gamma', type=float, default=0.99)
- parser.add_argument('--seed', '-s', type=int, default=0)
- parser.add_argument('--cpu', type=int, default=4)
- parser.add_argument('--steps', type=int, default=4000)
- parser.add_argument('--epochs', type=int, default=50)
- parser.add_argument('--exp_name', type=str, default='ppo')
- args = parser.parse_args()
-
- mpi_fork(args.cpu) # run parallel code with mpi
-
- from spinup.utils.run_utils import setup_logger_kwargs
- logger_kwargs = setup_logger_kwargs(args.exp_name, args.seed)
-
- ppo(lambda : gym.make(args.env), actor_critic=core.mlp_actor_critic,
- ac_kwargs=dict(hidden_sizes=[args.hid]*args.l), gamma=args.gamma,
- seed=args.seed, steps_per_epoch=args.steps, epochs=args.epochs,
- logger_kwargs=logger_kwargs)
-
-import numpy as np
-import tensorflow as tf
-import gym
-import time
-from spinup.algos.sac import core
-from spinup.algos.sac.core import get_vars
-from spinup.utils.logx import EpochLogger
-
-
-class ReplayBuffer:
- """
- A simple FIFO experience replay buffer for SAC agents.
- """
-
- def __init__(self, obs_dim, act_dim, size):
- self.obs1_buf = np.zeros([size, obs_dim], dtype=np.float32)
- self.obs2_buf = np.zeros([size, obs_dim], dtype=np.float32)
- self.acts_buf = np.zeros([size, act_dim], dtype=np.float32)
- self.rews_buf = np.zeros(size, dtype=np.float32)
- self.done_buf = np.zeros(size, dtype=np.float32)
- self.ptr, self.size, self.max_size = 0, 0, size
-
- def store(self, obs, act, rew, next_obs, done):
- self.obs1_buf[self.ptr] = obs
- self.obs2_buf[self.ptr] = next_obs
- self.acts_buf[self.ptr] = act
- self.rews_buf[self.ptr] = rew
- self.done_buf[self.ptr] = done
- self.ptr = (self.ptr+1) % self.max_size
- self.size = min(self.size+1, self.max_size)
-
- def sample_batch(self, batch_size=32):
- idxs = np.random.randint(0, self.size, size=batch_size)
- return dict(obs1=self.obs1_buf[idxs],
- obs2=self.obs2_buf[idxs],
- acts=self.acts_buf[idxs],
- rews=self.rews_buf[idxs],
- done=self.done_buf[idxs])
-
-"""
-
-Soft Actor-Critic
-
-(With slight variations that bring it closer to TD3)
-
-"""
-[docs]def sac(env_fn, actor_critic=core.mlp_actor_critic, ac_kwargs=dict(), seed=0,
- steps_per_epoch=5000, epochs=100, replay_size=int(1e6), gamma=0.99,
- polyak=0.995, lr=1e-3, alpha=0.2, batch_size=100, start_steps=10000,
- max_ep_len=1000, logger_kwargs=dict(), save_freq=1):
- """
-
- 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``.
- ``q1_pi`` (batch,) | Gives the composition of ``q1`` and
- | ``pi`` for states in ``x_ph``:
- | q1(x, pi(x)).
- ``q2_pi`` (batch,) | Gives the composition of ``q2`` and
- | ``pi`` for states in ``x_ph``:
- | q2(x, pi(x)).
- ``v`` (batch,) | Gives the value estimate for states
- | in ``x_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.
-
- 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, q1_pi, q2_pi, v = actor_critic(x_ph, a_ph, **ac_kwargs)
-
- # Target value network
- with tf.variable_scope('target'):
- _, _, _, _, _, _, _, v_targ = actor_critic(x2_ph, a_ph, **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/v', 'main'])
- print(('\nNumber of parameters: \t pi: %d, \t' + \
- 'q1: %d, \t q2: %d, \t v: %d, \t total: %d\n')%var_counts)
-
- # Min Double-Q:
- min_q_pi = tf.minimum(q1_pi, q2_pi)
-
- # Targets for Q and V regression
- q_backup = tf.stop_gradient(r_ph + gamma*(1-d_ph)*v_targ)
- v_backup = tf.stop_gradient(min_q_pi - alpha * logp_pi)
-
- # Soft actor-critic losses
- pi_loss = tf.reduce_mean(alpha * logp_pi - q1_pi)
- q1_loss = 0.5 * tf.reduce_mean((q_backup - q1)**2)
- q2_loss = 0.5 * tf.reduce_mean((q_backup - q2)**2)
- v_loss = 0.5 * tf.reduce_mean((v_backup - v)**2)
- value_loss = q1_loss + q2_loss + v_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') + get_vars('main/v')
- 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, v_loss, q1, q2, v, 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, 'v': v})
-
- 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(n=10):
- global sess, mu, pi, q1, q2, q1_pi, q2_pi
- for j in range(n):
- o, r, d, ep_ret, ep_len = test_env.reset(), 0, 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, r, d, ep_ret, ep_len = env.reset(), 0, False, 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
-
- if d or (ep_len == max_ep_len):
- """
- Perform all SAC updates at the end of the trajectory.
- This is a slight difference from the SAC specified in the
- original paper.
- """
- for j in range(ep_len):
- 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],
- LossV=outs[3], Q1Vals=outs[4], Q2Vals=outs[5],
- VVals=outs[6], LogPi=outs[7])
-
- logger.store(EpRet=ep_ret, EpLen=ep_len)
- o, r, d, ep_ret, ep_len = env.reset(), 0, False, 0, 0
-
-
- # End of epoch wrap-up
- if t > 0 and t % steps_per_epoch == 0:
- epoch = t // steps_per_epoch
-
- # Save model
- if (epoch % save_freq == 0) or (epoch == epochs-1):
- 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('VVals', 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('LossV', average_only=True)
- logger.log_tabular('Time', time.time()-start_time)
- logger.dump_tabular()
-
-if __name__ == '__main__':
- import argparse
- parser = argparse.ArgumentParser()
- parser.add_argument('--env', type=str, default='HalfCheetah-v2')
- parser.add_argument('--hid', type=int, default=300)
- parser.add_argument('--l', type=int, default=1)
- parser.add_argument('--gamma', type=float, default=0.99)
- parser.add_argument('--seed', '-s', type=int, default=0)
- parser.add_argument('--epochs', type=int, default=50)
- parser.add_argument('--exp_name', type=str, default='sac')
- args = parser.parse_args()
-
- from spinup.utils.run_utils import setup_logger_kwargs
- logger_kwargs = setup_logger_kwargs(args.exp_name, args.seed)
-
- sac(lambda : gym.make(args.env), actor_critic=core.mlp_actor_critic,
- ac_kwargs=dict(hidden_sizes=[args.hid]*args.l),
- gamma=args.gamma, seed=args.seed, epochs=args.epochs,
- logger_kwargs=logger_kwargs)
-
-import numpy as np
-import tensorflow as tf
-import gym
-import time
-from spinup.algos.td3 import core
-from spinup.algos.td3.core import get_vars
-from spinup.utils.logx import EpochLogger
-
-
-class ReplayBuffer:
- """
- A simple FIFO experience replay buffer for TD3 agents.
- """
-
- def __init__(self, obs_dim, act_dim, size):
- self.obs1_buf = np.zeros([size, obs_dim], dtype=np.float32)
- self.obs2_buf = np.zeros([size, obs_dim], dtype=np.float32)
- self.acts_buf = np.zeros([size, act_dim], dtype=np.float32)
- self.rews_buf = np.zeros(size, dtype=np.float32)
- self.done_buf = np.zeros(size, dtype=np.float32)
- self.ptr, self.size, self.max_size = 0, 0, size
-
- def store(self, obs, act, rew, next_obs, done):
- self.obs1_buf[self.ptr] = obs
- self.obs2_buf[self.ptr] = next_obs
- self.acts_buf[self.ptr] = act
- self.rews_buf[self.ptr] = rew
- self.done_buf[self.ptr] = done
- self.ptr = (self.ptr+1) % self.max_size
- self.size = min(self.size+1, self.max_size)
-
- def sample_batch(self, batch_size=32):
- idxs = np.random.randint(0, self.size, size=batch_size)
- return dict(obs1=self.obs1_buf[idxs],
- obs2=self.obs2_buf[idxs],
- acts=self.acts_buf[idxs],
- rews=self.rews_buf[idxs],
- done=self.done_buf[idxs])
-
-"""
-
-TD3 (Twin Delayed DDPG)
-
-"""
-[docs]def td3(env_fn, actor_critic=core.mlp_actor_critic, ac_kwargs=dict(), seed=0,
- steps_per_epoch=5000, 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,
- act_noise=0.1, target_noise=0.2, noise_clip=0.5, policy_delay=2,
- max_ep_len=1000, logger_kwargs=dict(), save_freq=1):
- """
-
- 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.
-
- 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.
-
- 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(n=10):
- for j in range(n):
- o, r, d, ep_ret, ep_len = test_env.reset(), 0, 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, r, d, ep_ret, ep_len = env.reset(), 0, False, 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
-
- if d or (ep_len == max_ep_len):
- """
- Perform all TD3 updates at the end of the trajectory
- (in accordance with source code of TD3 published by
- original authors).
- """
- for j in range(ep_len):
- 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])
-
- logger.store(EpRet=ep_ret, EpLen=ep_len)
- o, r, d, ep_ret, ep_len = env.reset(), 0, False, 0, 0
-
- # End of epoch wrap-up
- if t > 0 and t % steps_per_epoch == 0:
- epoch = t // steps_per_epoch
-
- # Save model
- if (epoch % save_freq == 0) or (epoch == epochs-1):
- 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()
-
-if __name__ == '__main__':
- import argparse
- parser = argparse.ArgumentParser()
- parser.add_argument('--env', type=str, default='HalfCheetah-v2')
- parser.add_argument('--hid', type=int, default=300)
- parser.add_argument('--l', type=int, default=1)
- parser.add_argument('--gamma', type=float, default=0.99)
- parser.add_argument('--seed', '-s', type=int, default=0)
- parser.add_argument('--epochs', type=int, default=50)
- parser.add_argument('--exp_name', type=str, default='td3')
- args = parser.parse_args()
-
- from spinup.utils.run_utils import setup_logger_kwargs
- logger_kwargs = setup_logger_kwargs(args.exp_name, args.seed)
-
- td3(lambda : gym.make(args.env), actor_critic=core.mlp_actor_critic,
- ac_kwargs=dict(hidden_sizes=[args.hid]*args.l),
- gamma=args.gamma, seed=args.seed, epochs=args.epochs,
- logger_kwargs=logger_kwargs)
-
-import numpy as np
-import tensorflow as tf
-import gym
-import time
-import spinup.algos.trpo.core as core
-from spinup.utils.logx import EpochLogger
-from spinup.utils.mpi_tf import MpiAdamOptimizer, sync_all_params
-from spinup.utils.mpi_tools import mpi_fork, mpi_avg, proc_id, mpi_statistics_scalar, num_procs
-
-
-EPS = 1e-8
-
-class GAEBuffer:
- """
- A buffer for storing trajectories experienced by a TRPO agent interacting
- with the environment, and using Generalized Advantage Estimation (GAE-Lambda)
- for calculating the advantages of state-action pairs.
- """
-
- def __init__(self, obs_dim, act_dim, size, info_shapes, gamma=0.99, lam=0.95):
- self.obs_buf = np.zeros(core.combined_shape(size, obs_dim), dtype=np.float32)
- self.act_buf = np.zeros(core.combined_shape(size, act_dim), dtype=np.float32)
- self.adv_buf = np.zeros(size, dtype=np.float32)
- self.rew_buf = np.zeros(size, dtype=np.float32)
- self.ret_buf = np.zeros(size, dtype=np.float32)
- self.val_buf = np.zeros(size, dtype=np.float32)
- self.logp_buf = np.zeros(size, dtype=np.float32)
- self.info_bufs = {k: np.zeros([size] + list(v), dtype=np.float32) for k,v in info_shapes.items()}
- self.sorted_info_keys = core.keys_as_sorted_list(self.info_bufs)
- self.gamma, self.lam = gamma, lam
- self.ptr, self.path_start_idx, self.max_size = 0, 0, size
-
- def store(self, obs, act, rew, val, logp, info):
- """
- Append one timestep of agent-environment interaction to the buffer.
- """
- assert self.ptr < self.max_size # buffer has to have room so you can store
- self.obs_buf[self.ptr] = obs
- self.act_buf[self.ptr] = act
- self.rew_buf[self.ptr] = rew
- self.val_buf[self.ptr] = val
- self.logp_buf[self.ptr] = logp
- for i, k in enumerate(self.sorted_info_keys):
- self.info_bufs[k][self.ptr] = info[i]
- self.ptr += 1
-
- def finish_path(self, last_val=0):
- """
- Call this at the end of a trajectory, or when one gets cut off
- by an epoch ending. This looks back in the buffer to where the
- trajectory started, and uses rewards and value estimates from
- the whole trajectory to compute advantage estimates with GAE-Lambda,
- as well as compute the rewards-to-go for each state, to use as
- the targets for the value function.
-
- The "last_val" argument should be 0 if the trajectory ended
- because the agent reached a terminal state (died), and otherwise
- should be V(s_T), the value function estimated for the last state.
- This allows us to bootstrap the reward-to-go calculation to account
- for timesteps beyond the arbitrary episode horizon (or epoch cutoff).
- """
-
- path_slice = slice(self.path_start_idx, self.ptr)
- rews = np.append(self.rew_buf[path_slice], last_val)
- vals = np.append(self.val_buf[path_slice], last_val)
-
- # the next two lines implement GAE-Lambda advantage calculation
- deltas = rews[:-1] + self.gamma * vals[1:] - vals[:-1]
- self.adv_buf[path_slice] = core.discount_cumsum(deltas, self.gamma * self.lam)
-
- # the next line computes rewards-to-go, to be targets for the value function
- self.ret_buf[path_slice] = core.discount_cumsum(rews, self.gamma)[:-1]
-
- self.path_start_idx = self.ptr
-
- def get(self):
- """
- Call this at the end of an epoch to get all of the data from
- the buffer, with advantages appropriately normalized (shifted to have
- mean zero and std one). Also, resets some pointers in the buffer.
- """
- assert self.ptr == self.max_size # buffer has to be full before you can get
- self.ptr, self.path_start_idx = 0, 0
- # the next two lines implement the advantage normalization trick
- adv_mean, adv_std = mpi_statistics_scalar(self.adv_buf)
- self.adv_buf = (self.adv_buf - adv_mean) / adv_std
- return [self.obs_buf, self.act_buf, self.adv_buf, self.ret_buf,
- self.logp_buf] + core.values_as_sorted_list(self.info_bufs)
-
-"""
-
-Trust Region Policy Optimization
-
-(with support for Natural Policy Gradient)
-
-"""
-[docs]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'):
- """
-
- 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, 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):
- 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:]
-
- # save and log
- buf.store(o, a, r, v_t, logp_t, info_t)
- logger.store(VVals=v_t)
-
- o, r, d, _ = env.step(a)
- ep_ret += r
- ep_len += 1
-
- 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 = r 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, r, d, ep_ret, ep_len = env.reset(), 0, False, 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()
-
-if __name__ == '__main__':
- import argparse
- parser = argparse.ArgumentParser()
- parser.add_argument('--env', type=str, default='HalfCheetah-v2')
- parser.add_argument('--hid', type=int, default=64)
- parser.add_argument('--l', type=int, default=2)
- parser.add_argument('--gamma', type=float, default=0.99)
- parser.add_argument('--seed', '-s', type=int, default=0)
- parser.add_argument('--cpu', type=int, default=4)
- parser.add_argument('--steps', type=int, default=4000)
- parser.add_argument('--epochs', type=int, default=50)
- parser.add_argument('--exp_name', type=str, default='trpo')
- args = parser.parse_args()
-
- mpi_fork(args.cpu) # run parallel code with mpi
-
- from spinup.utils.run_utils import setup_logger_kwargs
- logger_kwargs = setup_logger_kwargs(args.exp_name, args.seed)
-
- trpo(lambda : gym.make(args.env), actor_critic=core.mlp_actor_critic,
- ac_kwargs=dict(hidden_sizes=[args.hid]*args.l), gamma=args.gamma,
- seed=args.seed, steps_per_epoch=args.steps, epochs=args.epochs,
- logger_kwargs=logger_kwargs)
-
-import numpy as np
-import tensorflow as tf
-import gym
-import time
-import spinup.algos.vpg.core as core
-from spinup.utils.logx import EpochLogger
-from spinup.utils.mpi_tf import MpiAdamOptimizer, sync_all_params
-from spinup.utils.mpi_tools import mpi_fork, mpi_avg, proc_id, mpi_statistics_scalar, num_procs
-
-
-class VPGBuffer:
- """
- A buffer for storing trajectories experienced by a VPG agent interacting
- with the environment, and using Generalized Advantage Estimation (GAE-Lambda)
- for calculating the advantages of state-action pairs.
- """
-
- def __init__(self, obs_dim, act_dim, size, gamma=0.99, lam=0.95):
- self.obs_buf = np.zeros(core.combined_shape(size, obs_dim), dtype=np.float32)
- self.act_buf = np.zeros(core.combined_shape(size, act_dim), dtype=np.float32)
- self.adv_buf = np.zeros(size, dtype=np.float32)
- self.rew_buf = np.zeros(size, dtype=np.float32)
- self.ret_buf = np.zeros(size, dtype=np.float32)
- self.val_buf = np.zeros(size, dtype=np.float32)
- self.logp_buf = np.zeros(size, dtype=np.float32)
- self.gamma, self.lam = gamma, lam
- self.ptr, self.path_start_idx, self.max_size = 0, 0, size
-
- def store(self, obs, act, rew, val, logp):
- """
- Append one timestep of agent-environment interaction to the buffer.
- """
- assert self.ptr < self.max_size # buffer has to have room so you can store
- self.obs_buf[self.ptr] = obs
- self.act_buf[self.ptr] = act
- self.rew_buf[self.ptr] = rew
- self.val_buf[self.ptr] = val
- self.logp_buf[self.ptr] = logp
- self.ptr += 1
-
- def finish_path(self, last_val=0):
- """
- Call this at the end of a trajectory, or when one gets cut off
- by an epoch ending. This looks back in the buffer to where the
- trajectory started, and uses rewards and value estimates from
- the whole trajectory to compute advantage estimates with GAE-Lambda,
- as well as compute the rewards-to-go for each state, to use as
- the targets for the value function.
-
- The "last_val" argument should be 0 if the trajectory ended
- because the agent reached a terminal state (died), and otherwise
- should be V(s_T), the value function estimated for the last state.
- This allows us to bootstrap the reward-to-go calculation to account
- for timesteps beyond the arbitrary episode horizon (or epoch cutoff).
- """
-
- path_slice = slice(self.path_start_idx, self.ptr)
- rews = np.append(self.rew_buf[path_slice], last_val)
- vals = np.append(self.val_buf[path_slice], last_val)
-
- # the next two lines implement GAE-Lambda advantage calculation
- deltas = rews[:-1] + self.gamma * vals[1:] - vals[:-1]
- self.adv_buf[path_slice] = core.discount_cumsum(deltas, self.gamma * self.lam)
-
- # the next line computes rewards-to-go, to be targets for the value function
- self.ret_buf[path_slice] = core.discount_cumsum(rews, self.gamma)[:-1]
-
- self.path_start_idx = self.ptr
-
- def get(self):
- """
- Call this at the end of an epoch to get all of the data from
- the buffer, with advantages appropriately normalized (shifted to have
- mean zero and std one). Also, resets some pointers in the buffer.
- """
- assert self.ptr == self.max_size # buffer has to be full before you can get
- self.ptr, self.path_start_idx = 0, 0
- # the next two lines implement the advantage normalization trick
- adv_mean, adv_std = mpi_statistics_scalar(self.adv_buf)
- self.adv_buf = (self.adv_buf - adv_mean) / adv_std
- return [self.obs_buf, self.act_buf, self.adv_buf,
- self.ret_buf, self.logp_buf]
-
-
-"""
-
-Vanilla Policy Gradient
-
-(with GAE-Lambda for advantage estimation)
-
-"""
-[docs]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):
- """
-
- 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, 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)})
-
- # save and log
- buf.store(o, a, r, v_t, logp_t)
- logger.store(VVals=v_t)
-
- o, r, d, _ = env.step(a[0])
- ep_ret += r
- ep_len += 1
-
- 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 = r 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, r, d, ep_ret, ep_len = env.reset(), 0, False, 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()
-
-if __name__ == '__main__':
- import argparse
- parser = argparse.ArgumentParser()
- parser.add_argument('--env', type=str, default='HalfCheetah-v2')
- parser.add_argument('--hid', type=int, default=64)
- parser.add_argument('--l', type=int, default=2)
- parser.add_argument('--gamma', type=float, default=0.99)
- parser.add_argument('--seed', '-s', type=int, default=0)
- parser.add_argument('--cpu', type=int, default=4)
- parser.add_argument('--steps', type=int, default=4000)
- parser.add_argument('--epochs', type=int, default=50)
- parser.add_argument('--exp_name', type=str, default='vpg')
- args = parser.parse_args()
-
- mpi_fork(args.cpu) # run parallel code with mpi
-
- from spinup.utils.run_utils import setup_logger_kwargs
- logger_kwargs = setup_logger_kwargs(args.exp_name, args.seed)
-
- vpg(lambda : gym.make(args.env), actor_critic=core.mlp_actor_critic,
- ac_kwargs=dict(hidden_sizes=[args.hid]*args.l), gamma=args.gamma,
- seed=args.seed, steps_per_epoch=args.steps, epochs=args.epochs,
- logger_kwargs=logger_kwargs)
-
-"""
-
-Some simple logging functionality, inspired by rllab's logging.
-
-Logs to a tab-separated-values file (path/to/output_directory/progress.txt)
-
-"""
-import json
-import joblib
-import shutil
-import numpy as np
-import tensorflow as tf
-import os.path as osp, time, atexit, os
-from spinup.utils.mpi_tools import proc_id, mpi_statistics_scalar
-from spinup.utils.serialization_utils import convert_json
-
-color2num = dict(
- gray=30,
- red=31,
- green=32,
- yellow=33,
- blue=34,
- magenta=35,
- cyan=36,
- white=37,
- crimson=38
-)
-
-def colorize(string, color, bold=False, highlight=False):
- """
- Colorize a string.
-
- This function was originally written by John Schulman.
- """
- attr = []
- num = color2num[color]
- if highlight: num += 10
- attr.append(str(num))
- if bold: attr.append('1')
- return '\x1b[%sm%s\x1b[0m' % (';'.join(attr), string)
-
-[docs]def restore_tf_graph(sess, fpath):
- """
- Loads graphs saved by Logger.
-
- Will output a dictionary whose keys and values are from the 'inputs'
- and 'outputs' dict you specified with logger.setup_tf_saver().
-
- Args:
- sess: A Tensorflow session.
- fpath: Filepath to save directory.
-
- Returns:
- A dictionary mapping from keys to tensors in the computation graph
- loaded from ``fpath``.
- """
- tf.saved_model.loader.load(
- sess,
- [tf.saved_model.tag_constants.SERVING],
- fpath
- )
- model_info = joblib.load(osp.join(fpath, 'model_info.pkl'))
- graph = tf.get_default_graph()
- model = dict()
- model.update({k: graph.get_tensor_by_name(v) for k,v in model_info['inputs'].items()})
- model.update({k: graph.get_tensor_by_name(v) for k,v in model_info['outputs'].items()})
- return model
-
-[docs]class Logger:
- """
- A general-purpose logger.
-
- Makes it easy to save diagnostics, hyperparameter configurations, the
- state of a training run, and the trained model.
- """
-
-[docs] def __init__(self, output_dir=None, output_fname='progress.txt', exp_name=None):
- """
- Initialize a Logger.
-
- Args:
- output_dir (string): A directory for saving results to. If
- ``None``, defaults to a temp directory of the form
- ``/tmp/experiments/somerandomnumber``.
-
- output_fname (string): Name for the tab-separated-value file
- containing metrics logged throughout a training run.
- Defaults to ``progress.txt``.
-
- exp_name (string): Experiment name. If you run multiple training
- runs and give them all the same ``exp_name``, the plotter
- will know to group them. (Use case: if you run the same
- hyperparameter configuration with multiple random seeds, you
- should give them all the same ``exp_name``.)
- """
- if proc_id()==0:
- self.output_dir = output_dir or "/tmp/experiments/%i"%int(time.time())
- if osp.exists(self.output_dir):
- print("Warning: Log dir %s already exists! Storing info there anyway."%self.output_dir)
- else:
- os.makedirs(self.output_dir)
- self.output_file = open(osp.join(self.output_dir, output_fname), 'w')
- atexit.register(self.output_file.close)
- print(colorize("Logging data to %s"%self.output_file.name, 'green', bold=True))
- else:
- self.output_dir = None
- self.output_file = None
- self.first_row=True
- self.log_headers = []
- self.log_current_row = {}
- self.exp_name = exp_name
-
-[docs] def log(self, msg, color='green'):
- """Print a colorized message to stdout."""
- if proc_id()==0:
- print(colorize(msg, color, bold=True))
-
-[docs] def log_tabular(self, key, val):
- """
- Log a value of some diagnostic.
-
- Call this only once for each diagnostic quantity, each iteration.
- After using ``log_tabular`` to store values for each diagnostic,
- make sure to call ``dump_tabular`` to write them out to file and
- stdout (otherwise they will not get saved anywhere).
- """
- if self.first_row:
- self.log_headers.append(key)
- else:
- assert key in self.log_headers, "Trying to introduce a new key %s that you didn't include in the first iteration"%key
- assert key not in self.log_current_row, "You already set %s this iteration. Maybe you forgot to call dump_tabular()"%key
- self.log_current_row[key] = val
-
-[docs] def save_config(self, config):
- """
- Log an experiment configuration.
-
- Call this once at the top of your experiment, passing in all important
- config vars as a dict. This will serialize the config to JSON, while
- handling anything which can't be serialized in a graceful way (writing
- as informative a string as possible).
-
- Example use:
-
- .. code-block:: python
-
- logger = EpochLogger(**logger_kwargs)
- logger.save_config(locals())
- """
- config_json = convert_json(config)
- if self.exp_name is not None:
- config_json['exp_name'] = self.exp_name
- if proc_id()==0:
- output = json.dumps(config_json, separators=(',',':\t'), indent=4, sort_keys=True)
- print(colorize('Saving config:\n', color='cyan', bold=True))
- print(output)
- with open(osp.join(self.output_dir, "config.json"), 'w') as out:
- out.write(output)
-
-[docs] def save_state(self, state_dict, itr=None):
- """
- Saves the state of an experiment.
-
- To be clear: this is about saving *state*, not logging diagnostics.
- All diagnostic logging is separate from this function. This function
- will save whatever is in ``state_dict``---usually just a copy of the
- environment---and the most recent parameters for the model you
- previously set up saving for with ``setup_tf_saver``.
-
- Call with any frequency you prefer. If you only want to maintain a
- single state and overwrite it at each call with the most recent
- version, leave ``itr=None``. If you want to keep all of the states you
- save, provide unique (increasing) values for 'itr'.
-
- Args:
- state_dict (dict): Dictionary containing essential elements to
- describe the current state of training.
-
- itr: An int, or None. Current iteration of training.
- """
- if proc_id()==0:
- fname = 'vars.pkl' if itr is None else 'vars%d.pkl'%itr
- try:
- joblib.dump(state_dict, osp.join(self.output_dir, fname))
- except:
- self.log('Warning: could not pickle state_dict.', color='red')
- if hasattr(self, 'tf_saver_elements'):
- self._tf_simple_save(itr)
-
-[docs] def setup_tf_saver(self, sess, inputs, outputs):
- """
- Set up easy model saving for tensorflow.
-
- Call once, after defining your computation graph but before training.
-
- Args:
- sess: The Tensorflow session in which you train your computation
- graph.
-
- inputs (dict): A dictionary that maps from keys of your choice
- to the tensorflow placeholders that serve as inputs to the
- computation graph. Make sure that *all* of the placeholders
- needed for your outputs are included!
-
- outputs (dict): A dictionary that maps from keys of your choice
- to the outputs from your computation graph.
- """
- self.tf_saver_elements = dict(session=sess, inputs=inputs, outputs=outputs)
- self.tf_saver_info = {'inputs': {k:v.name for k,v in inputs.items()},
- 'outputs': {k:v.name for k,v in outputs.items()}}
-
- def _tf_simple_save(self, itr=None):
- """
- Uses simple_save to save a trained model, plus info to make it easy
- to associated tensors to variables after restore.
- """
- if proc_id()==0:
- assert hasattr(self, 'tf_saver_elements'), \
- "First have to setup saving with self.setup_tf_saver"
- fpath = 'simple_save' + ('%d'%itr if itr is not None else '')
- fpath = osp.join(self.output_dir, fpath)
- if osp.exists(fpath):
- # simple_save refuses to be useful if fpath already exists,
- # so just delete fpath if it's there.
- shutil.rmtree(fpath)
- tf.saved_model.simple_save(export_dir=fpath, **self.tf_saver_elements)
- joblib.dump(self.tf_saver_info, osp.join(fpath, 'model_info.pkl'))
-
-[docs] def dump_tabular(self):
- """
- Write all of the diagnostics from the current iteration.
-
- Writes both to stdout, and to the output file.
- """
- if proc_id()==0:
- vals = []
- key_lens = [len(key) for key in self.log_headers]
- max_key_len = max(15,max(key_lens))
- keystr = '%'+'%d'%max_key_len
- fmt = "| " + keystr + "s | %15s |"
- n_slashes = 22 + max_key_len
- print("-"*n_slashes)
- for key in self.log_headers:
- val = self.log_current_row.get(key, "")
- valstr = "%8.3g"%val if hasattr(val, "__float__") else val
- print(fmt%(key, valstr))
- vals.append(val)
- print("-"*n_slashes)
- if self.output_file is not None:
- if self.first_row:
- self.output_file.write("\t".join(self.log_headers)+"\n")
- self.output_file.write("\t".join(map(str,vals))+"\n")
- self.output_file.flush()
- self.log_current_row.clear()
- self.first_row=False
-
-[docs]class EpochLogger(Logger):
- """
- A variant of Logger tailored for tracking average values over epochs.
-
- Typical use case: there is some quantity which is calculated many times
- throughout an epoch, and at the end of the epoch, you would like to
- report the average / std / min / max value of that quantity.
-
- With an EpochLogger, each time the quantity is calculated, you would
- use
-
- .. code-block:: python
-
- epoch_logger.store(NameOfQuantity=quantity_value)
-
- to load it into the EpochLogger's state. Then at the end of the epoch, you
- would use
-
- .. code-block:: python
-
- epoch_logger.log_tabular(NameOfQuantity, **options)
-
- to record the desired values.
- """
-
- def __init__(self, *args, **kwargs):
- super().__init__(*args, **kwargs)
- self.epoch_dict = dict()
-
-[docs] def store(self, **kwargs):
- """
- Save something into the epoch_logger's current state.
-
- Provide an arbitrary number of keyword arguments with numerical
- values.
- """
- for k,v in kwargs.items():
- if not(k in self.epoch_dict.keys()):
- self.epoch_dict[k] = []
- self.epoch_dict[k].append(v)
-
-[docs] def log_tabular(self, key, val=None, with_min_and_max=False, average_only=False):
- """
- Log a value or possibly the mean/std/min/max values of a diagnostic.
-
- Args:
- key (string): The name of the diagnostic. If you are logging a
- diagnostic whose state has previously been saved with
- ``store``, the key here has to match the key you used there.
-
- val: A value for the diagnostic. If you have previously saved
- values for this key via ``store``, do *not* provide a ``val``
- here.
-
- with_min_and_max (bool): If true, log min and max values of the
- diagnostic over the epoch.
-
- average_only (bool): If true, do not log the standard deviation
- of the diagnostic over the epoch.
- """
- if val is not None:
- super().log_tabular(key,val)
- else:
- v = self.epoch_dict[key]
- vals = np.concatenate(v) if isinstance(v[0], np.ndarray) and len(v[0].shape)>0 else v
- stats = mpi_statistics_scalar(vals, with_min_and_max=with_min_and_max)
- super().log_tabular(key if average_only else 'Average' + key, stats[0])
- if not(average_only):
- super().log_tabular('Std'+key, stats[1])
- if with_min_and_max:
- super().log_tabular('Max'+key, stats[3])
- super().log_tabular('Min'+key, stats[2])
- self.epoch_dict[key] = []
-
-[docs] def get_stats(self, key):
- """
- Lets an algorithm ask the logger for mean/std/min/max of a diagnostic.
- """
- v = self.epoch_dict[key]
- vals = np.concatenate(v) if isinstance(v[0], np.ndarray) and len(v[0].shape)>0 else v
- return mpi_statistics_scalar(vals)
-
-import numpy as np
-import tensorflow as tf
-from mpi4py import MPI
-from spinup.utils.mpi_tools import broadcast
-
-
-def flat_concat(xs):
- return tf.concat([tf.reshape(x,(-1,)) for x in xs], axis=0)
-
-def assign_params_from_flat(x, params):
- flat_size = lambda p : int(np.prod(p.shape.as_list())) # the 'int' is important for scalars
- splits = tf.split(x, [flat_size(p) for p in params])
- new_params = [tf.reshape(p_new, p.shape) for p, p_new in zip(params, splits)]
- return tf.group([tf.assign(p, p_new) for p, p_new in zip(params, new_params)])
-
-def sync_params(params):
- get_params = flat_concat(params)
- def _broadcast(x):
- broadcast(x)
- return x
- synced_params = tf.py_func(_broadcast, [get_params], tf.float32)
- return assign_params_from_flat(synced_params, params)
-
-[docs]def sync_all_params():
- """Sync all tf variables across MPI processes."""
- return sync_params(tf.global_variables())
-
-
-[docs]class MpiAdamOptimizer(tf.train.AdamOptimizer):
- """
- Adam optimizer that averages gradients across MPI processes.
-
- The compute_gradients method is taken from Baselines `MpiAdamOptimizer`_.
- For documentation on method arguments, see the Tensorflow docs page for
- the base `AdamOptimizer`_.
-
- .. _`MpiAdamOptimizer`: https://github.com/openai/baselines/blob/master/baselines/common/mpi_adam_optimizer.py
- .. _`AdamOptimizer`: https://www.tensorflow.org/api_docs/python/tf/train/AdamOptimizer
- """
-
- def __init__(self, **kwargs):
- self.comm = MPI.COMM_WORLD
- tf.train.AdamOptimizer.__init__(self, **kwargs)
-
-[docs] def compute_gradients(self, loss, var_list, **kwargs):
- """
- Same as normal compute_gradients, except average grads over processes.
- """
- grads_and_vars = super().compute_gradients(loss, var_list, **kwargs)
- grads_and_vars = [(g, v) for g, v in grads_and_vars if g is not None]
- flat_grad = flat_concat([g for g, v in grads_and_vars])
- shapes = [v.shape.as_list() for g, v in grads_and_vars]
- sizes = [int(np.prod(s)) for s in shapes]
-
- num_tasks = self.comm.Get_size()
- buf = np.zeros(flat_grad.shape, np.float32)
-
- def _collect_grads(flat_grad):
- self.comm.Allreduce(flat_grad, buf, op=MPI.SUM)
- np.divide(buf, float(num_tasks), out=buf)
- return buf
-
- avg_flat_grad = tf.py_func(_collect_grads, [flat_grad], tf.float32)
- avg_flat_grad.set_shape(flat_grad.shape)
- avg_grads = tf.split(avg_flat_grad, sizes, axis=0)
- avg_grads_and_vars = [(tf.reshape(g, v.shape), v)
- for g, (_, v) in zip(avg_grads, grads_and_vars)]
-
- return avg_grads_and_vars
-
-[docs] def apply_gradients(self, grads_and_vars, global_step=None, name=None):
- """
- Same as normal apply_gradients, except sync params after update.
- """
- opt = super().apply_gradients(grads_and_vars, global_step, name)
- with tf.control_dependencies([opt]):
- sync = sync_params([v for g,v in grads_and_vars])
- return tf.group([opt, sync])
-
-from mpi4py import MPI
-import os, subprocess, sys
-import numpy as np
-
-
-[docs]def mpi_fork(n, bind_to_core=False):
- """
- Re-launches the current script with workers linked by MPI.
-
- Also, terminates the original process that launched it.
-
- Taken almost without modification from the Baselines function of the
- `same name`_.
-
- .. _`same name`: https://github.com/openai/baselines/blob/master/baselines/common/mpi_fork.py
-
- Args:
- n (int): Number of process to split into.
-
- bind_to_core (bool): Bind each MPI process to a core.
- """
- if n<=1:
- return
- if os.getenv("IN_MPI") is None:
- env = os.environ.copy()
- env.update(
- MKL_NUM_THREADS="1",
- OMP_NUM_THREADS="1",
- IN_MPI="1"
- )
- args = ["mpirun", "-np", str(n)]
- if bind_to_core:
- args += ["-bind-to", "core"]
- args += [sys.executable] + sys.argv
- subprocess.check_call(args, env=env)
- sys.exit()
-
-
-def msg(m, string=''):
- print(('Message from %d: %s \t '%(MPI.COMM_WORLD.Get_rank(), string))+str(m))
-
-
-
-def allreduce(*args, **kwargs):
- return MPI.COMM_WORLD.Allreduce(*args, **kwargs)
-
-
-
-def broadcast(x, root=0):
- MPI.COMM_WORLD.Bcast(x, root=root)
-
-def mpi_op(x, op):
- x, scalar = ([x], True) if np.isscalar(x) else (x, False)
- x = np.asarray(x, dtype=np.float32)
- buff = np.zeros_like(x, dtype=np.float32)
- allreduce(x, buff, op=op)
- return buff[0] if scalar else buff
-
-def mpi_sum(x):
- return mpi_op(x, MPI.SUM)
-
-[docs]def mpi_avg(x):
- """Average a scalar or vector over MPI processes."""
- return mpi_sum(x) / num_procs()
-
-[docs]def mpi_statistics_scalar(x, with_min_and_max=False):
- """
- Get mean/std and optional min/max of scalar x across MPI processes.
-
- Args:
- x: An array containing samples of the scalar to produce statistics
- for.
-
- with_min_and_max (bool): If true, return min and max of x in
- addition to mean and std.
- """
- x = np.array(x, dtype=np.float32)
- global_sum, global_n = mpi_sum([np.sum(x), len(x)])
- mean = global_sum / global_n
-
- global_sum_sq = mpi_sum(np.sum((x - mean)**2))
- std = np.sqrt(global_sum_sq / global_n) # compute global std
-
- if with_min_and_max:
- global_min = mpi_op(np.min(x) if len(x) > 0 else np.inf, op=MPI.MIN)
- global_max = mpi_op(np.max(x) if len(x) > 0 else -np.inf, op=MPI.MAX)
- return mean, std, global_min, global_max
- return mean, std
-
-from spinup.user_config import DEFAULT_DATA_DIR, FORCE_DATESTAMP, \
- DEFAULT_SHORTHAND, WAIT_BEFORE_LAUNCH
-from spinup.utils.logx import colorize
-from spinup.utils.mpi_tools import mpi_fork, msg
-from spinup.utils.serialization_utils import convert_json
-import base64
-from copy import deepcopy
-import cloudpickle
-import json
-import numpy as np
-import os
-import os.path as osp
-import psutil
-import string
-import subprocess
-from subprocess import CalledProcessError
-import sys
-from textwrap import dedent
-import time
-from tqdm import trange
-import zlib
-
-DIV_LINE_WIDTH = 80
-
-[docs]def setup_logger_kwargs(exp_name, seed=None, data_dir=None, datestamp=False):
- """
- Sets up the output_dir for a logger and returns a dict for logger kwargs.
-
- If no seed is given and datestamp is false,
-
- ::
-
- output_dir = data_dir/exp_name
-
- If a seed is given and datestamp is false,
-
- ::
-
- output_dir = data_dir/exp_name/exp_name_s[seed]
-
- If datestamp is true, amend to
-
- ::
-
- output_dir = data_dir/YY-MM-DD_exp_name/YY-MM-DD_HH-MM-SS_exp_name_s[seed]
-
- You can force datestamp=True by setting ``FORCE_DATESTAMP=True`` in
- ``spinup/user_config.py``.
-
- Args:
-
- exp_name (string): Name for experiment.
-
- seed (int): Seed for random number generators used by experiment.
-
- data_dir (string): Path to folder where results should be saved.
- Default is the ``DEFAULT_DATA_DIR`` in ``spinup/user_config.py``.
-
- datestamp (bool): Whether to include a date and timestamp in the
- name of the save directory.
-
- Returns:
-
- logger_kwargs, a dict containing output_dir and exp_name.
- """
-
- # Datestamp forcing
- datestamp = datestamp or FORCE_DATESTAMP
-
- # Make base path
- ymd_time = time.strftime("%Y-%m-%d_") if datestamp else ''
- relpath = ''.join([ymd_time, exp_name])
-
- if seed is not None:
- # Make a seed-specific subfolder in the experiment directory.
- if datestamp:
- hms_time = time.strftime("%Y-%m-%d_%H-%M-%S")
- subfolder = ''.join([hms_time, '-', exp_name, '_s', str(seed)])
- else:
- subfolder = ''.join([exp_name, '_s', str(seed)])
- relpath = osp.join(relpath, subfolder)
-
- data_dir = data_dir or DEFAULT_DATA_DIR
- logger_kwargs = dict(output_dir=osp.join(data_dir, relpath),
- exp_name=exp_name)
- return logger_kwargs
-
-
-[docs]def call_experiment(exp_name, thunk, seed=0, num_cpu=1, data_dir=None,
- datestamp=False, **kwargs):
- """
- Run a function (thunk) with hyperparameters (kwargs), plus configuration.
-
- This wraps a few pieces of functionality which are useful when you want
- to run many experiments in sequence, including logger configuration and
- splitting into multiple processes for MPI.
-
- There's also a SpinningUp-specific convenience added into executing the
- thunk: if ``env_name`` is one of the kwargs passed to call_experiment, it's
- assumed that the thunk accepts an argument called ``env_fn``, and that
- the ``env_fn`` should make a gym environment with the given ``env_name``.
-
- The way the experiment is actually executed is slightly complicated: the
- function is serialized to a string, and then ``run_entrypoint.py`` is
- executed in a subprocess call with the serialized string as an argument.
- ``run_entrypoint.py`` unserializes the function call and executes it.
- We choose to do it this way---instead of just calling the function
- directly here---to avoid leaking state between successive experiments.
-
- Args:
-
- exp_name (string): Name for experiment.
-
- thunk (callable): A python function.
-
- seed (int): Seed for random number generators.
-
- num_cpu (int): Number of MPI processes to split into. Also accepts
- 'auto', which will set up as many procs as there are cpus on
- the machine.
-
- data_dir (string): Used in configuring the logger, to decide where
- to store experiment results. Note: if left as None, data_dir will
- default to ``DEFAULT_DATA_DIR`` from ``spinup/user_config.py``.
-
- **kwargs: All kwargs to pass to thunk.
-
- """
-
- # Determine number of CPU cores to run on
- num_cpu = psutil.cpu_count(logical=False) if num_cpu=='auto' else num_cpu
-
- # Send random seed to thunk
- kwargs['seed'] = seed
-
- # Be friendly and print out your kwargs, so we all know what's up
- print(colorize('Running experiment:\n', color='cyan', bold=True))
- print(exp_name + '\n')
- print(colorize('with kwargs:\n', color='cyan', bold=True))
- kwargs_json = convert_json(kwargs)
- print(json.dumps(kwargs_json, separators=(',',':\t'), indent=4, sort_keys=True))
- print('\n')
-
- # Set up logger output directory
- if 'logger_kwargs' not in kwargs:
- kwargs['logger_kwargs'] = setup_logger_kwargs(exp_name, seed, data_dir, datestamp)
- else:
- print('Note: Call experiment is not handling logger_kwargs.\n')
-
- def thunk_plus():
- # Make 'env_fn' from 'env_name'
- if 'env_name' in kwargs:
- import gym
- env_name = kwargs['env_name']
- kwargs['env_fn'] = lambda : gym.make(env_name)
- del kwargs['env_name']
-
- # Fork into multiple processes
- mpi_fork(num_cpu)
-
- # Run thunk
- thunk(**kwargs)
-
- # Prepare to launch a script to run the experiment
- pickled_thunk = cloudpickle.dumps(thunk_plus)
- encoded_thunk = base64.b64encode(zlib.compress(pickled_thunk)).decode('utf-8')
-
- entrypoint = osp.join(osp.abspath(osp.dirname(__file__)),'run_entrypoint.py')
- cmd = [sys.executable if sys.executable else 'python', entrypoint, encoded_thunk]
- try:
- subprocess.check_call(cmd, env=os.environ)
- except CalledProcessError:
- err_msg = '\n'*3 + '='*DIV_LINE_WIDTH + '\n' + dedent("""
-
- There appears to have been an error in your experiment.
-
- Check the traceback above to see what actually went wrong. The
- traceback below, included for completeness (but probably not useful
- for diagnosing the error), shows the stack leading up to the
- experiment launch.
-
- """) + '='*DIV_LINE_WIDTH + '\n'*3
- print(err_msg)
- raise
-
- # Tell the user about where results are, and how to check them
- logger_kwargs = kwargs['logger_kwargs']
-
- plot_cmd = 'python -m spinup.run plot '+logger_kwargs['output_dir']
- plot_cmd = colorize(plot_cmd, 'green')
-
- test_cmd = 'python -m spinup.run test_policy '+logger_kwargs['output_dir']
- test_cmd = colorize(test_cmd, 'green')
-
- output_msg = '\n'*5 + '='*DIV_LINE_WIDTH +'\n' + dedent("""\
- End of experiment.
-
-
- Plot results from this run with:
-
- %s
-
-
- Watch the trained agent with:
-
- %s
-
-
- """%(plot_cmd,test_cmd)) + '='*DIV_LINE_WIDTH + '\n'*5
-
- print(output_msg)
-
-
-def all_bools(vals):
- return all([isinstance(v,bool) for v in vals])
-
-def valid_str(v):
- """
- Convert a value or values to a string which could go in a filepath.
-
- Partly based on `this gist`_.
-
- .. _`this gist`: https://gist.github.com/seanh/93666
-
- """
- if hasattr(v, '__name__'):
- return valid_str(v.__name__)
-
- if isinstance(v, tuple) or isinstance(v, list):
- return '-'.join([valid_str(x) for x in v])
-
- # Valid characters are '-', '_', and alphanumeric. Replace invalid chars
- # with '-'.
- str_v = str(v).lower()
- valid_chars = "-_%s%s" % (string.ascii_letters, string.digits)
- str_v = ''.join(c if c in valid_chars else '-' for c in str_v)
- return str_v
-
-
-[docs]class ExperimentGrid:
- """
- Tool for running many experiments given hyperparameter ranges.
- """
-
- def __init__(self, name=''):
- self.keys = []
- self.vals = []
- self.shs = []
- self.in_names = []
- self.name(name)
-
- def name(self, _name):
- assert isinstance(_name, str), "Name has to be a string."
- self._name = _name
-
-[docs] def print(self):
- """Print a helpful report about the experiment grid."""
- print('='*DIV_LINE_WIDTH)
-
- # Prepare announcement at top of printing. If the ExperimentGrid has a
- # short name, write this as one line. If the name is long, break the
- # announcement over two lines.
- base_msg = 'ExperimentGrid %s runs over parameters:\n'
- name_insert = '['+self._name+']'
- if len(base_msg%name_insert) <= 80:
- msg = base_msg%name_insert
- else:
- msg = base_msg%(name_insert+'\n')
- print(colorize(msg, color='green', bold=True))
-
- # List off parameters, shorthands, and possible values.
- for k, v, sh in zip(self.keys, self.vals, self.shs):
- color_k = colorize(k.ljust(40), color='cyan', bold=True)
- print('', color_k, '['+sh+']' if sh is not None else '', '\n')
- for i, val in enumerate(v):
- print('\t' + str(convert_json(val)))
- print()
-
- # Count up the number of variants. The number counting seeds
- # is the total number of experiments that will run; the number not
- # counting seeds is the total number of otherwise-unique configs
- # being investigated.
- nvars_total = int(np.prod([len(v) for v in self.vals]))
- if 'seed' in self.keys:
- num_seeds = len(self.vals[self.keys.index('seed')])
- nvars_seedless = int(nvars_total / num_seeds)
- else:
- nvars_seedless = nvars_total
- print(' Variants, counting seeds: '.ljust(40), nvars_total)
- print(' Variants, not counting seeds: '.ljust(40), nvars_seedless)
- print()
- print('='*DIV_LINE_WIDTH)
-
-
- def _default_shorthand(self, key):
- # Create a default shorthand for the key, built from the first
- # three letters of each colon-separated part.
- # But if the first three letters contains something which isn't
- # alphanumeric, shear that off.
- valid_chars = "%s%s" % (string.ascii_letters, string.digits)
- def shear(x):
- return ''.join(z for z in x[:3] if z in valid_chars)
- sh = '-'.join([shear(x) for x in key.split(':')])
- return sh
-
-[docs] def add(self, key, vals, shorthand=None, in_name=False):
- """
- Add a parameter (key) to the grid config, with potential values (vals).
-
- By default, if a shorthand isn't given, one is automatically generated
- from the key using the first three letters of each colon-separated
- term. To disable this behavior, change ``DEFAULT_SHORTHAND`` in the
- ``spinup/user_config.py`` file to ``False``.
-
- Args:
- key (string): Name of parameter.
-
- vals (value or list of values): Allowed values of parameter.
-
- shorthand (string): Optional, shortened name of parameter. For
- example, maybe the parameter ``steps_per_epoch`` is shortened
- to ``steps``.
-
- in_name (bool): When constructing variant names, force the
- inclusion of this parameter into the name.
- """
- assert isinstance(key, str), "Key must be a string."
- assert shorthand is None or isinstance(shorthand, str), \
- "Shorthand must be a string."
- if not isinstance(vals, list):
- vals = [vals]
- if DEFAULT_SHORTHAND and shorthand is None:
- shorthand = self._default_shorthand(key)
- self.keys.append(key)
- self.vals.append(vals)
- self.shs.append(shorthand)
- self.in_names.append(in_name)
-
-[docs] def variant_name(self, variant):
- """
- Given a variant (dict of valid param/value pairs), make an exp_name.
-
- A variant's name is constructed as the grid name (if you've given it
- one), plus param names (or shorthands if available) and values
- separated by underscores.
-
- Note: if ``seed`` is a parameter, it is not included in the name.
- """
-
- def get_val(v, k):
- # Utility method for getting the correct value out of a variant
- # given as a nested dict. Assumes that a parameter name, k,
- # describes a path into the nested dict, such that k='a:b:c'
- # corresponds to value=variant['a']['b']['c']. Uses recursion
- # to get this.
- if k in v:
- return v[k]
- else:
- splits = k.split(':')
- k0, k1 = splits[0], ':'.join(splits[1:])
- return get_val(v[k0], k1)
-
- # Start the name off with the name of the variant generator.
- var_name = self._name
-
- # Build the rest of the name by looping through all parameters,
- # and deciding which ones need to go in there.
- for k, v, sh, inn in zip(self.keys, self.vals, self.shs, self.in_names):
-
- # Include a parameter in a name if either 1) it can take multiple
- # values, or 2) the user specified that it must appear in the name.
- # Except, however, when the parameter is 'seed'. Seed is handled
- # differently so that runs of the same experiment, with different
- # seeds, will be grouped by experiment name.
- if (len(v)>1 or inn) and not(k=='seed'):
-
- # Use the shorthand if available, otherwise the full name.
- param_name = sh if sh is not None else k
- param_name = valid_str(param_name)
-
- # Get variant value for parameter k
- variant_val = get_val(variant, k)
-
- # Append to name
- if all_bools(v):
- # If this is a param which only takes boolean values,
- # only include in the name if it's True for this variant.
- var_name += ('_' + param_name) if variant_val else ''
- else:
- var_name += '_' + param_name + valid_str(variant_val)
-
- return var_name.lstrip('_')
-
- def _variants(self, keys, vals):
- """
- Recursively builds list of valid variants.
- """
- if len(keys)==1:
- pre_variants = [dict()]
- else:
- pre_variants = self._variants(keys[1:], vals[1:])
-
- variants = []
- for val in vals[0]:
- for pre_v in pre_variants:
- v = {}
- v[keys[0]] = val
- v.update(pre_v)
- variants.append(v)
- return variants
-
-[docs] def variants(self):
- """
- Makes a list of dicts, where each dict is a valid config in the grid.
-
- There is special handling for variant parameters whose names take
- the form
-
- ``'full:param:name'``.
-
- The colons are taken to indicate that these parameters should
- have a nested dict structure. eg, if there are two params,
-
- ==================== ===
- Key Val
- ==================== ===
- ``'base:param:one'`` 1
- ``'base:param:two'`` 2
- ==================== ===
-
- the variant dict will have the structure
-
- .. parsed-literal::
-
- variant = {
- base: {
- param : {
- a : 1,
- b : 2
- }
- }
- }
- """
- flat_variants = self._variants(self.keys, self.vals)
-
- def unflatten_var(var):
- """
- Build the full nested dict version of var, based on key names.
- """
- new_var = dict()
- unflatten_set = set()
-
- for k,v in var.items():
- if ':' in k:
- splits = k.split(':')
- k0 = splits[0]
- assert k0 not in new_var or isinstance(new_var[k0], dict), \
- "You can't assign multiple values to the same key."
-
- if not(k0 in new_var):
- new_var[k0] = dict()
-
- sub_k = ':'.join(splits[1:])
- new_var[k0][sub_k] = v
- unflatten_set.add(k0)
- else:
- assert not(k in new_var), \
- "You can't assign multiple values to the same key."
- new_var[k] = v
-
- # Make sure to fill out the nested dicts.
- for k in unflatten_set:
- new_var[k] = unflatten_var(new_var[k])
-
- return new_var
-
- new_variants = [unflatten_var(var) for var in flat_variants]
- return new_variants
-
-[docs] def run(self, thunk, num_cpu=1, data_dir=None, datestamp=False):
- """
- Run each variant in the grid with function 'thunk'.
-
- Note: 'thunk' must be either a callable function, or a string. If it is
- a string, it must be the name of a parameter whose values are all
- callable functions.
-
- Uses ``call_experiment`` to actually launch each experiment, and gives
- each variant a name using ``self.variant_name()``.
-
- Maintenance note: the args for ExperimentGrid.run should track closely
- to the args for call_experiment. However, ``seed`` is omitted because
- we presume the user may add it as a parameter in the grid.
- """
-
- # Print info about self.
- self.print()
-
- # Make the list of all variants.
- variants = self.variants()
-
- # Print variant names for the user.
- var_names = set([self.variant_name(var) for var in variants])
- var_names = sorted(list(var_names))
- line = '='*DIV_LINE_WIDTH
- preparing = colorize('Preparing to run the following experiments...',
- color='green', bold=True)
- joined_var_names = '\n'.join(var_names)
- announcement = f"\n{preparing}\n\n{joined_var_names}\n\n{line}"
- print(announcement)
-
-
- if WAIT_BEFORE_LAUNCH > 0:
- delay_msg = colorize(dedent("""
- Launch delayed to give you a few seconds to review your experiments.
-
- To customize or disable this behavior, change WAIT_BEFORE_LAUNCH in
- spinup/user_config.py.
-
- """), color='cyan', bold=True)+line
- print(delay_msg)
- wait, steps = WAIT_BEFORE_LAUNCH, 100
- prog_bar = trange(steps, desc='Launching in...',
- leave=False, ncols=DIV_LINE_WIDTH,
- mininterval=0.25,
- bar_format='{desc}: {bar}| {remaining} {elapsed}')
- for _ in prog_bar:
- time.sleep(wait/steps)
-
- # Run the variants.
- for var in variants:
- exp_name = self.variant_name(var)
-
- # Figure out what the thunk is.
- if isinstance(thunk, str):
- # Assume one of the variant parameters has the same
- # name as the string you passed for thunk, and that
- # variant[thunk] is a valid callable function.
- thunk_ = var[thunk]
- del var[thunk]
- else:
- # Assume thunk is given as a function.
- thunk_ = thunk
-
- call_experiment(exp_name, thunk_, num_cpu=num_cpu,
- data_dir=data_dir, datestamp=datestamp, **var)
-
-
-def test_eg():
- eg = ExperimentGrid()
- eg.add('test:a', [1,2,3], 'ta', True)
- eg.add('test:b', [1,2,3])
- eg.add('some', [4,5])
- eg.add('why', [True,False])
- eg.add('huh', 5)
- eg.add('no', 6, in_name=True)
- return eg.variants()
-