def load_policy(model_path,
input_dim,
output_dim,
num_hidden,
num_layers,
init_logstd=1.,
discrete=False,
beta=1.0):
observation_space = Box(low=-np.inf, high=np.inf, shape=(input_dim, ))
if discrete:
action_space = Discrete(n=output_dim)
else:
action_space = Box(low=-np.inf, high=np.inf, shape=(output_dim, ))
tf.reset_default_graph()
config = tf.ConfigProto(allow_soft_placement=True,
inter_op_parallelism_threads=8,
intra_op_parallelism_threads=8,
device_count={'CPU': 8})
config.gpu_options.allow_growth = True
sess = U.make_session(make_default=True, config=config)
network = mlp(num_hidden=num_hidden, num_layers=num_layers)
# policy_train = build_policy(observation_space, action_space, network, trainable_variance=True,
# state_dependent_variance=True, beta=beta, init_logstd=init_logstd)()
ob_space = env.observation_space
ac_space = env.action_space
ob = observation_placeholder(ob_space)
policy_train = build_policy_trpo(
env, network, value_network='copy')(observ_placeholder=ob)
U.initialize()
if model_path != '':
policy_train.load(model_path)
return policy_train
def runner(env,
policy_func,
load_model_path,
timesteps_per_batch,
number_trajs,
stochastic_policy,
save=False,
reuse=False):
# Setup network
# ----------------------------------------
ob_space = env.observation_space
ac_space = env.action_space
U.initialize()
policy = build_policy(env, 'mlp', value_network='copy')
ob = observation_placeholder(ob_space)
with tf.variable_scope('pi'):
pi = policy(observ_placeholder=ob)
saver = tf.train.Saver()
ckpt = tf.train.get_checkpoint_state(load_model_path)
saver.restore(U.get_session(), ckpt.model_checkpoint_path)
obs_list = []
acs_list = []
len_list = []
ret_list = []
from tqdm import tqdm
for _ in tqdm(range(number_trajs)):
traj = traj_1_generator(pi,
env,
timesteps_per_batch,
stochastic=stochastic_policy)
obs, acs, ep_len, ep_ret = traj['ob'], traj['ac'], traj[
'ep_len'], traj['ep_ret']
obs_list.append(obs)
acs_list.append(acs)
len_list.append(ep_len)
ret_list.append(ep_ret)
if stochastic_policy:
print('stochastic policy:')
else:
print('deterministic policy:')
if save:
filename = load_model_path.split('/')[-1] + '.' + env.spec.id
np.savez(filename,
obs=np.array(obs_list),
acs=np.array(acs_list),
lens=np.array(len_list),
rets=np.array(ret_list))
avg_len = sum(len_list) / len(len_list)
avg_ret = sum(ret_list) / len(ret_list)
# print("Average length:", avg_len)
# print("Average return:", avg_ret)
return avg_len, avg_ret
def runner(env,
policy_func,
load_model_path,
timesteps_per_batch,
number_trajs,
stochastic_policy,
save=False,
reuse=False):
# Setup network
# ----------------------------------------
ob_space = env.observation_space
ac_space = env.action_space
pi = policy_func("pi", ob_space, ac_space, reuse=reuse)
U.initialize()
# Prepare for rollouts
# ----------------------------------------
U.load_state(load_model_path)
obs_list = []
acs_list = []
len_list = []
ret_list = []
from tqdm import tqdm
for _ in tqdm(range(number_trajs)):
traj = traj_1_generator(pi,
env,
timesteps_per_batch,
stochastic=stochastic_policy)
obs, acs, ep_len, ep_ret = traj['ob'], traj['ac'], traj[
'ep_len'], traj['ep_ret']
obs_list.append(obs)
acs_list.append(acs)
len_list.append(ep_len)
ret_list.append(ep_ret)
if stochastic_policy:
print('stochastic policy:')
else:
print('deterministic policy:')
if save:
filename = load_model_path.split('/')[-1] + '.' + env.spec.id
np.savez(filename,
obs=np.array(obs_list),
acs=np.array(acs_list),
lens=np.array(len_list),
rets=np.array(ret_list))
avg_len = sum(len_list) / len(len_list)
avg_ret = sum(ret_list) / len(ret_list)
print("Average length:", avg_len)
print("Average return:", avg_ret)
return avg_len, avg_ret
def test_runningmeanstd():
for (x1, x2, x3) in [
(np.random.randn(3), np.random.randn(4), np.random.randn(5)),
(np.random.randn(3, 2), np.random.randn(4, 2), np.random.randn(5, 2)),
]:
rms = RunningMeanStd(epsilon=0.0, shape=x1.shape[1:])
U.initialize()
x = np.concatenate([x1, x2, x3], axis=0)
ms1 = [x.mean(axis=0), x.std(axis=0)]
rms.update(x1)
rms.update(x2)
rms.update(x3)
ms2 = [rms.mean.eval(), rms.std.eval()]
assert np.allclose(ms1, ms2)
def run_eval(config):
def create_loss():
test_dataset_args = create_dataset(config,
split="test",
shuffle=True,
repeat=True)
test_inputs, test_targets, test_lengths, _, _, _ = test_dataset_args
cell = create_cell(config, test_dataset_args)
test_ll_per_seq, kl, log_weight, log_ess, trajectories = \
basic_bounds.iwae(cell, (test_inputs, test_targets), test_lengths, num_samples=config.num_samples)
test_ll_per_t = tf.reduce_mean(test_ll_per_seq /
tf.to_float(test_lengths))
return test_ll_per_t, trajectories, cell, log_weight, test_lengths
def create_graph():
global_step = tf.train.get_or_create_global_step()
test_bound, trajectories, cell, log_weight, test_lengths = create_loss
return cell, test_bound, global_step, trajectories, test_lengths
cell, test_bound, global_step, trajectories, test_lengths = create_graph()
sess = U.get_session()
U.initialize()
cur_step = 0
# saver
saver = tf.train.Saver(max_to_keep=1)
if not os.path.exists(config.logdir):
assert False
ckpt = tf.train.get_checkpoint_state(config.logdir + '/valid_best')
if ckpt and ckpt.model_checkpoint_path:
saver.restore(sess, ckpt.model_checkpoint_path)
cur_step = int(ckpt.model_checkpoint_path.split('-')[-1])
print('Model and log loaded! (checkpoint_path=%s, cur_step=%d)' %
(ckpt.model_checkpoint_path, cur_step))
test_bound_value = sess.run([test_bound])
print("##################################")
print("VALID_BEST_STEP: %s" % cur_step)
print("PARTICLE_NUM: %s" % config.num_samples)
print("TEST_BOUND: %s" % test_bound_value[0])
print("##################################")
def evaluate(env,
bc_agent_wrapper,
num_trajs,
render,
exact_model_path=None,
model_ckpt_dir=None):
"""Evaluate a trained SAM agent"""
# Only one of the two arguments can be provided
assert sum([exact_model_path is None, model_ckpt_dir is None]) == 1
# Rebuild the computational graph
pol = bc_agent_wrapper('pol')
# Create episode generator
traj_gen = traj_ep_generator(env, pol, render)
# Initialize and load the previously learned weights into the freshly re-built graph
U.initialize()
if exact_model_path is not None:
U.load_model(exact_model_path)
logger.info(
"model loaded from exact path:\n {}".format(exact_model_path))
else: # `exact_model_path` is None -> `model_ckpt_dir` is not None
U.load_latest_checkpoint(model_ckpt_dir)
logger.info("model loaded from ckpt dir:\n {}".format(model_ckpt_dir))
# Initialize the history data structures
ep_lens = []
ep_env_rets = []
# Collect trajectories
for i in range(num_trajs):
logger.info("evaluating [{}/{}]".format(i + 1, num_trajs))
traj = traj_gen.__next__()
ep_len, ep_env_ret = traj['ep_len'], traj['ep_env_ret']
# Aggregate to the history data structures
ep_lens.append(ep_len)
ep_env_rets.append(ep_env_ret)
# Log some statistics of the collected trajectories
ep_len_mean = np.mean(ep_lens)
ep_env_ret_mean = np.mean(ep_env_rets)
logger.record_tabular("ep_len_mean", ep_len_mean)
logger.record_tabular("ep_env_ret_mean", ep_env_ret_mean)
logger.dump_tabular()
def test_dist():
np.random.seed(0)
p1, p2, p3 = (np.random.randn(3, 1), np.random.randn(4, 1),
np.random.randn(5, 1))
q1, q2, q3 = (np.random.randn(6, 1), np.random.randn(7, 1),
np.random.randn(8, 1))
# p1,p2,p3=(np.random.randn(3), np.random.randn(4), np.random.randn(5))
# q1,q2,q3=(np.random.randn(6), np.random.randn(7), np.random.randn(8))
comm = MPI.COMM_WORLD
assert comm.Get_size() == 2
if comm.Get_rank() == 0:
x1, x2, x3 = p1, p2, p3
elif comm.Get_rank() == 1:
x1, x2, x3 = q1, q2, q3
else:
assert False
rms = RunningMeanStd(epsilon=0.0, shape=(1, ))
U.initialize()
rms.update(x1)
rms.update(x2)
rms.update(x3)
bigvec = np.concatenate([p1, p2, p3, q1, q2, q3])
def checkallclose(x, y):
print(x, y)
return np.allclose(x, y)
assert checkallclose(
bigvec.mean(axis=0),
rms.mean.eval(),
)
assert checkallclose(
bigvec.std(axis=0),
rms.std.eval(),
)
def learn(env,
policy_func,
reward_giver,
reward_guidance,
expert_dataset,
rank,
pretrained,
pretrained_weight,
*,
g_step,
d_step,
entcoeff,
save_per_iter,
ckpt_dir,
log_dir,
timesteps_per_batch,
task_name,
gamma,
lam,
algo,
max_kl,
cg_iters,
cg_damping=1e-2,
vf_stepsize=3e-4,
d_stepsize=1e-4,
vf_iters=3,
max_timesteps=0,
max_episodes=0,
max_iters=0,
loss_percent=0.0,
callback=None):
nworkers = MPI.COMM_WORLD.Get_size()
rank = MPI.COMM_WORLD.Get_rank()
np.set_printoptions(precision=3)
# Setup losses and stuff
# ----------------------------------------
ob_space = env.observation_space
ac_space = env.action_space
policy = build_policy(env, 'mlp', value_network='copy')
ob = observation_placeholder(ob_space)
with tf.variable_scope('pi'):
pi = policy(observ_placeholder=ob)
with tf.variable_scope('oldpi'):
oldpi = policy(observ_placeholder=ob)
atarg = tf.placeholder(
dtype=tf.float32,
shape=[None]) # Target advantage function (if applicable)
ret = tf.placeholder(dtype=tf.float32, shape=[None]) # Empirical return
ac = pi.pdtype.sample_placeholder([None])
kloldnew = oldpi.pd.kl(pi.pd)
ent = pi.pd.entropy()
meankl = tf.reduce_mean(kloldnew)
meanent = tf.reduce_mean(ent)
entbonus = entcoeff * meanent
vferr = tf.reduce_mean(tf.square(pi.vf - ret))
ratio = tf.exp(pi.pd.logp(ac) -
oldpi.pd.logp(ac)) # advantage * pnew / pold
surrgain = tf.reduce_mean(ratio * atarg)
optimgain = surrgain + entbonus
losses = [optimgain, meankl, entbonus, surrgain, meanent]
loss_names = ["optimgain", "meankl", "entloss", "surrgain", "entropy"]
dist = meankl
all_var_list = get_trainable_variables('pi')
# var_list = [v for v in all_var_list if v.name.startswith("pi/pol") or v.name.startswith("pi/logstd")]
# vf_var_list = [v for v in all_var_list if v.name.startswith("pi/vff")]
var_list = get_pi_trainable_variables("pi")
vf_var_list = get_vf_trainable_variables("pi")
# assert len(var_list) == len(vf_var_list) + 1
d_adam = MpiAdam(reward_giver.get_trainable_variables())
guidance_adam = MpiAdam(reward_guidance.get_trainable_variables())
vfadam = MpiAdam(vf_var_list)
get_flat = U.GetFlat(var_list)
set_from_flat = U.SetFromFlat(var_list)
klgrads = tf.gradients(dist, var_list)
flat_tangent = tf.placeholder(dtype=tf.float32,
shape=[None],
name="flat_tan")
shapes = [var.get_shape().as_list() for var in var_list]
start = 0
tangents = []
for shape in shapes:
sz = U.intprod(shape)
tangents.append(tf.reshape(flat_tangent[start:start + sz], shape))
start += sz
gvp = tf.add_n([
tf.reduce_sum(g * tangent)
for (g, tangent) in zipsame(klgrads, tangents)
]) # pylint: disable=E1111
fvp = U.flatgrad(gvp, var_list)
assign_old_eq_new = U.function(
[], [],
updates=[
tf.assign(oldv, newv)
for (oldv,
newv) in zipsame(get_variables('oldpi'), get_variables('pi'))
])
compute_losses = U.function([ob, ac, atarg], losses)
compute_lossandgrad = U.function([ob, ac, atarg], losses +
[U.flatgrad(optimgain, var_list)])
compute_fvp = U.function([flat_tangent, ob, ac, atarg], fvp)
compute_vflossandgrad = U.function([ob, ret],
U.flatgrad(vferr, vf_var_list))
@contextmanager
def timed(msg):
if rank == 0:
print(colorize(msg, color='magenta'))
tstart = time.time()
yield
print(
colorize("done in %.3f seconds" % (time.time() - tstart),
color='magenta'))
else:
yield
def allmean(x):
assert isinstance(x, np.ndarray)
out = np.empty_like(x)
MPI.COMM_WORLD.Allreduce(x, out, op=MPI.SUM)
out /= nworkers
return out
U.initialize()
th_init = get_flat()
MPI.COMM_WORLD.Bcast(th_init, root=0)
set_from_flat(th_init)
d_adam.sync()
guidance_adam.sync()
vfadam.sync()
if rank == 0:
print("Init param sum", th_init.sum(), flush=True)
# Prepare for rollouts
# ----------------------------------------
seg_gen = traj_segment_generator(pi,
env,
reward_giver,
reward_guidance,
timesteps_per_batch,
stochastic=True,
algo=algo,
loss_percent=loss_percent)
episodes_so_far = 0
timesteps_so_far = 0
iters_so_far = 0
tstart = time.time()
lenbuffer = deque(maxlen=40) # rolling buffer for episode lengths
rewbuffer = deque(maxlen=40) # rolling buffer for episode rewards
true_rewbuffer = deque(maxlen=40)
assert sum([max_iters > 0, max_timesteps > 0, max_episodes > 0]) == 1
g_loss_stats = stats(loss_names)
d_loss_stats = stats(reward_giver.loss_name)
ep_stats = stats(["True_rewards", "Rewards", "Episode_length"])
# if provide pretrained weight
if pretrained_weight is not None:
U.load_state(pretrained_weight, var_list=pi.get_variables())
while True:
if callback: callback(locals(), globals())
if max_timesteps and timesteps_so_far >= max_timesteps:
break
elif max_episodes and episodes_so_far >= max_episodes:
break
elif max_iters and iters_so_far >= max_iters:
break
# Save model
if rank == 0 and iters_so_far % save_per_iter == 0 and ckpt_dir is not None:
fname = os.path.join(ckpt_dir, task_name)
os.makedirs(os.path.dirname(fname), exist_ok=True)
saver = tf.train.Saver()
saver.save(tf.get_default_session(), fname)
logger.log("********** Iteration %i ************" % iters_so_far)
# global flag_render
# if iters_so_far > 0 and iters_so_far % 10 ==0:
# flag_render = True
# else:
# flag_render = False
def fisher_vector_product(p):
return allmean(compute_fvp(p, *fvpargs)) + cg_damping * p
# ------------------ Update G ------------------
logger.log("Optimizing Policy...")
for _ in range(g_step):
with timed("sampling"):
seg = seg_gen.__next__()
print('rewards', seg['rew'])
add_vtarg_and_adv(seg, gamma, lam)
# ob, ac, atarg, ret, td1ret = map(np.concatenate, (obs, acs, atargs, rets, td1rets))
ob, ac, atarg, tdlamret = seg["ob"], seg["ac"], seg["adv"], seg[
"tdlamret"]
vpredbefore = seg[
"vpred"] # predicted value function before udpate
atarg = (atarg - atarg.mean()) / atarg.std(
) # standardized advantage function estimate
if hasattr(pi, "ob_rms"):
pi.ob_rms.update(ob) # update running mean/std for policy
args = seg["ob"], seg["ac"], atarg
fvpargs = [arr[::5] for arr in args]
assign_old_eq_new(
) # set old parameter values to new parameter values
with timed("computegrad"):
*lossbefore, g = compute_lossandgrad(*args)
lossbefore = allmean(np.array(lossbefore))
g = allmean(g)
if np.allclose(g, 0):
logger.log("Got zero gradient. not updating")
else:
with timed("cg"):
stepdir = cg(fisher_vector_product,
g,
cg_iters=cg_iters,
verbose=rank == 0)
assert np.isfinite(stepdir).all()
shs = .5 * stepdir.dot(fisher_vector_product(stepdir))
lm = np.sqrt(shs / max_kl)
# logger.log("lagrange multiplier:", lm, "gnorm:", np.linalg.norm(g))
fullstep = stepdir / lm
expectedimprove = g.dot(fullstep)
surrbefore = lossbefore[0]
stepsize = 1.0
thbefore = get_flat()
for _ in range(10):
thnew = thbefore + fullstep * stepsize
set_from_flat(thnew)
meanlosses = surr, kl, *_ = allmean(
np.array(compute_losses(*args)))
improve = surr - surrbefore
logger.log("Expected: %.3f Actual: %.3f" %
(expectedimprove, improve))
if not np.isfinite(meanlosses).all():
logger.log("Got non-finite value of losses -- bad!")
elif kl > max_kl * 1.5:
logger.log("violated KL constraint. shrinking step.")
elif improve < 0:
logger.log("surrogate didn't improve. shrinking step.")
else:
logger.log("Stepsize OK!")
break
stepsize *= .5
else:
logger.log("couldn't compute a good step")
set_from_flat(thbefore)
if nworkers > 1 and iters_so_far % 20 == 0:
paramsums = MPI.COMM_WORLD.allgather(
(thnew.sum(),
vfadam.getflat().sum())) # list of tuples
assert all(
np.allclose(ps, paramsums[0]) for ps in paramsums[1:])
with timed("vf"):
for _ in range(vf_iters):
for (mbob, mbret) in dataset.iterbatches(
(seg["ob"], seg["tdlamret"]),
include_final_partial_batch=False,
batch_size=128):
if hasattr(pi, "ob_rms"):
pi.ob_rms.update(
mbob) # update running mean/std for policy
g = allmean(compute_vflossandgrad(mbob, mbret))
vfadam.update(g, vf_stepsize)
g_losses = meanlosses
for (lossname, lossval) in zip(loss_names, meanlosses):
logger.record_tabular(lossname, lossval)
logger.record_tabular("ev_tdlam_before",
explained_variance(vpredbefore, tdlamret))
# ------------------ Update D ------------------
logger.log("Optimizing Discriminator...")
logger.log(fmt_row(13, reward_giver.loss_name))
ob_expert, ac_expert = expert_dataset.get_next_batch(
batch_size=len(ob))
batch_size = 128
d_losses = [
] # list of tuples, each of which gives the loss for a minibatch
with timed("Discriminator"):
for (ob_batch, ac_batch) in dataset.iterbatches(
(ob, ac),
include_final_partial_batch=False,
batch_size=batch_size):
ob_expert, ac_expert = expert_dataset.get_next_batch(
batch_size=batch_size)
# update running mean/std for reward_giver
if hasattr(reward_giver, "obs_rms"):
reward_giver.obs_rms.update(
np.concatenate((ob_batch, ob_expert), 0))
*newlosses, g = reward_giver.lossandgrad(ob_batch, ob_expert)
d_adam.update(allmean(g), d_stepsize)
d_losses.append(newlosses)
logger.log(fmt_row(13, np.mean(d_losses, axis=0)))
# ------------------ Update Guidance ------------
logger.log("Optimizing Guidance...")
logger.log(fmt_row(13, reward_guidance.loss_name))
batch_size = 128
guidance_losses = [
] # list of tuples, each of which gives the loss for a minibatch
with timed("Guidance"):
for ob_batch, ac_batch in dataset.iterbatches(
(ob, ac),
include_final_partial_batch=False,
batch_size=batch_size):
ob_expert, ac_expert = expert_dataset.get_next_batch(
batch_size=batch_size)
idx_condition = process_expert(ob_expert, ac_expert)
pick_idx = (idx_condition >= loss_percent)
# pick_idx = idx_condition
ob_expert_p = ob_expert[pick_idx]
ac_expert_p = ac_expert[pick_idx]
ac_batch_p = []
for each_ob in ob_expert_p:
tmp_ac, _, _, _ = pi.step(each_ob, stochastic=True)
ac_batch_p.append(tmp_ac)
# update running mean/std for reward_giver
if hasattr(reward_guidance, "obs_rms"):
reward_guidance.obs_rms.update(ob_expert_p)
# reward_guidance.train(expert_s=ob_batch_p, agent_a=ac_batch_p, expert_a=ac_expert_p)
*newlosses, g = reward_guidance.lossandgrad(
ob_expert_p, ac_batch_p, ac_expert_p)
guidance_adam.update(allmean(g), d_stepsize)
guidance_losses.append(newlosses)
logger.log(fmt_row(13, np.mean(guidance_losses, axis=0)))
lrlocal = (seg["ep_lens"], seg["ep_rets"], seg["ep_true_rets"]
) # local values
listoflrpairs = MPI.COMM_WORLD.allgather(lrlocal) # list of tuples
lens, rews, true_rets = map(flatten_lists, zip(*listoflrpairs))
true_rewbuffer.extend(true_rets)
lenbuffer.extend(lens)
rewbuffer.extend(rews)
logger.record_tabular("EpLenMean", np.mean(lenbuffer))
logger.record_tabular("EpRewMean", np.mean(rewbuffer))
logger.record_tabular("EpTrueRewMean", np.mean(true_rewbuffer))
logger.record_tabular("EpThisIter", len(lens))
episodes_so_far += len(lens)
timesteps_so_far += sum(lens) * g_step
iters_so_far += 1
logger.record_tabular("EpisodesSoFar", episodes_so_far)
logger.record_tabular("TimestepsSoFar", timesteps_so_far)
logger.record_tabular("TimeElapsed", time.time() - tstart)
if rank == 0:
logger.dump_tabular()
def main(_):
# create visualizer
#visualizer = TensorboardVisualizer()
monitor = Monitor(FLAGS)
#log_dir = monitor.log_dir
#visualizer.initialize(log_dir, None)
saved_mean_reward = None
# openAI logger
L.configure(monitor.log_dir, format_strs=['stdout', 'csv'])
# initialize env
atari_env = AtariEnv(monitor)
#screen_shot_subgoal(atari_env)
# we should probably follow deepmind style env
# stack 4 frames and scale float
env = wrapper.wrap_deepmind(atari_env, frame_stack=True, scale=True)
# get default tf_session
sess = U.get_session()
# create q networks for controller
controller_optimizer = tf.train.AdamOptimizer(learning_rate=LEARNING_RATE)
controller_network = Q_network(env.observation_space, env.action_space.n, controller_optimizer, scope='controller')
controller = Controller(controller_network, env.action_space.n)
# create q networks for meta-controller
num_goals = env.unwrapped.goals_space.n
metacontroller_optimizer = tf.train.AdamOptimizer(learning_rate=LEARNING_RATE)
metacontroller_network = Q_network(env.observation_space, num_goals, metacontroller_optimizer, scope='metacontroller')
metacontroller = MetaController(metacontroller_network, num_goals)
# Create the schedule for exploration starting from 1.
exploration2 = LinearSchedule(schedule_timesteps=int(EXPLORATION_FRACTION * monitor.num_timesteps),
initial_p=1.0,
final_p=EXPLORATION_FINAL_EPS)
# initialize experience replay
controller_replay_buffer = ReplayBuffer(D1_MEMORY_SIZE)
metacontroller_replay_buffer = ReplayBuffer(D2_MEMORY_SIZE)
# initialize critic
critic = Critic(env.unwrapped)
total_extrinsic_reward = []
# for success rate
total_goal_reached = np.zeros(num_goals, dtype=np.int32)
total_goal_sampled = np.zeros(num_goals, dtype=np.int32)
total_goal_epsilon = np.ones(num_goals, dtype=np.float32)
ep = 0
total_step = 0
init_ob = env.reset()
U.initialize()
# initialize target network in both controller and meta
sess.run(metacontroller.network.update_target_op)
sess.run(controller.network.update_target_op)
# load ckpt if presence
model_path = tf.train.latest_checkpoint(monitor.ckpt_dir)
model_saved = False
model_file = os.path.join(monitor.ckpt_dir, 'model')
if model_path is not None:
U.load_variables(model_file)
L.log('loaded model from %s' % model_file)
model_saved = True
while ep < MAX_EPISODE: # count number of steps
# init environment game play variables
init_ob = env.reset()
observation = np.reshape(init_ob['observation'], (1, )+init_ob['observation'].shape)
desired_goal = metacontroller.sample_act(sess, observation, update_eps=1.0)[0]
env.unwrapped.desired_goal = desired_goal
total_goal_sampled[desired_goal] += 1
# given predicted goal, we encode this goal bounding mask to the observation np array
ob_with_g = env.unwrapped._add_goal_mask(init_ob['observation'], desired_goal)
# NOTE: Below code verify added mask correctly
# for i in range(ob_with_g.shape[-1]):
# ob = ob_with_g[:,:,i]
# image = Image.fromarray(ob)
# image = image.convert('RGB')
# image.save('test_%i.png' % i)
done = False
reached_goal = False
while not done:
extrinsic_rewards = 0
s0 = init_ob['observation']
while not (done or reached_goal):
update_eps1_with_respect_to_g = get_epsilon(total_goal_epsilon, total_goal_reached, total_goal_sampled, desired_goal, total_step, EXPLORATION_WARM_UP)
ob_with_g_reshaped = np.reshape(ob_with_g, (1, )+ob_with_g.shape)
primitive_action_t = controller.sample_act(sess, ob_with_g_reshaped, update_eps=update_eps1_with_respect_to_g)[0]
# obtain extrinsic reward from environment
ob_tp1, extrinsic_reward_t, done_t, info = env.step(primitive_action_t)
reached_goal = env.unwrapped.reached_goal(desired_goal)
ob_with_g_tp1 = env.unwrapped._add_goal_mask(ob_tp1['observation'], desired_goal)
intrinsic_reward_t = critic.criticize(desired_goal, reached_goal, primitive_action_t, done_t)
controller_replay_buffer.add(ob_with_g, primitive_action_t, intrinsic_reward_t, ob_with_g_tp1, done_t)
# sample from replay_buffer1 to train controller
obs_with_g_t, primitive_actions_t, intrinsic_rewards_t, obs_with_g_tp1, dones_t = controller_replay_buffer.sample(TRAIN_BATCH_SIZE)
weights, batch_idxes = np.ones_like(intrinsic_rewards_t), None
# get q estimate for tp1 as 'supervised'
ob_with_g_tp1_reshaped = np.reshape(ob_with_g_tp1, (1, )+ob_with_g.shape)
q_tp1 = controller.get_q(sess, ob_with_g_tp1_reshaped)[0]
td_error = controller.train(sess, obs_with_g_t, primitive_actions_t, intrinsic_rewards_t, obs_with_g_tp1, dones_t, weights, q_tp1)
# join train meta-controller only sample from replay_buffer2 to train meta-controller
if total_step >= WARMUP_STEPS:
L.log('join train has started ----- step %d', total_step)
# sample from replay_buffer2 to train meta-controller
init_obs, goals_t, extrinsic_rewards_t, obs_terminate_in_g, dones_t = metacontroller_replay_buffer.sample(TRAIN_BATCH_SIZE)
weights, batch_idxes = np.ones_like(extrinsic_rewards_t), None
# get q estimate for tp1 as 'supervised'
obs_terminate_in_g_reshaped = np.reshape(obs_terminate_in_g, (1, )+obs_terminate_in_g.shape)
q_tp1 = metacontroller.get_q(sess, obs_terminate_in_g_reshaped)[0]
td_error = metacontroller.train(sess, init_obs, goals_t, extrinsic_rewards_t, obs_terminate_in_g, dones_t, weights, q_tp1)
if total_step % UPDATE_TARGET_NETWORK_FREQ == 0:
#L.log('UPDATE BOTH CONTROLLER Q NETWORKS ----- step %d', step)
sess.run(controller.network.update_target_op)
# its fine, we aren't really training meta dqn until after certain steps.
sess.run(metacontroller.network.update_target_op)
extrinsic_rewards += extrinsic_reward_t
ob_with_g = ob_with_g_tp1
done = done_t
total_step += 1
# we are done / reached_goal
# store transitions of init_ob, goal, all the extrinsic rewards, current ob in D2
# print("ep %d : step %d, goal extrinsic total %d" % (ep, step, extrinsic_rewards))
# clean observation without goal encoded
metacontroller_replay_buffer.add(init_ob['observation'], desired_goal, extrinsic_rewards, ob_tp1['observation'], done)
# if we are here then we have finished the desired goal
if not done:
#print("ep %d : goal %d reached, not yet done, extrinsic %d" % (ep, desired_goal, extrinsic_rewards))
exploration_ep = 1.0
total_goal_reached[env.unwrapped.achieved_goal] += 1
if total_step >= WARMUP_STEPS:
t = total_step - WARMUP_STEPS
exploration_ep = exploration2.value(t)
ob_with_g_reshaped = np.reshape(ob_with_g, (1, )+ob_with_g.shape)
while env.unwrapped.achieved_goal == desired_goal:
desired_goal = metacontroller.sample_act(sess, ob_with_g_reshaped, update_eps=exploration_ep)[0]
env.unwrapped.desired_goal = desired_goal
total_goal_sampled[desired_goal] += 1
L.log('ep %d : achieved goal was %d ----- new goal --- %d' % (ep, env.unwrapped.achieved_goal, desired_goal))
# start again
reached_goal = False
# finish an episode
total_extrinsic_reward.append(extrinsic_rewards)
ep += 1
mean_100ep_reward = round(np.mean(total_extrinsic_reward[-101:-1]), 1)
if ep % monitor.print_freq == 0 :
L.record_tabular("steps", total_step)
L.record_tabular("episodes", ep)
L.record_tabular("mean 100 episode reward", mean_100ep_reward)
L.dump_tabular()
if total_step % monitor.ckpt_freq == 0:
if saved_mean_reward is None or mean_100ep_reward > saved_mean_reward:
L.log("Saving model due to mean reward increase: {} -> {}".format(
saved_mean_reward, mean_100ep_reward))
U.save_variables(model_file)
model_saved = True
saved_mean_reward = mean_100ep_reward
# verified our model was saved
if model_saved:
L.log('restored model with mean reward: %d' % saved_mean_reward)
U.load_variables(model_file)
def learn(comm, env, bc_agent_wrapper, experiment_name, ckpt_dir, summary_dir,
expert_dataset, lr, batch_size, max_iters):
rank = comm.Get_rank()
# Create the BC agent
pol = bc_agent_wrapper('pol')
# Create mpi adam optimizer for the policy
pol_optimizer = MpiAdamOptimizer(comm,
clip_norm=pol.hps.clip_norm,
learning_rate=lr,
name='pol_adam')
_optimize_pol = pol_optimizer.minimize(pol.loss,
var_list=pol.trainable_vars)
# Retrieve already-existing placeholders
e_obs = U.get_placeholder_cached(name='e_obs')
e_acs = U.get_placeholder_cached(name='e_acs')
# Create Theano-like ops
optimize_pol = U.function([e_obs, e_acs], _optimize_pol)
# Initialize variables
U.initialize()
# Sync params of all processes with the params of the root process
pol_optimizer.sync_from_root(pol.trainable_vars)
if rank == 0:
# Create summary writer
writer = U.file_writer(summary_dir)
# Create the summary
_names = ['train_loss', 'val_loss']
_summary = CustomSummary(scalar_keys=_names, family="bc")
# Define the origin of time
tstart = time.time()
# Define rolling buffers for loss collection
maxlen = 100
pol_train_loss_buffer = deque(maxlen=maxlen)
pol_val_loss_buffer = deque(maxlen=maxlen)
for iters_so_far in range(max_iters):
# Verify that the processes are still in sync
if iters_so_far > 0 and iters_so_far % 10 == 0:
pol_optimizer.check_synced(pol.trainable_vars)
# Save the model
if rank == 0 and iters_so_far % int(1e4) == 0 and ckpt_dir is not None:
model_path = osp.join(ckpt_dir, experiment_name)
U.save_state(model_path, iters_so_far=iters_so_far)
logger.info("saving model")
logger.info(" @: {}".format(model_path))
# Make non-zero-rank workers wait for rank zero
comm.Barrier()
# Go through mini-batches of the demonstration dataset, training fraction
obs, acs = expert_dataset.get_next_pair_batch(batch_size, 'train')
# Update running mean and std on states
if hasattr(pol, "obs_rms"):
pol.obs_rms.update(obs, comm)
# Perform a gradient step to update the policy parameters
optimize_pol(obs, acs)
# Compute training loss
pol_train_loss = pol.compute_pol_loss(obs, acs)
pol_train_loss_buffer.append(pol_train_loss)
# Go through mini-batches of the demonstration dataset, validation fraction
obs, acs = expert_dataset.get_next_pair_batch(-1, 'val')
# Compute validation loss
pol_val_loss = pol.compute_pol_loss(obs, acs)
pol_val_loss_buffer.append(pol_val_loss)
if iters_so_far % 100 == 0:
# Log training and validation losses
logger.info(
('iter #{} '
'| train loss: {} '
'| val loss: {} '
'| elapsed: {}').format(iters_so_far, pol_train_loss,
pol_val_loss,
prettify_time(time.time() - tstart)))
# Prepare losses to be dumped in summaries
all_summaries = [
np.mean(pol_train_loss_buffer),
np.mean(pol_val_loss_buffer)
] # must be visible by all workers
if rank == 0:
assert len(_names) == len(
all_summaries), "mismatch in list lengths"
_summary.add_all_summaries(writer, all_summaries, iters_so_far)
def learn(
*,
network,
env,
eval_policy,
total_timesteps,
timesteps_per_batch=1024, # what to train on
max_kl=0.001,
cg_iters=10,
gamma=0.99,
lam=1.0, # advantage estimation
seed=None,
ent_coef=0.0,
cg_damping=1e-2,
vf_stepsize=3e-4,
vf_iters=3,
max_episodes=0,
max_iters=0, # time constraint
callback=None,
load_path=None,
checkpoint_path_in=None,
checkpoint_dir_out=None,
checkpoint_freq=100, # In iterations!!,
from_iter=0,
eval_episodes=20,
**network_kwargs):
'''
learn a policy function with TRPO algorithm
Parameters:
----------
network neural network to learn. Can be either string ('mlp', 'cnn', 'lstm', 'lnlstm' for basic types)
or function that takes input placeholder and returns tuple (output, None) for feedforward nets
or (output, (state_placeholder, state_output, mask_placeholder)) for recurrent nets
env environment (one of the gym environments or wrapped via baselines.common.vec_env.VecEnv-type class
timesteps_per_batch timesteps per gradient estimation batch
max_kl max KL divergence between old policy and new policy ( KL(pi_old || pi) )
ent_coef coefficient of policy entropy term in the optimization objective
cg_iters number of iterations of conjugate gradient algorithm
cg_damping conjugate gradient damping
vf_stepsize learning rate for adam optimizer used to optimie value function loss
vf_iters number of iterations of value function optimization iterations per each policy optimization step
total_timesteps max number of timesteps
max_episodes max number of episodes
max_iters maximum number of policy optimization iterations
callback function to be called with (locals(), globals()) each policy optimization step
load_path str, path to load the model from (default: None, i.e. no model is loaded)
**network_kwargs keyword arguments to the policy / network builder. See baselines.common/policies.py/build_policy and arguments to a particular type of network
Returns:
-------
learnt model
'''
nworkers = MPI.COMM_WORLD.Get_size()
rank = MPI.COMM_WORLD.Get_rank()
cpus_per_worker = 1
U.get_session(
config=tf.ConfigProto(allow_soft_placement=True,
inter_op_parallelism_threads=cpus_per_worker,
intra_op_parallelism_threads=cpus_per_worker))
policy = build_policy(env, network, value_network='copy', **network_kwargs)
set_global_seeds(seed)
np.set_printoptions(precision=3)
# Setup losses and stuff
# ----------------------------------------
# ob_space = Box(low=-np.inf, high=np.inf, shape=(env.observation_space.n,))
ob_space = env.observation_space
ac_space = env.action_space
ob = observation_placeholder(ob_space)
with tf.variable_scope("pi"):
pi = policy(observ_placeholder=ob)
with tf.variable_scope("oldpi"):
oldpi = policy(observ_placeholder=ob)
# Loading checkpoint
if checkpoint_path_in is not None and os.path.isfile(checkpoint_path_in):
pi.load(checkpoint_path_in)
logger.log('Loaded policy weights from %s' % checkpoint_path_in)
atarg = tf.placeholder(
dtype=tf.float32,
shape=[None]) # Target advantage function (if applicable)
ret = tf.placeholder(dtype=tf.float32, shape=[None]) # Empirical return
ac = pi.pdtype.sample_placeholder([None])
kloldnew = oldpi.pd.kl(pi.pd)
ent = pi.pd.entropy()
meankl = tf.reduce_mean(kloldnew)
meanent = tf.reduce_mean(ent)
entbonus = ent_coef * meanent
vferr = tf.reduce_mean(tf.square(pi.vf - ret))
ratio = tf.exp(pi.pd.logp(ac) -
oldpi.pd.logp(ac)) # advantage * pnew / pold
surrgain = tf.reduce_mean(ratio * atarg)
optimgain = surrgain + entbonus
losses = [optimgain, meankl, entbonus, surrgain, meanent]
loss_names = ["optimgain", "meankl", "entloss", "surrgain", "entropy"]
dist = meankl
all_var_list = get_trainable_variables("pi")
# var_list = [v for v in all_var_list if v.name.split("/")[1].startswith("pol")]
# vf_var_list = [v for v in all_var_list if v.name.split("/")[1].startswith("vf")]
var_list = get_pi_trainable_variables("pi")
vf_var_list = get_vf_trainable_variables("pi")
vfadam = MpiAdam(vf_var_list)
get_flat = U.GetFlat(var_list)
set_from_flat = U.SetFromFlat(var_list)
klgrads = tf.gradients(dist, var_list)
flat_tangent = tf.placeholder(dtype=tf.float32,
shape=[None],
name="flat_tan")
shapes = [var.get_shape().as_list() for var in var_list]
start = 0
tangents = []
for shape in shapes:
sz = U.intprod(shape)
tangents.append(tf.reshape(flat_tangent[start:start + sz], shape))
start += sz
gvp = tf.add_n([
tf.reduce_sum(g * tangent)
for (g, tangent) in zipsame(klgrads, tangents)
]) # pylint: disable=E1111
fvp = U.flatgrad(gvp, var_list)
assign_old_eq_new = U.function(
[], [],
updates=[
tf.assign(oldv, newv)
for (oldv,
newv) in zipsame(get_variables("oldpi"), get_variables("pi"))
])
compute_losses = U.function([ob, ac, atarg], losses)
compute_lossandgrad = U.function([ob, ac, atarg], losses +
[U.flatgrad(optimgain, var_list)])
compute_fvp = U.function([flat_tangent, ob, ac, atarg], fvp)
compute_vflossandgrad = U.function([ob, ret],
U.flatgrad(vferr, vf_var_list))
@contextmanager
def timed(msg):
if rank == 0:
print(colorize(msg, color='magenta'))
tstart = time.time()
yield
print(
colorize("done in %.3f seconds" % (time.time() - tstart),
color='magenta'))
else:
yield
def allmean(x):
assert isinstance(x, np.ndarray)
out = np.empty_like(x)
MPI.COMM_WORLD.Allreduce(x, out, op=MPI.SUM)
out /= nworkers
return out
U.initialize()
# s = env.reset()
# start = time.time()
# for i in range(10000):
# pi.step(s, stochastic=True)
# duration = time.time() - start
# print(duration)
# return
if load_path is not None:
pi.load(load_path)
th_init = get_flat()
MPI.COMM_WORLD.Bcast(th_init, root=0)
set_from_flat(th_init)
vfadam.sync()
print("Init param sum", th_init.sum(), flush=True)
# Prepare for rollouts
# ----------------------------------------
seg_gen = traj_segment_generator(pi,
env,
timesteps_per_batch,
stochastic=True,
gamma=gamma)
episodes_so_far = 0
timesteps_so_far = 0
iters_so_far = 0
iters_eval = 0
all_logs = []
best_rew = -np.inf
tstart = time.time()
lenbuffer = deque(maxlen=40) # rolling buffer for episode lengths
rewbuffer = deque(maxlen=40) # rolling buffer for episode rewards
online_scores = []
offline_scores = []
if sum([max_iters > 0, total_timesteps > 0, max_episodes > 0]) == 0:
# noththing to be done
return pi
assert sum([max_iters > 0, total_timesteps > 0, max_episodes > 0]) < 2, \
'out of max_iters, total_timesteps, and max_episodes only one should be specified'
while True:
if callback: callback(locals(), globals())
if total_timesteps and timesteps_so_far >= total_timesteps:
break
elif max_episodes and episodes_so_far >= max_episodes:
break
elif max_iters and iters_so_far >= max_iters:
break
logger.log("********** Iteration %i ************" % iters_so_far)
if iters_so_far % checkpoint_freq == 0 and checkpoint_dir_out is not None:
if not os.path.exists(checkpoint_dir_out):
os.makedirs(checkpoint_dir_out)
pi.save(
os.path.join(checkpoint_dir_out,
'checkpoint_%d' % iters_so_far))
logger.log('Saved policy weights as %s' % os.path.join(
checkpoint_dir_out, 'checkpoint_%d.npy' % iters_so_far))
def pi_wrapper(ob):
ac, vpred, _, _ = pi.step(ob, stochastic=True)
return ac
rew, _, logs, disc_rets, num_stops, avg_damages = eval_policy(
pi=pi_wrapper, n_episodes=eval_episodes, verbose=True)
offline_scores.append(
[np.mean(disc_rets),
np.mean(num_stops),
np.mean(avg_damages)])
np.save(os.path.join(checkpoint_dir_out, 'offline_scores.npy'),
offline_scores)
for log in logs:
log['iter'] = iters_eval
all_logs = all_logs + logs
iters_eval += 1
with timed("sampling"):
seg = seg_gen.__next__()
add_vtarg_and_adv(seg, gamma, lam)
# ob, ac, atarg, ret, td1ret = map(np.concatenate, (obs, acs, atargs, rets, td1rets))
ob, ac, atarg, tdlamret = seg["ob"], seg["ac"], seg["adv"], seg[
"tdlamret"]
vpredbefore = seg["vpred"] # predicted value function before udpate
atarg = (atarg - atarg.mean()
) / atarg.std() # standardized advantage function estimate
if hasattr(pi, "ret_rms"): pi.ret_rms.update(tdlamret)
if hasattr(pi, "ob_rms"):
pi.ob_rms.update(ob) # update running mean/std for policy
args = seg["ob"], seg["ac"], atarg
fvpargs = [arr[::5] for arr in args]
def fisher_vector_product(p):
return allmean(compute_fvp(p, *fvpargs)) + cg_damping * p
assign_old_eq_new() # set old parameter values to new parameter values
with timed("computegrad"):
*lossbefore, g = compute_lossandgrad(*args)
lossbefore = allmean(np.array(lossbefore))
g = allmean(g)
if np.allclose(g, 0):
logger.log("Got zero gradient. not updating")
else:
with timed("cg"):
stepdir = cg(fisher_vector_product,
g,
cg_iters=cg_iters,
verbose=rank == 0)
assert np.isfinite(stepdir).all()
shs = .5 * stepdir.dot(fisher_vector_product(stepdir))
lm = np.sqrt(shs / max_kl)
# logger.log("lagrange multiplier:", lm, "gnorm:", np.linalg.norm(g))
fullstep = stepdir / lm
expectedimprove = g.dot(fullstep)
surrbefore = lossbefore[0]
stepsize = 1.0
thbefore = get_flat()
for _ in range(10):
thnew = thbefore + fullstep * stepsize
set_from_flat(thnew)
meanlosses = surr, kl, *_ = allmean(
np.array(compute_losses(*args)))
improve = surr - surrbefore
logger.log("Expected: %.3f Actual: %.3f" %
(expectedimprove, improve))
if not np.isfinite(meanlosses).all():
logger.log("Got non-finite value of losses -- bad!")
elif kl > max_kl * 1.5:
logger.log("violated KL constraint. shrinking step.")
elif improve < 0:
logger.log("surrogate didn't improve. shrinking step.")
else:
logger.log("Stepsize OK!")
break
stepsize *= .5
else:
logger.log("couldn't compute a good step")
set_from_flat(thbefore)
if nworkers > 1 and iters_so_far % 20 == 0:
paramsums = MPI.COMM_WORLD.allgather(
(thnew.sum(), vfadam.getflat().sum())) # list of tuples
assert all(
np.allclose(ps, paramsums[0]) for ps in paramsums[1:])
for (lossname, lossval) in zip(loss_names, meanlosses):
logger.record_tabular(lossname, lossval)
with timed("vf"):
for _ in range(vf_iters):
for (mbob, mbret) in dataset.iterbatches(
(seg["ob"], seg["tdlamret"]),
include_final_partial_batch=False,
batch_size=64):
g = allmean(compute_vflossandgrad(mbob, mbret))
vfadam.update(g, vf_stepsize)
logger.record_tabular("ev_tdlam_before",
explained_variance(vpredbefore, tdlamret))
lrlocal = (seg["ep_lens"], seg["ep_rets"]) # local values
listoflrpairs = MPI.COMM_WORLD.allgather(lrlocal) # list of tuples
lens, rews = map(flatten_lists, zip(*listoflrpairs))
lenbuffer.extend(lens)
rewbuffer.extend(rews)
logger.record_tabular("EpLenMean", np.mean(lenbuffer))
logger.record_tabular("EpRewMean", np.mean(rewbuffer))
logger.record_tabular("EpThisIter", len(lens))
episodes_so_far += len(lens)
timesteps_so_far += sum(lens)
iters_so_far += 1
logger.record_tabular("EpisodesSoFar", episodes_so_far)
logger.record_tabular("TimestepsSoFar", timesteps_so_far)
logger.record_tabular("TimeElapsed", time.time() - tstart)
ep_rew_mean = np.mean(rewbuffer)
online_scores.append(ep_rew_mean)
np.save(os.path.join(checkpoint_dir_out, 'online_scores.npy'),
online_scores)
# Saving best
if iters_so_far % checkpoint_freq == 0 and ep_rew_mean > best_rew and checkpoint_dir_out is not None:
pi.save(os.path.join(checkpoint_dir_out, 'best'))
best_rew = ep_rew_mean
logger.log('Saved policy weights as %s' %
os.path.join(checkpoint_dir_out, 'best.npy'))
if rank == 0:
logger.dump_tabular()
return pi
def learn(
env,
policy_fn,
*,
timesteps_per_actorbatch, # timesteps per actor per update
optim_stepsize,
optim_batchsize, # optimization hypers
gamma,
lam, # advantage estimation
entcoeff=0.0,
max_episodes=0,
max_iters=0,
max_seconds=0, # time constraint
callback=None, # you can do anything in the callback, since it takes locals(), globals()
adam_epsilon=1e-5,
schedule='constant', # annealing for stepsize parameters (epsilon and adam)
args):
# Setup losses and stuff`
# ----------------------------------------
ob_space = env.observation_space
ac_space = env.action_space
pi = policy_fn("pi", ob_space,
ac_space) # Construct network for new policy
oldpi = policy_fn("oldpi", ob_space, ac_space) # Network for old policy
# Ops to reassign params from new to old
assign_old_eq_new = U.function(
[], [],
updates=[
tf.assign(oldv, newv)
for (oldv,
newv) in zipsame(oldpi.get_variables(), pi.get_variables())
])
atarg = tf.placeholder(
dtype=tf.float32,
shape=[None]) # Target advantage function (if applicable)
ret = tf.placeholder(dtype=tf.float32, shape=[None]) # Empirical return
lrmult = tf.placeholder(
name='lrmult', dtype=tf.float32,
shape=[]) # learning rate multiplier, updated with schedule
ob = U.get_placeholder_cached(name="ob")
ac = pi.pdtype.sample_placeholder([None])
kloldnew = oldpi.pd.kl(pi.pd)
ent = pi.pd.entropy()
meankl = tf.reduce_mean(kloldnew)
meanent = tf.reduce_mean(ent)
pol_entpen = (-entcoeff) * meanent
newprob = tf.exp(pi.pd.logp(ac))
oldprob = tf.exp(oldpi.pd.logp(ac))
ratio = newprob / oldprob
kl = pi.pd.kl(oldpi.pd)
mean_kl = tf.reduce_mean(kl)
get_kl = U.function([ob, ac], kl)
get_mean_kl = U.function([ob, ac], mean_kl)
threshold = kl < args.kl_threshold
threshold = tf.cast(threshold, tf.float32)
pol_surr = (kl - ratio * atarg / args.sepg_lam) * threshold
pol_surr = tf.reduce_mean(pol_surr)
vf_loss = tf.reduce_mean(tf.square(pi.vpred - ret))
total_loss = pol_surr + pol_entpen + vf_loss
losses = [pol_surr, pol_entpen, vf_loss, meankl, meanent]
loss_names = ["pol_surr", "pol_entpen", "vf_loss", "kl", "ent"]
var_list = pi.get_trainable_variables()
lossandgrad = U.function([ob, ac, atarg, ret, lrmult],
losses + [U.flatgrad(total_loss, var_list)])
adam = MpiAdam(var_list, epsilon=adam_epsilon)
compute_losses = U.function([ob, ac, atarg, ret, lrmult], losses)
U.initialize()
adam.sync()
# Prepare for rollouts
# ----------------------------------------
seg_gen = traj_segment_generator(pi,
env,
timesteps_per_actorbatch,
stochastic=True)
episodes_so_far = 0
timesteps_so_far = 0
iters_so_far = 0
tstart = time.time()
lenbuffer = deque(maxlen=100) # rolling buffer for episode lengths
rewbuffer = deque(maxlen=100) # rolling buffer for episode rewards
running_scores = []
assert sum([
max_iters > 0, args.num_timesteps > 0, max_episodes > 0,
max_seconds > 0
]) == 1, "Only one time constraint permitted"
while True:
if callback: callback(locals(), globals())
if args.num_timesteps and timesteps_so_far >= args.num_timesteps:
break
elif max_episodes and episodes_so_far >= max_episodes:
break
elif max_iters and iters_so_far >= max_iters:
break
elif max_seconds and time.time() - tstart >= max_seconds:
break
if schedule == 'constant':
cur_lrmult = 1.0
elif schedule == 'linear':
cur_lrmult = max(
1.0 - float(timesteps_so_far) / args.num_timesteps, 0)
else:
raise NotImplementedError
if MPI.COMM_WORLD.Get_rank() == 0:
logger.log("********** Iteration %i ************" % iters_so_far)
seg = seg_gen.__next__()
add_vtarg_and_adv(seg, gamma, lam)
# ob, ac, atarg, ret, td1ret = map(np.concatenate, (obs, acs, atargs, rets, td1rets))
ob, ac, atarg, tdlamret = seg["ob"], seg["ac"], seg["adv"], seg[
"tdlamret"]
vpredbefore = seg["vpred"] # predicted value function before udpate
atarg = (atarg - atarg.mean()) / (
atarg.std() + 1e-8) # standardized advantage function estimate
optim_batchsize = optim_batchsize or ob.shape[0]
if hasattr(pi, "ob_rms"):
pi.ob_rms.update(ob) # update running mean/std for policy
assign_old_eq_new() # set old parameter values to new parameter values
d = Dataset(dict(ob=ob, ac=ac, atarg=atarg, vtarg=tdlamret),
shuffle=not pi.recurrent)
# Here we do a bunch of optimization epochs over the data
for num_epoch in count():
losses = [
] # list of tuples, each of which gives the loss for a minibatch
for batch in d.iterate_once(optim_batchsize):
*newlosses, g = lossandgrad(batch["ob"], batch["ac"],
batch["atarg"], batch["vtarg"],
cur_lrmult)
g = np.nan_to_num(g)
adam.update(g, optim_stepsize * cur_lrmult)
losses.append(newlosses)
agg_mean_kl = get_mean_kl(ob, ac)
if agg_mean_kl > args.agg_kl_threshold or num_epoch == args.optim_epochs:
break
lrlocal = (seg["ep_lens"], seg["ep_rets"]) # local values
listoflrpairs = MPI.COMM_WORLD.allgather(lrlocal) # list of tuples
lens, rews = map(flatten_lists, zip(*listoflrpairs))
rewbuffer.extend(rews)
mean_score = None
if rewbuffer:
mean_score = np.mean(rewbuffer)
running_scores.append((timesteps_so_far, mean_score))
episodes_so_far += len(lens)
timesteps_so_far += sum(lens)
iters_so_far += 1
if MPI.COMM_WORLD.Get_rank() == 0:
logger.record_tabular("EpRewMean", mean_score)
logger.record_tabular("EpThisIter", len(lens))
logger.record_tabular("EpisodesSoFar", episodes_so_far)
logger.record_tabular("TimestepsSoFar", timesteps_so_far)
logger.record_tabular("TimeElapsed", time.time() - tstart)
logger.record_tabular("NumEpoch", num_epoch)
logger.dump_tabular()
return running_scores
def run(config):
def create_loss():
train_dataset_args = create_dataset(config,
split="train",
shuffle=True,
repeat=True)
test_dataset_args = create_dataset(config,
split="test",
shuffle=True,
repeat=True)
valid_dataset_args = create_dataset(config,
split="valid",
shuffle=True,
repeat=True)
inputs, targets, lengths, params, _, _ = train_dataset_args
test_inputs, test_targets, test_lengths, _, _, _ = test_dataset_args
valid_inputs, valid_targets, valid_lengths, _, _, _ = valid_dataset_args
cell = create_cell(config, train_dataset_args)
if config.bound == "iwae":
ll_per_seq, kl, log_weight, log_ess, trajectories = \
basic_bounds.iwae(cell, (inputs, targets), lengths, num_samples=config.num_samples)
else:
raise ValueError("Undefined bound %s" % config.bound)
if config.test_bound == "iwae":
valid_ll_per_seq, _, _, _, _ = \
basic_bounds.iwae(cell, (valid_inputs, valid_targets), valid_lengths, num_samples=config.test_num_samples)
else:
raise ValueError("Undefined bound %s" % config.test_bound)
ll_per_t = tf.reduce_mean(ll_per_seq / tf.to_float(lengths))
valid_ll_per_t = tf.reduce_mean(valid_ll_per_seq /
tf.to_float(valid_lengths))
return cell, ll_per_t, valid_ll_per_t, trajectories, lengths
def create_graph():
global_step = tf.train.get_or_create_global_step()
cell, bound, valid_bound, trajectories, lengths = create_loss()
loss = -bound
opt = tf.train.AdamOptimizer(config.learning_rate)
if config.model_train:
grad_theta = opt.compute_gradients(loss,
var_list=tf.trainable_variables(
"%s/theta" % config.cell))
train_op_theta = opt.apply_gradients(grad_theta,
global_step=global_step)
else:
train_op_theta = tf.constant(1)
if config.algorithm == 'reparam':
grad_phi = opt.compute_gradients(
loss, var_list=tf.trainable_variables('prop_phi'))
train_op_phi = opt.apply_gradients(grad_phi,
global_step=global_step)
else:
train_op_phi = tf.constant(1)
return cell, bound, valid_bound, trajectories, lengths, train_op_theta, train_op_phi, global_step
valid_best = -1000000
cell, bound, valid_bound, trajectories, lengths, train_op_theta, train_op_phi, global_step = create_graph(
)
sess = U.get_session()
U.initialize()
cur_step = 0
saver = tf.train.Saver(max_to_keep=1)
valid_saver = tf.train.Saver(max_to_keep=1)
model_savepath = config.logdir + '/model.ckpt'
valid_best_model_savepath = config.logdir + '/valid_best/valid_best_model.ckpt'
if not os.path.exists(config.logdir):
os.makedirs(config.logdir)
os.makedirs(config.logdir + '/valid_best')
ckpt = tf.train.get_checkpoint_state(config.logdir)
if ckpt and ckpt.model_checkpoint_path:
saver.restore(sess, ckpt.model_checkpoint_path)
valid_saver.restore(sess, ckpt.model_checkpoint_path)
cur_step = int(ckpt.model_checkpoint_path.split('-')[-1])
print('Model and log loaded! (checkpoint_path=%s, cur_step=%d)' %
(ckpt.model_checkpoint_path, cur_step))
while cur_step < config.max_iter + 1:
if config.algorithm == 'reparam':
_, _, bound_value, valid_bound_value = sess.run(
[train_op_theta, train_op_phi, bound, valid_bound])
elif "reinforce" in config.algorithm or "vimco" in config.algorithm or "vifle" in config.algorithm or "fr" in config.algorithm:
_, bound_value, raw_seg, valid_bound_value, run_lengths = sess.run(
[train_op_theta, bound, trajectories, valid_bound, lengths])
cell.prop_update.update(raw_seg, run_lengths)
else:
raise ValueError("Undefined algorithm %s" % config.algorithm)
if valid_bound_value > valid_best and cur_step > config.init_steps:
valid_best = valid_bound_value
valid_best_model_saved_path = valid_saver.save(
sess, valid_best_model_savepath, global_step=cur_step)
print('Model saved: %s' % valid_best_model_saved_path)
# for save - current work
if cur_step % config.save_every == 0:
model_saved_path = saver.save(sess,
model_savepath,
global_step=cur_step)
print('Model saved: %s' % model_saved_path)
cur_step += 1
def main():
L.configure('/home/metalabadmin/exp/freeway',
format_strs=['stdout', 'csv', 'tensorboard'])
env = gym.make('Freeway-v0')
env = wrapper.wrap_deepmind(env, frame_stack=True, scale=True)
optimizer = tf.train.AdamOptimizer(learning_rate=0.0001)
network = Q_network(env.observation_space,
env.action_space.n,
optimizer,
gamma=0.99,
scope='freeway')
m_controller = MetaController(network, env.action_space.n)
# Create the schedule for exploration starting from 1.
exploration = LinearSchedule(schedule_timesteps=int(0.1 * 1e7),
initial_p=1.0,
final_p=0.02)
replay = ReplayBuffer(50000)
# get default tf_session
sess = U.get_session()
U.initialize()
sess.run(m_controller.network.update_target_op)
step = 0
episodes = 0
rewards = 0
mean_100ep_reward = 0
total_reward = []
saved_mean_reward = None
ob = env.reset()
while step <= 1e7:
ep = exploration.value(step)
ob_reshaped = np.reshape(ob, (1, ) + env.observation_space.shape)
act = m_controller.sample_act(sess, ob_reshaped, update_eps=ep)[0]
ob_tp1, reward_t, done_t, info = env.step(act)
env.render()
rewards += reward_t
replay.add(ob, act, reward_t, ob_tp1, float(done_t))
ob = ob_tp1
# train every 4 steps
if step >= 1000 and step % 4 == 0:
obs, acts, rewards_t, obs_tp1, dones_t = replay.sample(64)
weights, batch_idxes = np.ones_like(rewards_t), None
# get q estimate for tp1 as 'supervised'
obs_tp1_reshaped = np.reshape(obs_tp1,
(64, ) + env.observation_space.shape)
q_tp1 = m_controller.get_q(sess, obs_tp1_reshaped)[0]
td_error = m_controller.train(sess, obs, acts, rewards_t, obs_tp1,
dones_t, weights, q_tp1)
step += 1
if step >= 1000 and step % 1000 == 0:
sess.run(m_controller.network.update_target_op)
if done_t:
ob = env.reset()
total_reward.append(rewards)
episodes += 1
rewards = 0
print('step %d done %s, ep %.2f' % (step, str(done_t), ep))
mean_100ep_reward = round(np.mean(total_reward[-101:-1]), 1)
if episodes % 10 == 0 and episodes != 0:
print('date time %s' % str(datetime.now()))
L.record_tabular("steps", step)
L.record_tabular("episodes", episodes)
L.record_tabular("mean 100 episode reward", mean_100ep_reward)
L.dump_tabular()
if step % 1000 == 0:
if saved_mean_reward is None or mean_100ep_reward > saved_mean_reward:
L.log("Saving model due to mean reward increase: {} -> {}".
format(saved_mean_reward, mean_100ep_reward))
U.save_variables('./freewaymodel.ckpt')
model_saved = True
saved_mean_reward = mean_100ep_reward
