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)
#=========================================================================#
# #
# All of your code goes in the space below. #
# #
#=========================================================================#
# 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
pi_noise_targ = get_pi_noise_clipped(pi,
noise_scale=target_noise,
noise_clip=noise_clip,
act_limit=act_limit)
# Target Q-values, using action from smoothed target policy
_, q1_targ, q2_targ, _ = actor_critic(x2_ph, pi_noise_targ,
**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
q_targ = get_q_target(q1_targ, q2_targ, r_ph, d=d_ph, gamma=0.99)
# TD3 losses
pi_loss = -tf.reduce_mean(q1_pi)
q1_loss = tf.losses.mean_squared_error(q_targ, q1)
q2_loss = tf.losses.mean_squared_error(q_targ, q2)
q_loss = q1_loss + q2_loss
#=========================================================================#
# #
# All of your code goes in the space above. #
# #
#=========================================================================#
# 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)})
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()
def td3(env_fn, actor_critic=core.ActorCritic, 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())
torch.manual_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
# Main outputs from computation graph
main = actor_critic(in_features=obs_dim, **ac_kwargs)
# Target policy network
target = actor_critic(in_features=obs_dim, **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(module) for module in
[main.policy, main.q1, main.q2, main])
print('\nNumber of parameters: \t pi: %d, \t q1: %d, \t q2: %d, \t total: %d\n'%var_counts)
# Separate train ops for pi, q
pi_optimizer = torch.optim.Adam(main.policy.parameters(), lr=pi_lr)
q_params = list(main.q1.parameters()) + list(main.q2.parameters())
q_optimizer = torch.optim.Adam(q_params, lr=q_lr)
# Initializing targets to match main variables
target.load_state_dict(main.state_dict())
def get_action(o, noise_scale):
pi = main.policy(torch.Tensor(o.reshape(1,-1)))
a = pi.data.numpy()[0] + 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)
(obs1, obs2, acts, rews, done) = (torch.Tensor(batch['obs1']),
torch.Tensor(batch['obs2']),
torch.Tensor(batch['acts']),
torch.Tensor(batch['rews']),
torch.Tensor(batch['done']))
_, q1, q2, _ = main(obs1, acts)
pi_targ = target.policy(obs2)
# Target policy smoothing, by adding clipped noise to target actions
epsilon = torch.normal(torch.zeros_like(pi_targ),
target_noise*torch.ones_like(pi_targ))
epsilon = torch.clamp(epsilon, -noise_clip, noise_clip)
a2 = torch.clamp(pi_targ + epsilon, -act_limit, act_limit)
# Target Q-values, using action from target policy
_, q1_targ, q2_targ, _ = target(obs2, a2)
# Bellman backup for Q functions, using Clipped Double-Q targets
min_q_targ = torch.min(q1_targ, q2_targ)
backup = (rews + gamma * (1 - done) * min_q_targ).detach()
# TD3 Q losses
q1_loss = torch.mean((q1 - backup)**2)
q2_loss = torch.mean((q2 - backup)**2)
q_loss = q1_loss + q2_loss
q_optimizer.zero_grad()
q_loss.backward()
q_optimizer.step()
logger.store(LossQ=q_loss.item(), Q1Vals=q1.data.numpy(),
Q2Vals=q2.data.numpy())
if j % policy_delay == 0:
_, _, _, q1_pi = main(obs1, acts)
# TD3 policy loss
pi_loss = -torch.mean(q1_pi)
# Delayed policy update
pi_optimizer.zero_grad()
pi_loss.backward()
pi_optimizer.step()
# Polyak averaging for target variables
for p_main, p_target in zip(main.parameters(), target.parameters()):
p_target.data.copy_(polyak*p_target.data + (1 - polyak)*p_main.data)
logger.store(LossPi=pi_loss.item())
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()
def td3(env_fn,
actor_critic=core.mlp_actor_critic,
ac_kwargs=dict(),
seed=0,
steps_per_epoch=5000,
epochs=10000,
replay_size=int(5e3),
gamma=0.99,
polyak=0.995,
pi_lr=1e-4,
q_lr=1e-4,
batch_size=64,
start_epochs=0,
act_noise=0.1,
target_noise=0.2,
noise_clip=0.5,
policy_delay=2,
max_ep_len=500,
policy_path=None,
logger_kwargs=dict(),
save_freq=100,
UPDATE_STEP=10,
test_freq=100):
"""
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())
test_logger_kwargs = dict()
test_logger_kwargs['output_dir'] = osp.join(logger_kwargs['output_dir'],
"test")
test_logger_kwargs['exp_name'] = logger_kwargs['exp_name']
test_logger = EpochLogger(**test_logger_kwargs)
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, do not assumes all dimensions share the same bound!
act_limit = env.action_space.high
act_limit_low = env.action_space.low
act_noise_limit = act_noise * act_limit
# Share information about action space with policy architecture
ac_kwargs['action_space'] = env.action_space
if policy_path is None:
# 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_low, act_limit)
# Target Q-values, using action from target policy
_, q1_targ, q2_targ, _ = actor_critic(x2_ph, a2, **ac_kwargs)
sess = tf.Session()
else:
# sess, x_ph, a_ph, pi, q1, q2, q1_pi = load_policy(fpath=policy_path, itr='last', deterministic=False, **ac_kwargs)
# x2_ph, r_ph, d_ph = core.placeholders( obs_dim, None, None)
# # Target policy network
# # todo: copy parameters from main
# 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)
# 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
var_counts = tuple(
core.count_vars(scope) for scope in ['main/pi', 'main'])
print('\nNumber of parameters: \t pi: %d, \t total: %d\n' % var_counts)
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_low, act_limit)
# Target Q-values, using action from target policy
_, q1_targ, q2_targ, _ = actor_critic(x2_ph, a2, **ac_kwargs)
# sess = tf.Session()
# 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())
# uninitialized_vars = []
# for var in tf.all_variables():
# try:
# sess.run(var)
# except tf.errors.FailedPreconditionError:
# uninitialized_vars.append(var)
# init_new_vars_op = tf.initialize_variables(uninitialized_vars)
# todo: reinitialize pi
pi_ref_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES,
scope='main/pi')
# chkp.print_tensors_in_checkpoint_file("/model/model_ckpt", tensor_name='pi', all_tensors=False, all_tensor_names = '')
saver_pi = tf.train.Saver(pi_ref_vars)
# # model = restore_tf_graph(sess, osp.join(policy_path, 'simple_save'))
# with tf.variable_scope('main', reuse=True):
saver_pi.restore(sess, './model')
# pi = model['pi']
# pass
sess.run(target_init)
# Count variables
var_counts = tuple(
core.count_vars(scope)
for scope in ['main/pi', 'main/q1', 'main/q2', 'main'])
print(
'\n updated Number of parameters: \t pi: %d, \t q1: %d, \t q2: %d, \t total: %d\n'
% var_counts)
# 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]
# todo: add act_limit scale noise
a += noise_scale * np.random.randn(act_dim)
return np.clip(a, act_limit_low, act_limit)
start_time = time.time()
o, r, d, ep_ret, ep_len = env.reset(), 0, False, 0, 0
total_steps = steps_per_epoch * epochs
test_num = 0
test_policy_epochs = 91
episode_steps = 500
# Main loop: collect experience in env and update/log each epoch
for epoch in range(0, epochs):
"""
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).
"""
# test policy
if epoch % test_freq == 0:
env.unwrapped._set_test_mode(True)
for i in range(test_policy_epochs):
observation = env.reset()
policy_cumulated_reward = 0
for t in range(episode_steps):
action = get_action(np.array(observation), act_noise_limit)
newObservation, reward, done, info = env.step(action)
observation = newObservation
if (t == episode_steps - 1):
print("reached the end")
done = True
policy_cumulated_reward += reward
if done:
test_logger.store(
policy_reward=policy_cumulated_reward)
test_logger.store(policy_steps=t)
# print(info)
test_logger.store(arrive_des=info['arrive_des'])
break
else:
pass
test_logger.log_tabular('epoch', epoch)
test_logger.log_tabular('policy_reward', average_only=True)
test_logger.log_tabular('policy_steps', average_only=True)
test_logger.log_tabular('arrive_des', average_only=True)
test_logger.dump_tabular()
# train policy
env.unwrapped._set_test_mode(False)
o, r, d, ep_ret, ep_len = env.reset(), 0, False, 0, 0
for t in range(steps_per_epoch):
if epoch > start_epochs:
a = get_action(np.array(o), act_noise_limit)
else:
a = env.action_space.sample()
# Step the env
o2, r, d, _ = env.step(a)
ep_ret += r
ep_len += 1
if (t == steps_per_epoch - 1):
print("reached the end")
d = True
# 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:
"""
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
break
# End of epoch wrap-up
if epoch > 0 and (epoch % save_freq == 0) or (epoch == epochs - 1):
# Save model
logger.save_state({}, None)
# Test the performance of the deterministic version of the agent.
test_num += 1
# test_agent(test_num=test_num)
# Log info about epoch
logger.log_tabular('Epoch', epoch)
logger.log_tabular('EpRet', with_min_and_max=True)
logger.log_tabular('TestEpRet', with_min_and_max=True)
logger.log_tabular('EpLen', average_only=True)
logger.log_tabular('TestEpLen', average_only=True)
logger.log_tabular('TotalEnvInteracts', t)
logger.log_tabular('Q1Vals', with_min_and_max=True)
logger.log_tabular('Q2Vals', with_min_and_max=True)
logger.log_tabular('LossPi', average_only=True)
logger.log_tabular('LossQ', average_only=True)
logger.log_tabular('Time', time.time() - start_time)
logger.dump_tabular()
def __init__(self, obs_space, act_space, EP_MAX=1000, EP_LEN=250, GAMMA=0.99, A_LR = 0.0001, C_LR=0.0001, BATCH=32, UPDATE_STEP=10, logger_kwargs=dict()):
self.logger = EpochLogger(**logger_kwargs)
self.logger.save_config(locals())
self.EP_MAX = EP_MAX
self.EP_LEN = EP_LEN
self.GAMMA = GAMMA
self.LR = LR
self.BATCH = BATCH
self.UPDATE_STEP = UPDATE_STEP
self.obs_space = obs_space
self.act_space = act_space
self.obs_dim = self.obs_space.shape[0]
self.act_dim = self.act_space.shape[0]
self.act_limit = act_space.high
self.sess = tf.Session()
# self.tfs = tf.placeholder(tf.float32, [None, self.obs_dim], 'state')
# self.tfa = tf.placeholder(tf.float32, [None, self.act_dim], 'action')
tf.set_random_seed(seed)
np.random.seed(seed)
# Inputs to computation graph
self.x_ph, self.a_ph, self.x2_ph, self.r_ph, self.d_ph = core.placeholders(obs_dim, act_dim, obs_dim, None, None)
# Main outputs from computation graph
with tf.variable_scope('main'):
self.pi, self.q1, self.q2, self.q1_pi = actor_critic(x_ph, a_ph, **ac_kwargs)
# Target policy network
with tf.variable_scope('target'):
self.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
_, self.q1_targ, self.q2_targ, _ = actor_critic(x2_ph, a2, **ac_kwargs)
# # Experience buffer
# self.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'))])
def td3(env_fn,
expert=None,
policy_path=None,
actor_critic=core.mlp_actor_critic,
ac_kwargs=dict(),
seed=0,
steps_per_epoch=500,
epochs=1000,
replay_size=int(5e3),
gamma=0.99,
polyak=0.995,
pi_lr=1e-4,
q_lr=1e-4,
batch_size=64,
start_epochs=500,
dagger_epochs=500,
pretrain_epochs=50,
dagger_noise=0.02,
act_noise=0.02,
target_noise=0.02,
noise_clip=0.5,
policy_delay=2,
max_ep_len=500,
logger_kwargs=dict(),
save_freq=50,
UPDATE_STEP=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) | 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())
test_logger_kwargs = dict()
test_logger_kwargs['output_dir'] = osp.join(logger_kwargs['output_dir'],
"test")
test_logger_kwargs['exp_name'] = logger_kwargs['exp_name']
test_logger = EpochLogger(**test_logger_kwargs)
# test_logger_kwargs = dict()
# test_logger_kwargs['output_dir'] = osp.join(logger_kwargs['output_dir'], "test")
# test_logger_kwargs['exp_name'] = logger_kwargs['exp_name']
# test_logger = EpochLogger(**test_logger_kwargs)
# pretrain_logger_kwargs = dict()
# pretrain_logger_kwargs['output_dir'] = osp.join(logger_kwargs['output_dir'], "pretrain")
# pretrain_logger_kwargs['exp_name'] = logger_kwargs['exp_name']
# pretrain_logger = EpochLogger(**pretrain_logger_kwargs)
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, do not assumes all dimensions share the same bound!
act_limit = env.action_space.high / 2
act_high_limit = env.action_space.high
act_low_limit = env.action_space.low
act_noise_limit = act_noise * act_limit
sess = tf.Session()
if policy_path is None:
# 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)
tfa_ph = core.placeholder(act_dim)
# 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_low_limit, act_high_limit)
# Target Q-values, using action from target policy
_, q1_targ, q2_targ, _ = actor_critic(x2_ph, a2, **ac_kwargs)
else:
# sess = tf.Session()
model = restore_tf_graph(sess, osp.join(policy_path, 'simple_save'))
x_ph, a_ph, x2_ph, r_ph, d_ph = model['x_ph'], model['a_ph'], model[
'x2_ph'], model['r_ph'], model['d_ph']
pi, q1, q2, q1_pi = model['pi'], model['q1'], model['q2'], model[
'q1_pi']
pi_targ, q1_targ, q2_targ = model['pi_targ'], model['q1_targ'], model[
'q2_targ']
tfa_ph = core.placeholder(act_dim)
dagger_epochs = 0
# Experience buffer
replay_buffer = ReplayBuffer(obs_dim=obs_dim,
act_dim=act_dim,
size=replay_size)
dagger_replay_buffer = DaggerReplayBuffer(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)
if policy_path is None:
# 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)
# dagger loss
dagger_pi_loss = tf.reduce_mean(tf.square(pi - tfa_ph))
# 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 = tf.add(q1_loss, q2_loss)
pi_loss = tf.identity(pi_loss, name="pi_loss")
q1_loss = tf.identity(q1_loss, name="q1_loss")
q2_loss = tf.identity(q2_loss, name="q2_loss")
q_loss = tf.identity(q_loss, name="q_loss")
# Separate train ops for pi, q
dagger_pi_optimizer = tf.train.AdamOptimizer(learning_rate=pi_lr)
pi_optimizer = tf.train.AdamOptimizer(learning_rate=pi_lr)
q_optimizer = tf.train.AdamOptimizer(learning_rate=q_lr)
train_dagger_pi_op = dagger_pi_optimizer.minimize(
dagger_pi_loss,
var_list=get_vars('main/pi'),
name='train_dagger_pi_op')
train_pi_op = pi_optimizer.minimize(pi_loss,
var_list=get_vars('main/pi'),
name='train_pi_op')
train_q_op = q_optimizer.minimize(q_loss,
var_list=get_vars('main/q'),
name='train_q_op')
# 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.run(tf.global_variables_initializer())
else:
graph = tf.get_default_graph()
# opts = graph.get_operations()
# print (opts)
pi_loss = model['pi_loss']
q1_loss = model['q1_loss']
q2_loss = model['q2_loss']
q_loss = model['q_loss']
train_q_op = graph.get_operation_by_name('train_q_op')
train_pi_op = graph.get_operation_by_name('train_pi_op')
# target_update = graph.get_operation_by_name('target_update')
# target_init = graph.get_operation_by_name('target_init')
# 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_ph': x_ph, 'a_ph': a_ph, 'x2_ph': x2_ph, 'r_ph': r_ph, 'd_ph': d_ph}, \
outputs={'pi': pi, 'q1': q1, 'q2': q2, 'q1_pi': q1_pi, 'pi_targ': pi_targ, 'q1_targ': q1_targ, 'q2_targ': q2_targ, \
'pi_loss': pi_loss, 'q1_loss': q1_loss, 'q2_loss': q2_loss, 'q_loss': q_loss})
def get_action(o, noise_scale):
a = sess.run(pi, feed_dict={x_ph: o.reshape(1, -1)})[0]
# todo: add act_limit scale noise
a += noise_scale * np.random.randn(act_dim)
return np.clip(a, act_low_limit, act_high_limit)
def choose_action(s, add_noise=False):
s = s[np.newaxis, :]
a = sess.run(pi, {x_ph: s})[0]
if add_noise:
noise = dagger_noise * act_high_limit * np.random.normal(
size=a.shape)
a = a + noise
return np.clip(a, act_low_limit, act_high_limit)
def test_agent(n=81, test_num=1):
n = env.unwrapped._set_test_mode(True)
con_flag = False
for j in range(n):
o, r, d, ep_ret, ep_len = 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, info = env.step(choose_action(np.array(o), 0))
ep_ret += r
ep_len += 1
if d:
test_logger.store(TestEpRet=ep_ret, TestEpLen=ep_len)
test_logger.store(arrive_des=info['arrive_des'])
test_logger.store(
arrive_des_appro=info['arrive_des_appro'])
if not info['out_of_range']:
test_logger.store(converge_dis=info['converge_dis'])
con_flag = True
test_logger.store(out_of_range=info['out_of_range'])
# print(info)
# test_logger.dump_tabular()
# time.sleep(10)
if not con_flag:
test_logger.store(converge_dis=10000)
env.unwrapped._set_test_mode(False)
start_time = time.time()
env.unwrapped._set_test_mode(False)
o, r, d, ep_ret, ep_len = env.reset(), 0, False, 0, 0
total_steps = steps_per_epoch * epochs
test_num = 0
total_env_t = 0
print(colorize("begin dagger training", 'green', bold=True))
# Main loop for dagger pretrain
for epoch in range(1, dagger_epochs + 1, 1):
obs, acs, rewards = [], [], []
# number of timesteps
for t in range(steps_per_epoch):
# action = env.action_space.sample()
# action = ppo.choose_action(np.array(observation))
obs.append(o)
ref_action = call_ref_controller(env, expert)
if (epoch < pretrain_epochs):
action = ref_action
else:
action = choose_action(np.array(o), True)
o2, r, d, info = env.step(action)
ep_ret += r
ep_len += 1
total_env_t += 1
acs.append(ref_action)
rewards.append(r)
# Store experience to replay buffer
replay_buffer.store(o, action, r, o2, d)
o = o2
if (t == steps_per_epoch - 1):
# print ("reached the end")
d = True
if d:
# collected data to replaybuffer
max_step = len(np.array(rewards))
q = [
np.sum(
np.power(gamma, np.arange(max_step - t)) * rewards[t:])
for t in range(max_step)
]
dagger_replay_buffer.stores(obs, acs, rewards, q)
# update policy
for _ in range(int(max_step / 5)):
batch = dagger_replay_buffer.sample_batch(batch_size)
feed_dict = {x_ph: batch['obs1'], tfa_ph: batch['acts']}
q_step_ops = [dagger_pi_loss, train_dagger_pi_op]
for j in range(UPDATE_STEP):
outs = sess.run(q_step_ops, feed_dict)
logger.store(LossPi=outs[0])
# train q function
for j in range(int(max_step / 5)):
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]
# for _ in range(UPDATE_STEP):
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 target update
outs = sess.run([target_update], feed_dict)
# logger.store(LossPi=outs[0])
# logger.store(LossQ=1000000, Q1Vals=1000000, Q2Vals=1000000)
logger.store(EpRet=ep_ret, EpLen=ep_len)
o, r, d, ep_ret, ep_len = env.reset(), 0, False, 0, 0
break
# End of epoch wrap-up
if epoch > 0 and (epoch % save_freq == 0) or (epoch == dagger_epochs):
# Save model
logger.save_state({}, None)
# Test the performance of the deterministic version of the agent.
test_num += 1
test_agent(test_num=test_num)
# Log info about epoch
test_logger.log_tabular('epoch', epoch)
test_logger.log_tabular('TestEpRet', average_only=True)
test_logger.log_tabular('TestEpLen', average_only=True)
test_logger.log_tabular('arrive_des', average_only=True)
test_logger.log_tabular('converge_dis', average_only=True)
test_logger.log_tabular('out_of_range', average_only=True)
test_logger.dump_tabular()
sess.run(target_init)
print(colorize("begin td3 training", 'green', bold=True))
# Main loop: collect experience in env and update/log each epoch
# total_env_t = 0
for epoch in range(1, epochs + 1, 1):
# End of epoch wrap-up
if epoch > 0 and (epoch % save_freq == 0) or (epoch == epochs):
# Save model
logger.save_state({}, None)
# Test the performance of the deterministic version of the agent.
test_num += 1
test_agent(test_num=test_num)
# Log info about epoch
test_logger.log_tabular('epoch', epoch)
test_logger.log_tabular('TestEpRet', average_only=True)
test_logger.log_tabular('TestEpLen', average_only=True)
test_logger.log_tabular('arrive_des', average_only=True)
test_logger.log_tabular('converge_dis', average_only=True)
test_logger.log_tabular('out_of_range', average_only=True)
test_logger.dump_tabular()
"""
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).
"""
# o, r, d, ep_ret, ep_len = env.reset(), 0, False, 0, 0
for t in range(steps_per_epoch):
if epoch > start_epochs:
a = get_action(np.array(o), act_noise_limit)
else:
a = env.action_space.sample()
# ref_action = call_ref_controller(env, expert)
# Step the env
o2, r, d, _ = env.step(a)
ep_ret += r
ep_len += 1
total_env_t += 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 (t == steps_per_epoch - 1):
# print ("reached the end")
d = True
if d:
"""
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]
# for _ in range(UPDATE_STEP):
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
break
def td3(env_fn, actor_critic=core.MLPActorCritic, ac_kwargs=dict(), seed=0,
steps_per_epoch=4000, epochs=100, replay_size=int(1e6), gamma=0.99,
polyak=0.995, pi_lr=1e-3, q_lr=1e-3, batch_size=100, start_steps=10000,
update_after=1000, update_every=50, act_noise=0.1, target_noise=0.2,
noise_clip=0.5, policy_delay=2, num_test_episodes=10, max_ep_len=1000,
logger_kwargs=dict(), save_freq=1):
"""
Twin Delayed Deep Deterministic Policy Gradient (TD3)
Args:
env_fn : A function which creates a copy of the environment.
The environment must satisfy the OpenAI Gym API.
actor_critic: The constructor method for a PyTorch Module with an ``act``
method, a ``pi`` module, a ``q1`` module, and a ``q2`` module.
The ``act`` method and ``pi`` module should accept batches of
observations as inputs, and ``q1`` and ``q2`` should accept a batch
of observations and a batch of actions as inputs. When called,
these should return:
=========== ================ ======================================
Call Output Shape Description
=========== ================ ======================================
``act`` (batch, act_dim) | Numpy array of actions for each
| observation.
``pi`` (batch, act_dim) | Tensor containing actions from policy
| given observations.
``q1`` (batch,) | Tensor containing one current estimate
| of Q* for the provided observations
| and actions. (Critical: make sure to
| flatten this!)
``q2`` (batch,) | Tensor containing the other current
| estimate of Q* for the provided observations
| and actions. (Critical: make sure to
| flatten this!)
=========== ================ ======================================
ac_kwargs (dict): Any kwargs appropriate for the ActorCritic object
you provided to TD3.
seed (int): Seed for random number generators.
steps_per_epoch (int): Number of steps of interaction (state-action pairs)
for the agent and the environment in each epoch.
epochs (int): Number of epochs to run and train agent.
replay_size (int): Maximum length of replay buffer.
gamma (float): Discount factor. (Always between 0 and 1.)
polyak (float): Interpolation factor in polyak averaging for target
networks. Target networks are updated towards main networks
according to:
.. math:: \\theta_{\\text{targ}} \\leftarrow
\\rho \\theta_{\\text{targ}} + (1-\\rho) \\theta
where :math:`\\rho` is polyak. (Always between 0 and 1, usually
close to 1.)
pi_lr (float): Learning rate for policy.
q_lr (float): Learning rate for Q-networks.
batch_size (int): Minibatch size for SGD.
start_steps (int): Number of steps for uniform-random action selection,
before running real policy. Helps exploration.
update_after (int): Number of env interactions to collect before
starting to do gradient descent updates. Ensures replay buffer
is full enough for useful updates.
update_every (int): Number of env interactions that should elapse
between gradient descent updates. Note: Regardless of how long
you wait between updates, the ratio of env steps to gradient steps
is locked to 1.
act_noise (float): Stddev for Gaussian exploration noise added to
policy at training time. (At test time, no noise is added.)
target_noise (float): Stddev for smoothing noise added to target
policy.
noise_clip (float): Limit for absolute value of target policy
smoothing noise.
policy_delay (int): Policy will only be updated once every
policy_delay times for each update of the Q-networks.
num_test_episodes (int): Number of episodes to test the deterministic
policy at the end of each epoch.
max_ep_len (int): Maximum length of trajectory / episode / rollout.
logger_kwargs (dict): Keyword args for EpochLogger.
save_freq (int): How often (in terms of gap between epochs) to save
the current policy and value function.
"""
device = torch.device(
"cuda") if torch.cuda.is_available() else torch.device("cpu")
# device = 'cpu'
print(device)
logger = EpochLogger(**logger_kwargs)
print(logger_kwargs)
# exit(0)
logger.save_config(locals())
torch.manual_seed(seed)
np.random.seed(seed)
env, test_env = env_fn(), env_fn()
obs_dim = env.observation_space.shape
act_dim = env.action_space.shape[0]
# Action limit for clamping: critically, assumes all dimensions share the same bound!
act_limit = env.action_space.high[0]
# Create actor-critic module and target networks
ac = actor_critic(env.observation_space,
env.action_space, **ac_kwargs).to(device)
ac_targ = deepcopy(ac).to(device)
# Freeze target networks with respect to optimizers (only update via polyak averaging)
for p in ac_targ.parameters():
p.requires_grad = False
# List of parameters for both Q-networks (save this for convenience)
q_params = itertools.chain(ac.q1.parameters(), ac.q2.parameters())
# Experience buffer
replay_buffer = ReplayBuffer(
obs_dim=obs_dim, act_dim=act_dim, size=replay_size)
# Count variables (protip: try to get a feel for how different size networks behave!)
var_counts = tuple(core.count_vars(module)
for module in [ac.pi, ac.q1, ac.q2])
logger.log(
'\nNumber of parameters: \t pi: %d, \t q1: %d, \t q2: %d\n' % var_counts)
# Set up function for computing TD3 Q-losses
def compute_loss_q(data):
o, a, r, o2, d = data['obs'], data['act'], data['rew'], data['obs2'], data['done']
o, a, r, o2, d = o.to(device), a.to(device), r.to(
device), o2.to(device), d.to(device)
q1 = ac.q1(o, a)
q2 = ac.q2(o, a)
# Bellman backup for Q functions
with torch.no_grad():
pi_targ = ac_targ.pi(o2)
# Target policy smoothing
epsilon = torch.randn_like(pi_targ) * target_noise
epsilon = torch.clamp(epsilon, -noise_clip, noise_clip)
a2 = pi_targ + epsilon
a2 = torch.clamp(a2, -act_limit, act_limit)
# Target Q-values
q1_pi_targ = ac_targ.q1(o2, a2)
q2_pi_targ = ac_targ.q2(o2, a2)
q_pi_targ = torch.min(q1_pi_targ, q2_pi_targ)
backup = r + gamma * (1 - d) * q_pi_targ
# MSE loss against Bellman backup
loss_q1 = ((q1 - backup)**2).mean()
loss_q2 = ((q2 - backup)**2).mean()
loss_q = loss_q1 + loss_q2
# Useful info for logging
loss_info = dict(Q1Vals=q1.detach().cpu().numpy(),
Q2Vals=q2.detach().cpu().numpy())
return loss_q, loss_info
# Set up function for computing TD3 pi loss
def compute_loss_pi(data):
o = data['obs'].to(device)
q1_pi = ac.q1(o, ac.pi(o))
return -q1_pi.mean()
# Set up optimizers for policy and q-function
pi_optimizer = Adam(ac.pi.parameters(), lr=pi_lr)
q_optimizer = Adam(q_params, lr=q_lr)
# Set up model saving
logger.setup_pytorch_saver(ac)
def update(data, timer):
# First run one gradient descent step for Q1 and Q2
q_optimizer.zero_grad()
loss_q, loss_info = compute_loss_q(data)
loss_q.backward()
q_optimizer.step()
# Record things
logger.store(LossQ=loss_q.item(), **loss_info)
# Possibly update pi and target networks
if timer % policy_delay == 0:
# Freeze Q-networks so you don't waste computational effort
# computing gradients for them during the policy learning step.
for p in q_params:
p.requires_grad = False
# Next run one gradient descent step for pi.
pi_optimizer.zero_grad()
loss_pi = compute_loss_pi(data)
loss_pi.backward()
pi_optimizer.step()
# Unfreeze Q-networks so you can optimize it at next DDPG step.
for p in q_params:
p.requires_grad = True
# Record things
logger.store(LossPi=loss_pi.item())
# Finally, update target networks by polyak averaging.
with torch.no_grad():
for p, p_targ in zip(ac.parameters(), ac_targ.parameters()):
# NB: We use an in-place operations "mul_", "add_" to update target
# params, as opposed to "mul" and "add", which would make new tensors.
p_targ.data.mul_(polyak)
p_targ.data.add_((1 - polyak) * p.data)
def get_action(o, noise_scale):
a = ac.act(torch.as_tensor(o, dtype=torch.float32).to(device))
a += noise_scale * np.random.randn(act_dim)
return np.clip(a, -act_limit, act_limit)
def test_agent():
for j in range(num_test_episodes):
o, d, ep_ret, ep_len = test_env.reset(), False, 0, 0
while not(d or (ep_len == max_ep_len)):
# Take deterministic actions at test time (noise_scale=0)
o, r, d, _ = test_env.step(get_action(o, 0))
ep_ret += r
ep_len += 1
logger.store(TestEpRet=ep_ret, TestEpLen=ep_len)
# Prepare for interaction with environment
total_steps = steps_per_epoch * epochs
start_time = time.time()
o, ep_ret, ep_len = env.reset(), 0, 0
# Main loop: collect experience in env and update/log each epoch
for t in range(total_steps):
# Until start_steps have elapsed, randomly sample actions
# from a uniform distribution for better exploration. Afterwards,
# use the learned policy (with some noise, via act_noise).
if t > start_steps:
a = get_action(o, act_noise)
else:
a = env.action_space.sample()
# Step the env
o2, r, d, _ = env.step(a)
ep_ret += r
ep_len += 1
# Ignore the "done" signal if it comes from hitting the time
# horizon (that is, when it's an artificial terminal signal
# that isn't based on the agent's state)
d = False if ep_len == max_ep_len else d
# Store experience to replay buffer
replay_buffer.store(o, a, r, o2, d)
# Super critical, easy to overlook step: make sure to update
# most recent observation!
o = o2
# End of trajectory handling
if d or (ep_len == max_ep_len):
logger.store(EpRet=ep_ret, EpLen=ep_len)
o, ep_ret, ep_len = env.reset(), 0, 0
# Update handling
if t >= update_after and t % update_every == 0:
for j in range(update_every):
batch = replay_buffer.sample_batch(batch_size)
update(data=batch, timer=j)
# End of epoch handling
if (t+1) % steps_per_epoch == 0:
epoch = (t+1) // steps_per_epoch
# Save model
if (epoch % save_freq == 0) or (epoch == epochs):
logger.save_state({'env': env}, None)
# Test the performance of the deterministic version of the agent.
test_agent()
# Log info about epoch
logger.log_tabular('Epoch', epoch)
logger.log_tabular('EpRet', with_min_and_max=True)
logger.log_tabular('TestEpRet', with_min_and_max=True)
logger.log_tabular('EpLen', average_only=True)
logger.log_tabular('TestEpLen', average_only=True)
logger.log_tabular('TotalEnvInteracts', t)
logger.log_tabular('Q1Vals', with_min_and_max=True)
logger.log_tabular('Q2Vals', with_min_and_max=True)
logger.log_tabular('LossPi', average_only=True)
logger.log_tabular('LossQ', average_only=True)
logger.log_tabular('Time', time.time()-start_time)
logger.dump_tabular()
