def mpi_avg_grads(module):
if num_procs() == 1:
return
for p in module.parameters():
p_grad_np = p.grad.numpy()
avg_p_grad = mpi_avg(p_grad_np)
p_grad_np[:] = avg_p_grad[:]
def update():
inputs = {k:v for k,v in zip(all_phs, buf.get())}
# inputs[a_ph] = np.tril(np.transpose(np.repeat(inputs[a_ph], n).reshape(trials, n, n), [0, 2, 1]))
# inputs[rew_ph] = np.tril(np.transpose(np.repeat(inputs[rew_ph], n).reshape(trials, n, n), [0, 2, 1]))
inputs[a_ph] = inputs[a_ph].reshape(trials, n)
inputs[rew_ph] = inputs[rew_ph].reshape(trials, n)
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))
def update():
inputs = {k: v for k, v in zip(all_phs, buf.get())}
inputs[pi_state_ph] = np.zeros((trials_per_epoch, NUM_GRU_UNITS))
inputs[v_state_ph] = np.zeros((trials_per_epoch, NUM_GRU_UNITS))
pi_l_old, v_l_old, ent = sess.run([pi_loss, v_loss, approx_ent],
feed_dict=inputs)
print(pi_l_old, v_l_old)
# Training
for i in range(train_pi_iters):
# print(f'pi:{i}')
_, kl = sess.run([train_pi, approx_kl], feed_dict=inputs)
# print(sess.run(pi_loss, 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):
# print(f'v:{_}')
sess.run(train_v, feed_dict=inputs)
# Log changes from update
import datetime
print(f'finish one batch training at {datetime.datetime.now()}')
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))
def mpi_avg_grads(module):
""" Average contents of gradient buffers across MPI processes. """
if num_procs() == 1:
return
for p in module.parameters():
p_grad_numpy = p.grad.numpy() # numpy view of tensor data
avg_p_grad = mpi_avg(p.grad)
p_grad_numpy[:] = avg_p_grad[:]
def update():
data = buf.get()
pi_l_old, pi_info_old = compute_loss_pi(data)
pi_l_old = pi_l_old.item()
v_l_old = compute_loss_v(data).item()
# Train policy with multiple steps of gradient descent
for i in range(train_pi_iters):
pi_optimizer.zero_grad()
loss_pi, pi_info = compute_loss_pi(data)
# Entropy bonus
loss_pi += pi_info['ent'] * 0.0021
kl = mpi_avg(pi_info['kl'])
if kl > 1.5 * target_kl:
logger.log(
'Early stopping at step %d due to reaching max kl.' % i)
break
loss_pi.backward()
mpi_avg_grads(ac.pi) # average grads across MPI processes
pi_optimizer.step()
logger.store(StopIter=i)
# Value function learning
for i in range(train_v_iters):
vf_optimizer.zero_grad()
loss_v = compute_loss_v(data)
loss_v.backward()
mpi_avg_grads(ac.v) # average grads across MPI processes
vf_optimizer.step()
# Log changes from update
kl, ent, cf = pi_info['kl'], pi_info_old['ent'], pi_info['cf']
logger.store(LossPi=pi_l_old,
LossV=v_l_old,
KL=kl,
Entropy=ent,
ClipFrac=cf,
DeltaLossPi=(loss_pi.item() - pi_l_old),
DeltaLossV=(loss_v.item() - v_l_old))
def ppo(env_fn,
actor_critic=core.MLPActorCritic,
ac_kwargs=dict(),
seed=0,
steps_per_epoch=4000,
epochs=50,
gamma=0.99,
clip_ratio=0.2,
pi_lr=3e-4,
vf_lr=1e-3,
train_pi_iters=80,
train_v_iters=80,
lam=0.97,
max_ep_len=None,
target_kl=0.01,
logger_kwargs=dict(),
save_freq=10,
TensorBoard=True,
save_nn=True,
save_every=1000,
load_latest=False,
load_custom=False,
LoadPath=None,
RTA_type=None):
"""
Proximal Policy Optimization (by clipping),
with early stopping based on approximate KL
Args:
env_fn : A function which creates a copy of the environment.
The environment must satisfy the OpenAI Gym API.
actor_critic: The constructor method for a PyTorch Module with a
``step`` method, an ``act`` method, a ``pi`` module, and a ``v``
module. The ``step`` method should accept a batch of observations
and return:
=========== ================ ======================================
Symbol Shape Description
=========== ================ ======================================
``a`` (batch, act_dim) | Numpy array of actions for each
| observation.
``v`` (batch,) | Numpy array of value estimates
| for the provided observations.
``logp_a`` (batch,) | Numpy array of log probs for the
| actions in ``a``.
=========== ================ ======================================
The ``act`` method behaves the same as ``step`` but only returns ``a``.
The ``pi`` module's forward call should accept a batch of
observations and optionally a batch of actions, and return:
=========== ================ ======================================
Symbol Shape Description
=========== ================ ======================================
``pi`` N/A | Torch Distribution object, containing
| a batch of distributions describing
| the policy for the provided observations.
``logp_a`` (batch,) | Optional (only returned if batch of
| actions is given). Tensor containing
| the log probability, according to
| the policy, of the provided actions.
| If actions not given, will contain
| ``None``.
=========== ================ ======================================
The ``v`` module's forward call should accept a batch of observations
and return:
=========== ================ ======================================
Symbol Shape Description
=========== ================ ======================================
``v`` (batch,) | Tensor containing the value estimates
| for the provided observations. (Critical:
| make sure to flatten this!)
=========== ================ ======================================
ac_kwargs (dict): Any kwargs appropriate for the ActorCritic object
you provided to PPO.
seed (int): Seed for random number generators.
steps_per_epoch (int): Number of steps of interaction (state-action pairs)
for the agent and the environment in each epoch.
epochs (int): Number of epochs of interaction (equivalent to
number of policy updates) to perform.
gamma (float): Discount factor. (Always between 0 and 1.)
clip_ratio (float): Hyperparameter for clipping in the policy objective.
Roughly: how far can the new policy go from the old policy while
still profiting (improving the objective function)? The new policy
can still go farther than the clip_ratio says, but it doesn't help
on the objective anymore. (Usually small, 0.1 to 0.3.) Typically
denoted by :math:`\epsilon`.
pi_lr (float): Learning rate for policy optimizer.
vf_lr (float): Learning rate for value function optimizer.
train_pi_iters (int): Maximum number of gradient descent steps to take
on policy loss per epoch. (Early stopping may cause optimizer
to take fewer than this.)
train_v_iters (int): Number of gradient descent steps to take on
value function per epoch.
lam (float): Lambda for GAE-Lambda. (Always between 0 and 1,
close to 1.)
max_ep_len (int): Maximum length of trajectory / episode / rollout.
target_kl (float): Roughly what KL divergence we think is appropriate
between new and old policies after an update. This will get used
for early stopping. (Usually small, 0.01 or 0.05.)
logger_kwargs (dict): Keyword args for EpochLogger.
save_freq (int): How often (in terms of gap between epochs) to save
the current policy and value function.
TensorBoard (bool): True plots to TensorBoard, False does not
save_nn (bool): True saves neural network data, False does not
save_every (int): How often to save neural network
load_latest (bool): Load last saved neural network data before training
load_custom (bool): Load custom neural network data file before training
LoadPath (str): Path for custom neural network data file
RTA_type (str): RTA framework, either 'CBF', 'SVL', 'ASIF', or
'SBSF'
"""
# Special function to avoid certain slowdowns from PyTorch + MPI combo.
setup_pytorch_for_mpi()
# Set up logger and save configuration
logger = EpochLogger(**logger_kwargs)
logger.save_config(locals())
# Instantiate environment
env = env_fn()
obs_dim = env.observation_space.shape
act_dim = env.action_space.shape
# Random seed for each cpu
seed += 1 * proc_id()
env.seed(seed)
# Create actor-critic module
ac = actor_critic(env.observation_space, env.action_space, **ac_kwargs)
# Load model if True
if load_latest:
models = glob.glob(f"{PATH}/models/PPO/*")
LoadPath = max(models, key=os.path.getctime)
ac.load_state_dict(torch.load(LoadPath))
elif load_custom:
ac.load_state_dict(torch.load(LoadPath))
# Sync params across processes
sync_params(ac)
# Count variables
var_counts = tuple(core.count_vars(module) for module in [ac.pi, ac.v])
logger.log('\nNumber of parameters: \t pi: %d, \t v: %d\n' % var_counts)
# Set up experience buffer
local_steps_per_epoch = int(steps_per_epoch / num_procs())
buf = PPOBuffer(obs_dim, act_dim, local_steps_per_epoch, gamma, lam)
# Set up function for computing PPO policy loss
def compute_loss_pi(data):
obs, act, adv, logp_old = data['obs'], data['act'], data['adv'], data[
'logp']
# Policy loss
pi, logp = ac.pi(obs, act)
ratio = torch.exp(logp - logp_old)
clip_adv = torch.clamp(ratio, 1 - clip_ratio, 1 + clip_ratio) * adv
loss_pi = -(torch.min(ratio * adv, clip_adv)).mean()
# Useful extra info
approx_kl = (logp_old - logp).mean().item()
ent = pi.entropy().mean().item()
clipped = ratio.gt(1 + clip_ratio) | ratio.lt(1 - clip_ratio)
clipfrac = torch.as_tensor(clipped, dtype=torch.float32).mean().item()
pi_info = dict(kl=approx_kl, ent=ent, cf=clipfrac)
return loss_pi, pi_info
# Set up function for computing value loss
def compute_loss_v(data):
obs, ret = data['obs'], data['ret']
return ((ac.v(obs) - ret)**2).mean()
# Set up optimizers for policy and value function
pi_optimizer = Adam(ac.pi.parameters(), lr=pi_lr)
vf_optimizer = Adam(ac.v.parameters(), lr=vf_lr)
# Set up model saving
logger.setup_pytorch_saver(ac)
def update():
data = buf.get()
pi_l_old, pi_info_old = compute_loss_pi(data)
pi_l_old = pi_l_old.item()
v_l_old = compute_loss_v(data).item()
# Train policy with multiple steps of gradient descent
for i in range(train_pi_iters):
pi_optimizer.zero_grad()
loss_pi, pi_info = compute_loss_pi(data)
kl = mpi_avg(pi_info['kl'])
if kl > 1.5 * target_kl:
logger.log(
'Early stopping at step %d due to reaching max kl.' % i)
break
loss_pi.backward()
mpi_avg_grads(ac.pi) # average grads across MPI processes
pi_optimizer.step()
logger.store(StopIter=i)
# Value function learning
for i in range(train_v_iters):
vf_optimizer.zero_grad()
loss_v = compute_loss_v(data)
loss_v.backward()
mpi_avg_grads(ac.v) # average grads across MPI processes
vf_optimizer.step()
# Log changes from update
kl, ent, cf = pi_info['kl'], pi_info_old['ent'], pi_info['cf']
logger.store(LossPi=pi_l_old,
LossV=v_l_old,
KL=kl,
Entropy=ent,
ClipFrac=cf,
DeltaLossPi=(loss_pi.item() - pi_l_old),
DeltaLossV=(loss_v.item() - v_l_old))
# Import RTA
if RTA_type == 'CBF':
from CBF_for_speed_limit import RTA
elif RTA_type == 'SVL':
from Simple_velocity_limit import RTA
elif RTA_type == 'ASIF':
from IASIF import RTA
elif RTA_type == 'SBSF':
from ISimplex import RTA
# Call RTA, define action conversion
if RTA_type != 'off':
env.RTA_reward = RTA_type
rta = RTA(env)
def RTA_act(obs, act):
act = np.clip(act, -env.force_magnitude, env.force_magnitude)
x0 = [obs[0], obs[1], 0, obs[2], obs[3], 0]
u_des = np.array([[act[0]], [act[1]], [0]])
u = rta.main(x0, u_des)
new_act = [u[0, 0], u[1, 0]]
if np.sqrt((act[0] - new_act[0])**2 +
(act[1] - new_act[1])**2) < 0.0001:
env.RTA_on = False
else:
env.RTA_on = True
return new_act
# Prepare for interaction with environment
start_time = time.time()
o, ep_ret, ep_len = env.reset(), 0, 0
total_episodes = 0
RTA_percent = 0
# Create TensorBoard file if True
if TensorBoard and proc_id() == 0:
if env_name == 'spacecraft-docking-continuous-v0' or env_name == 'spacecraft-docking-v0':
Name = f"{PATH}/runs/Spacecraft-docking-" + current_time
elif env_name == 'dubins-aircraft-v0' or env_name == 'dubins-aircraft-continuous-v0':
Name = f"{PATH}/runs/Dubins-aircraft-" + current_time
writer = SummaryWriter(Name)
# Main loop: collect experience in env and update/log each epoch
for epoch in range(epochs):
batch_ret = [] # Track episode returns
batch_len = [] # Track episode lengths
batch_RTA_percent = [] # Track precentage of time RTA is on
env.success = 0 # Track episode success rate
env.failure = 0 # Track episode failure rate
env.crash = 0 # Track episode crash rate
env.overtime = 0 # Track episode over max time/control rate
episodes = 0 # Track episodes
delta_v = [] # Track episode total delta v
for t in range(local_steps_per_epoch):
a, v, logp = ac.step(torch.as_tensor(o, dtype=torch.float32))
if RTA_type != 'off': # If RTA is on, get RTA action
RTA_a = RTA_act(o, a)
if env.RTA_on:
RTA_percent += 1
next_o, r, d, _ = env.step(RTA_a)
else: # If RTA is off, pass through desired action
next_o, r, d, _ = env.step(a)
if env_name == 'spacecraft-docking-continuous-v0' or env_name == 'spacecraft-docking-v0':
over_max_vel, _, _ = env.check_velocity(a[0], a[1])
if over_max_vel:
RTA_percent += 1
ep_ret += r
ep_len += 1
# save and log
buf.store(o, a, r, v, logp)
logger.store(VVals=v)
# Update obs (critical!)
o = next_o
timeout = ep_len == max_ep_len
terminal = d or timeout
epoch_ended = t == local_steps_per_epoch - 1
if terminal or epoch_ended:
if epoch_ended and not (terminal):
print('Warning: trajectory cut off by epoch at %d steps.' %
ep_len,
flush=True)
# if trajectory didn't reach terminal state, bootstrap value target
if timeout or epoch_ended:
_, v, _ = ac.step(torch.as_tensor(o, dtype=torch.float32))
else:
v = 0
buf.finish_path(v)
if terminal:
# only save EpRet / EpLen if trajectory finished
logger.store(EpRet=ep_ret, EpLen=ep_len)
batch_ret.append(ep_ret)
batch_len.append(ep_len)
episodes += 1
if env_name == 'spacecraft-docking-continuous-v0' or env_name == 'spacecraft-docking-v0':
delta_v.append(env.control_input / env.mass_deputy)
batch_RTA_percent.append(RTA_percent / ep_len * 100)
RTA_percent = 0
o, ep_ret, ep_len = env.reset(), 0, 0
total_episodes += episodes
# Track success, failure, crash, overtime rates
if episodes != 0:
success_rate = env.success / episodes
failure_rate = env.failure / episodes
crash_rate = env.crash / episodes
overtime_rate = env.overtime / episodes
else:
success_rate = 0
failure_rate = 0
crash_rate = 0
overtime_rate = 0
raise (
"No completed episodes logging will break [increase steps per epoch]"
)
# 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()
# Average data over all cpus
avg_batch_ret = mpi_avg(np.mean(batch_ret))
avg_batch_len = mpi_avg(np.mean(batch_len))
avg_success_rate = mpi_avg(success_rate)
avg_failure_rate = mpi_avg(failure_rate)
avg_crash_rate = mpi_avg(crash_rate)
avg_overtime_rate = mpi_avg(overtime_rate)
if env_name == 'spacecraft-docking-continuous-v0' or env_name == 'spacecraft-docking-v0':
avg_delta_v = mpi_avg(np.mean(delta_v))
avg_RTA_percent = mpi_avg(np.mean(batch_RTA_percent))
if proc_id() == 0: # Only on one cpu
# Plot to TensorBoard if True, only on one cpu
if TensorBoard:
writer.add_scalar('Return', avg_batch_ret, epoch)
writer.add_scalar('Episode-Length', avg_batch_len * env.tau,
epoch)
writer.add_scalar('Success-Rate', avg_success_rate * 100,
epoch)
writer.add_scalar('Failure-Rate', avg_failure_rate * 100,
epoch)
writer.add_scalar('Crash-Rate', avg_crash_rate * 100, epoch)
writer.add_scalar('Overtime-Rate', avg_overtime_rate * 100,
epoch)
if env_name == 'spacecraft-docking-continuous-v0' or env_name == 'spacecraft-docking-v0':
writer.add_scalar('Delta-V', avg_delta_v, epoch)
writer.add_scalar('RTA-on-percent', avg_RTA_percent, epoch)
# Save neural network if true, can change to desired location
if save_nn and epoch % save_every == 0 and epoch != 0:
if not os.path.isdir(f"{PATH}/models"):
os.mkdir(f"{PATH}/models")
if not os.path.isdir(f"{PATH}/models/PPO"):
os.mkdir(f"{PATH}/models/PPO")
if env_name == 'spacecraft-docking-continuous-v0' or env_name == 'spacecraft-docking-v0':
Name2 = f"{PATH}/models/PPO/Spacecraft-docking-" + current_time + f"-epoch{epoch}.dat"
elif env_name == 'dubins-aircraft-v0' or env_name == 'dubins-aircraft-continuous-v0':
Name2 = f"{PATH}/models/PPO/Dubins-aircraft-" + current_time + f"-epoch{epoch}.dat"
torch.save(ac.state_dict(), Name2)
# Average episodes per hour, episode per epoch
ep_hr = mpi_avg(total_episodes) * args.cpu / (time.time() -
start_time) * 3600
ep_Ep = mpi_avg(total_episodes) * args.cpu / (epoch + 1)
# Plot on one cpu
if proc_id() == 0:
# Save neural network
if save_nn:
if not os.path.isdir(f"{PATH}/models"):
os.mkdir(f"{PATH}/models")
if not os.path.isdir(f"{PATH}/models/PPO"):
os.mkdir(f"{PATH}/models/PPO")
if env_name == 'spacecraft-docking-continuous-v0' or env_name == 'spacecraft-docking-v0':
Name2 = f"{PATH}/models/PPO/Spacecraft-docking-" + current_time + "-final.dat"
elif env_name == 'dubins-aircraft-v0' or env_name == 'dubins-aircraft-continuous-v0':
Name2 = f"{PATH}/models/PPO/Dubins-aircraft-" + current_time + "-final.dat"
torch.save(ac.state_dict(), Name2)
# Print statistics on episodes
print(
f"Episodes per hour: {ep_hr:.0f}, Episodes per epoch: {ep_Ep:.0f}, Epochs per hour: {(epoch+1)/(time.time()-start_time)*3600:.0f}"
)