agent_mtl = NPGFTW(e,
policy_mtl,
baseline_mtl,
normalized_step_size=1,
seed=SEED,
save_logs=True,
new_col_mode='max_k')
for task_id in range(num_tasks):
ts = timer.time()
train_agent(job_name=job_name_lpgftw_seed,
agent=agent_mtl,
seed=SEED,
niter=50,
gamma=0.995,
gae_lambda=0.97,
num_cpu=num_cpu,
sample_mode='trajectories',
num_traj=10,
save_freq=5,
evaluation_rollouts=0,
task_id=task_id)
agent_mtl.add_approximate_cost(N=10, task_id=task_id, num_cpu=num_cpu)
iterdir = job_name_lpgftw_seed + '/iterations/task_{}/'.format(task_id)
os.makedirs(iterdir, exist_ok=True)
policy_file = open(iterdir + 'policy_updated.pickle', 'wb')
pickle.dump(agent_mtl.policy, policy_file)
policy_file.close()
print("time taken for linear policy training = %f" %
(timer.time() - ts))
def single_process(job):
job_start_time = timer.time()
# Allow process to parallelize things internally
curr_proc = mp.current_process()
curr_proc.daemon = False
os.chdir(cwd)
dirpath = os.path.join(job['save_dir'], job['job_name'])
os.makedirs(dirpath, exist_ok=True)
# start job
os.chdir(cwd)
job_start_time = timer.time()
print('Started New Job : ', job['job_name'], '=======================')
print('Job specifications : \n', job)
# Make Env
e = GymEnv(job['env_name'])
# Make baseline
baseline = MLPBaseline(e.spec)
# save job details
job['horizon'] = e.horizon
job['ctrl_timestep'] = e.env.env.dt
job['sim_timestep'] = e.env.env.model.opt.timestep
# job['sim_skip'] = e.env.env.skip
job_data_file = open(dirpath + '/job_data.txt', 'w')
pprint.pprint(job, stream=job_data_file)
job_data_file.close()
# Make policy (???vik: sizes are hard coded)
if 'init_policy' in job:
policy = MLP(e.spec,
init_log_std=job['init_std'],
hidden_sizes=(32, 32),
seed=job['seed'])
loaded_policy = pickle.load(open(job['init_policy'], 'rb'))
loaded_params = loaded_policy.get_param_values()
print('log std values in loaded policy = ')
print(params[-policy.m:])
# NOTE: if the log std is too small
# (say <-2.0, it is problem dependent and intuition should be used)
# then we need to bump it up so that it explores
# params[-policy.m:] += 1.0
policy.set_param_values(loaded_params)
del job['init_policy']
else:
policy = MLP(e.spec,
init_log_std=job['init_std'],
hidden_sizes=(32, 32),
seed=job['seed'])
# Agent
agent = NPG(e, policy, baseline, seed=job['seed'], \
normalized_step_size=job['normalized_step_size'], \
save_logs=job['save_logs'], FIM_invert_args=job['FIM_invert_args'])
# Train Agent
train_agent(
job_name=dirpath,
agent=agent,
seed=job['seed'],
niter=job['niter'],
gamma=job['gamma'],
gae_lambda=job['gae_lambda'],
num_cpu=job['num_cpu'],
sample_mode=job['sample_mode'],
num_traj=job['num_traj'],
evaluation_rollouts=job['evaluation_rollouts'],
save_freq=job['save_freq'],
plot_keys={'stoc_pol_mean', 'stoc_pol_std'},
)
total_job_time = timer.time() - job_start_time
print('Job', job['job_name'],
'took %f seconds ==============' % total_job_time)
return total_job_time
SEED = 500
e = GymEnv('Walker2d-v2')
policy = LinearPolicy(e.spec, seed=SEED)
baseline = MLPBaseline(e.spec,
reg_coef=1e-3,
batch_size=64,
epochs=2,
learn_rate=1e-3)
agent = NPG(e,
policy,
baseline,
normalized_step_size=0.1,
seed=SEED,
save_logs=True)
ts = timer.time()
train_agent(job_name='walker_nominal',
agent=agent,
seed=SEED,
niter=500,
gamma=0.995,
gae_lambda=0.97,
num_cpu=4,
sample_mode='trajectories',
num_traj=50,
save_freq=5,
evaluation_rollouts=5)
print("time taken for linear policy training = %f" % (timer.time() - ts))
import time as timer
SEED = 500
e = GymEnv('mjrl_point_mass-v0')
policy = MLP(e.spec, hidden_sizes=(32, 32), seed=SEED)
baseline = QuadraticBaseline(e.spec)
agent = NPG(e,
policy,
baseline,
normalized_step_size=0.2,
seed=SEED,
save_logs=True)
ts = timer.time()
train_agent(job_name='vis_exp',
agent=agent,
seed=SEED,
niter=30,
gamma=0.95,
gae_lambda=0.97,
num_cpu=1,
sample_mode='trajectories',
num_traj=100,
save_freq=5,
evaluation_rollouts=None)
print("time taken = %f" % (timer.time() - ts))
e.visualize_policy(policy,
num_episodes=5,
horizon=e.horizon,
mode='evaluation')
import argparse
import time as timer
SEED = 500
#
e = GymEnv("half-cheetah-joint-v0")
policy = MLP(e.spec, hidden_sizes=(32, 32), seed=SEED)
baseline = MLPBaseline(e.spec,
reg_coef=1e-3,
batch_size=64,
epochs=2,
learn_rate=1e-3)
agent = PPO(e, policy, baseline, save_logs=True)
print("========================================")
print("Starting policy learning")
print("========================================")
ts = timer.time()
train_agent(job_name='beta_test',
agent=agent,
seed=SEED,
niter=2000,
gamma=0.995,
gae_lambda=0.97,
num_cpu=5,
sample_mode='trajectories',
num_traj=10,
save_freq=50,
evaluation_rollouts=5)
print("time taken = %f" % (timer.time() - ts))
def single_process(job):
job_start_time = timer.time()
# Allow process to parallelize things internally
curr_proc = mp.current_process()
curr_proc.daemon = False
# Create a directory for the job results.
job_dir = os.path.join(job['output_dir'])
if not os.path.isdir(job_dir):
os.mkdir(job_dir)
# start job
job_start_time = timer.time()
print('Started New Job : ', job['job_name'], '=======================')
print('Job specifications : \n', job)
# Make Env
env_name = job['env_name']
# adept_envs.global_config.set_config(env_name, {
# 'robot_params': job['robot'],
# **job.get('env_params', {}),
# })
e = GymEnv(env_name)
# Make baseline
baseline = MLPBaseline(e.spec)
# save job details
job['horizon'] = e.horizon
job['ctrl_timestep'] = e.env.env.dt
job['sim_timestep'] = e.env.env.model.opt.timestep
# job['sim_skip'] = e.env.env.skip
with open(os.path.join(job_dir, 'job_data.txt'), 'w') as job_data_file:
pprint.pprint(job, stream=job_data_file)
if 'init_policy' in job:
policy = MLP(e.spec, init_log_std=job['init_std'], hidden_sizes=(32,32), seed=job['seed'])
loaded_policy = pickle.load(open(job['init_policy'], 'rb'))
loaded_params = loaded_policy.get_param_values()
print("log std values in loaded policy = ")
print(loaded_params[-policy.m:])
# NOTE: if the log std is too small
# (say <-2.0, it is problem dependent and intuition should be used)
# then we need to bump it up so that it explores
loaded_params[-policy.m:] += job['init_std']
policy.set_param_values(loaded_params)
del job['init_policy']
else:
policy = MLP(
e.spec,
init_log_std=job['init_std'],
hidden_sizes=job['hidden_sizes'],
# hidden_sizes=(32, 32),
seed=job['seed'])
# Agent
agent = NPG(
e,
policy,
baseline,
seed=job['seed'],
normalized_step_size=job['normalized_step_size'],
save_logs=job['save_logs'],
FIM_invert_args=job['FIM_invert_args'])
# Train Agent
train_agent(
job_name=job['job_name'],
agent=agent,
# save_dir=job_dir,
seed=job['seed'],
niter=job['niter'],
gamma=job['gamma'],
gae_lambda=job['gae_lambda'],
num_cpu=job['num_cpu'],
sample_mode=job['sample_mode'],
num_traj=job.get('num_traj'),
num_samples=job.get('num_samples'),
evaluation_rollouts=job['evaluation_rollouts'],
save_freq=job['save_freq'],
plot_keys={'stoc_pol_mean', 'stoc_pol_std'},
)
total_job_time = timer.time() - job_start_time
print('Job', job['job_name'],
'took %f seconds ==============' % total_job_time)
return total_job_time
# RL Loop
# ===============================================================================
rl_agent = DAPG(e,
policy,
baseline,
demo_paths,
normalized_step_size=job_data['rl_step_size'],
lam_0=job_data['lam_0'],
lam_1=job_data['lam_1'],
seed=job_data['seed'],
save_logs=True)
print("========================================")
print("Starting reinforcement learning phase")
print("========================================")
ts = timer.time()
train_agent(job_name=JOB_DIR,
agent=rl_agent,
seed=job_data['seed'],
niter=job_data['rl_num_iter'],
gamma=job_data['rl_gamma'],
gae_lambda=job_data['rl_gae'],
num_cpu=job_data['num_cpu'],
sample_mode='trajectories',
num_traj=job_data['rl_num_traj'],
save_freq=job_data['save_freq'],
evaluation_rollouts=job_data['eval_rollouts'])
print("time taken = %f" % (timer.time() - ts))
SEED = 500
e = GymEnv('mjrl_swimmer-v0')
policy = MLP(e.spec, hidden_sizes=(32, 32), seed=SEED)
baseline = MLPBaseline(e.spec,
reg_coef=1e-3,
batch_size=64,
epochs=5,
learn_rate=1e-3)
agent = NPG(e,
policy,
baseline,
normalized_step_size=0.1,
seed=SEED,
save_logs=True)
ts = timer.time()
train_agent(
job_name='swimmer_exp1',
agent=agent,
seed=SEED,
niter=50,
gamma=0.995,
gae_lambda=0.97,
num_cpu=1,
sample_mode='trajectories',
num_traj=10, # samples = 10*500 = 5000
save_freq=5,
evaluation_rollouts=5)
print("time taken = %f" % (timer.time() - ts))
print("========================================")
print("Finetuning with DAPG")
baseline = MLPBaseline(e.spec,
reg_coef=1e-3,
batch_size=64,
epochs=2,
learn_rate=1e-3)
agent = DAPG(e,
policy,
baseline,
demo_paths=demo_paths,
normalized_step_size=0.1,
seed=SEED,
lam_0=1e-2,
lam_1=0.99,
save_logs=True)
ts = timer.time()
train_agent(job_name='relocate_demo_init_dapg',
agent=agent,
seed=SEED,
niter=100,
gamma=0.995,
gae_lambda=0.97,
num_cpu=5,
sample_mode='trajectories',
num_traj=200,
save_freq=25,
evaluation_rollouts=20)
print("time taken = %f" % (timer.time() - ts))
def launch_job(tag, variant):
print(len(variant))
seed, env, algo, optim, curv_type, lr, batch_size, cg_iters, cg_residual_tol, cg_prev_init_coef, \
cg_precondition_empirical, cg_precondition_regu_coef, cg_precondition_exp, \
shrinkage_method, lanczos_amortization, lanczos_iters, approx_adaptive, betas, use_nn_policy, gn_vfn_opt, total_samples = variant
beta1, beta2 = betas
iters = int(total_samples / batch_size)
# NN policy
# ==================================
e = GymEnv(env)
if use_nn_policy:
policy = MLP(e.spec, hidden_sizes=(64, ), seed=seed)
else:
policy = LinearPolicy(e.spec, seed=seed)
vfn_batch_size = 256 if gn_vfn_opt else 64
vfn_epochs = 2 if gn_vfn_opt else 2
# baseline = MLPBaseline(e.spec, reg_coef=1e-3, batch_size=64, epochs=2, learn_rate=1e-3)
baseline = MLPBaseline(e.spec,
reg_coef=1e-3,
batch_size=vfn_batch_size,
epochs=2,
learn_rate=1e-3,
use_gauss_newton=gn_vfn_opt)
# agent = NPG(e, policy, baseline, normalized_step_size=0.005, seed=SEED, save_logs=True)
common_kwargs = dict(lr=lr,
curv_type=curv_type,
cg_iters=cg_iters,
cg_residual_tol=cg_residual_tol,
cg_prev_init_coef=cg_prev_init_coef,
cg_precondition_empirical=cg_precondition_empirical,
cg_precondition_regu_coef=cg_precondition_regu_coef,
cg_precondition_exp=cg_precondition_exp,
shrinkage_method=shrinkage_method,
lanczos_amortization=lanczos_amortization,
lanczos_iters=lanczos_iters,
batch_size=batch_size)
if optim == 'ngd':
optimizer = fisher_optim.NGD(policy.trainable_params, **common_kwargs)
elif optim == 'natural_adam':
optimizer = fisher_optim.NaturalAdam(
policy.trainable_params,
**common_kwargs,
betas=(beta1, beta2),
assume_locally_linear=approx_adaptive)
elif optim == 'natural_adagrad':
optimizer = fisher_optim.NaturalAdagrad(
policy.trainable_params,
**common_kwargs,
betas=(beta1, beta2),
assume_locally_linear=approx_adaptive)
elif optim == 'natural_amsgrad':
optimizer = fisher_optim.NaturalAmsgrad(
policy.trainable_params,
**common_kwargs,
betas=(beta1, beta2),
assume_locally_linear=approx_adaptive)
if algo == 'trpo':
from mjrl.algos.trpo_delta import TRPO
agent = TRPO(e, policy, baseline, optimizer, seed=seed, save_logs=True)
# agent = TRPO(e, policy, baseline, seed=seed, save_logs=True)
else:
from mjrl.algos.npg_cg_delta import NPG
agent = NPG(e, policy, baseline, optimizer, seed=seed, save_logs=True)
save_dir = build_log_dir(tag, variant)
try:
os.makedirs(save_dir)
except:
pass
# print ("Iters:", iters, ", num_traj: ", str(batch_size//1000))
train_agent(job_name=save_dir,
agent=agent,
seed=seed,
niter=iters,
gamma=0.995,
gae_lambda=0.97,
num_cpu=1,
sample_mode='samples',
num_samples=batch_size,
save_freq=5,
evaluation_rollouts=5,
verbose=False) #True)
def train(cfg, run_no, multiple_runs, seed):
# ===============================================================================
# Train Loop
# ===============================================================================
gpus_available = setup_gpus()
env_name, job_name = parse_task(cfg)
env = GymEnv(env_name, **cfg['env_kwargs'])
policy = MLP(env.spec, hidden_sizes=tuple(cfg['policy_size']), seed=seed)
baseline = MLPBaseline(env.spec,
reg_coef=1e-3,
batch_size=cfg['value_function']['batch_size'],
epochs=cfg['value_function']['epochs'],
learn_rate=cfg['value_function']['lr'],
use_gpu=False)
# Get demonstration data if necessary and behavior clone
print("========================================")
print("Collecting expert demonstrations")
print("========================================")
demo_filename = cfg['demo_file']
if cfg['demo_file'] != None:
demo_paths = pickle.load(open(demo_filename, 'rb'))
else:
demo_paths = None
if 'demo_file' in cfg['BC'] and cfg['BC']['demo_file'] != 'default':
bc_demo_file_path = cfg['BC']['demo_file']
if cfg['train']['use_timestamp']:
bc_demo_file_path = bc_demo_file_path.replace(
'v0', 'v0_timestamp_inserted')
bc_demo_paths = pickle.load(open(bc_demo_file_path, 'rb'))
else:
bc_demo_paths = demo_paths
if 'num_demo' in cfg and cfg['num_demo']:
demo_paths = demo_paths[:cfg['num_demo']]
if cfg['algorithm'] == 'DAPG_based_IRL':
if 'get_paths_for_initialisation' in cfg['based_IRL']:
if cfg['based_IRL']['get_paths_for_initialisation']:
bc_demo_paths = add_dumped_paths_for_BC(demo_paths, cfg)
ts = timer.time()
if bc_demo_paths is not None and cfg['BC']['epochs'] > 0:
print("========================================")
print("Running BC with expert demonstrations")
print("========================================")
bc_agent = BC(bc_demo_paths[:25],
policy=policy,
epochs=cfg['BC']['epochs'],
batch_size=cfg['BC']['batch_size'],
lr=cfg['BC']['lr'],
loss_type='MSE',
set_transforms=True)
bc_agent.train()
print("========================================")
print("BC training complete !!!")
print("time taken = %f" % (timer.time() - ts))
print("========================================")
if cfg['algorithm'] == 'IRL' or cfg['algorithm'] == 'DAPG_based_IRL':
IRL_cfg = cfg
if cfg['algorithm'] == 'DAPG_based_IRL':
IRL_job_cfg_path = os.path.join("Runs",
cfg['based_IRL']['IRL_job'],
"config.yaml")
IRL_cfg = yamlreader.yaml_load(IRL_job_cfg_path)
irl_model = get_irl_model(env, demo_paths, IRL_cfg, seed)
if cfg['algorithm'] == 'DAPG_based_IRL':
full_irl_model_checkpoint_path = os.path.join(
'Runs', cfg['based_IRL']['IRL_job'])
if cfg['based_IRL']['IRL_run_no'] is not None:
full_irl_model_checkpoint_path = os.path.join(
full_irl_model_checkpoint_path,
'run_' + str(cfg['based_IRL']['IRL_run_no']))
if cfg['based_IRL']['IRL_step'] is not None:
irl_model.load_iteration(
path=full_irl_model_checkpoint_path,
iteration=cfg['based_IRL']['IRL_step'])
else:
irl_model.load_last(path=full_irl_model_checkpoint_path)
irl_model.eval(
demo_paths
) # required to load model completely from the given path before changin to different path during training
if cfg['eval_rollouts'] > 0:
score = env.evaluate_policy(policy,
num_episodes=cfg['eval_rollouts'],
mean_action=True)
print("Score with behavior cloning = %f" % score[0][0])
if not cfg['use_DAPG']:
# We throw away the demo data when training from scratch or fine-tuning with RL without explicit augmentation
demo_paths = None
# ===============================================================================
# RL Loop
# ===============================================================================
irl_kwargs = None
if cfg['algorithm'] == 'IRL' or cfg['algorithm'] == 'DAPG_based_IRL':
if cfg['algorithm'] == 'DAPG_based_IRL' or cfg['IRL'][
'generator_alg'] == 'DAPG':
generator_algorithm = DAPG
generator_args = dict(
demo_paths=demo_paths,
normalized_step_size=cfg['RL']['step_size'],
seed=seed,
lam_0=cfg['RL']['lam_0'],
lam_1=cfg['RL']['lam_1'],
save_logs=cfg['save_logs'],
augmentation=cfg['train']['augmentation'],
entropy_weight=cfg['train']['entropy_weight'])
elif cfg['IRL']['generator_alg'] == 'PPO':
generator_algorithm = PPO
generator_args = dict(
demo_paths=demo_paths,
epochs=cfg['PPO']['epochs'],
mb_size=cfg['PPO']['batch_size'],
target_kl_dist=cfg['PPO']['target_kl_dist'],
seed=seed,
lam_0=cfg['RL']['lam_0'],
lam_1=cfg['RL']['lam_1'],
save_logs=cfg['save_logs'],
clip_coef=cfg['PPO']['clip_coef'],
learn_rate=cfg['PPO']['lr'],
augmentation=cfg['train']['augmentation'],
entropy_weight=cfg['train']['entropy_weight'])
else:
raise ValueError("Generator algorithm name",
cfg['IRL']['generator_alg'], "not supported")
irl_class = irl_training_class(generator_algorithm)
rl_agent = irl_class(
env,
policy,
baseline,
train_irl=cfg['algorithm'] != 'DAPG_based_IRL',
discr_lr=IRL_cfg['IRL']['discr']['lr'],
irl_batch_size=IRL_cfg['IRL']['discr']['batch_size'],
lower_lr_on_main_loop_percentage=IRL_cfg['IRL']['discr']
['lower_lr_on_main_loop_percentage'],
irl_model=irl_model,
**generator_args)
irl_kwargs = dict(policy=dict(
min_updates=1,
max_updates=IRL_cfg['IRL']['max_gen_updates']
if cfg['algorithm'] != 'DAPG_based_IRL' else 0,
steps_till_max=IRL_cfg['IRL']['steps_till_max_gen_updates']))
elif cfg['algorithm'] == 'DAPG':
rl_agent = DAPG(env,
policy,
baseline,
demo_paths=demo_paths,
normalized_step_size=cfg['RL']['step_size'],
lam_0=cfg['RL']['lam_0'],
lam_1=cfg['RL']['lam_1'],
seed=seed,
save_logs=cfg['save_logs'],
augmentation=cfg['train']['augmentation'],
entropy_weight=cfg['train']['entropy_weight'])
elif cfg['algorithm'] == 'PPO':
rl_agent = PPO(env,
policy,
baseline,
demo_paths=demo_paths,
epochs=cfg['PPO']['epochs'],
mb_size=cfg['PPO']['batch_size'],
target_kl_dist=cfg['PPO']['target_kl_dist'],
seed=seed,
lam_0=cfg['RL']['lam_0'],
lam_1=cfg['RL']['lam_1'],
save_logs=cfg['save_logs'],
clip_coef=cfg['PPO']['clip_coef'],
learn_rate=cfg['PPO']['lr'],
augmentation=cfg['train']['augmentation'],
entropy_weight=cfg['train']['entropy_weight'])
else:
raise ValueError("Algorithm name", cfg['algorithm'], "not supported")
# get IRL model kwargs if doing DAPG based on IRL
env_kwargs = cfg['env_kwargs']
if cfg['algorithm'] == 'DAPG_based_IRL':
rl_agent.irl_model = irl_model
# dump YAML config file
job_path = os.path.join("Runs", job_name)
if not os.path.isdir(job_path):
os.makedirs(job_path)
with open(os.path.join(job_path, 'config.yaml'), 'w') as f:
dump(cfg, f)
print("========================================")
print("Starting reinforcement learning phase")
print("========================================")
ts = timer.time()
train_agent(
job_name=job_name,
agent=rl_agent,
seed=seed,
niter=cfg['train']['steps'],
gamma=cfg['train']['gamma'],
gae_lambda=cfg['train']['gae_lambda'],
num_cpu=cfg['num_cpu'],
sample_mode='trajectories',
num_traj=cfg['train']['num_traj'],
save_freq=cfg['train']['save_freq'],
evaluation_rollouts=cfg['eval_rollouts'],
should_fresh_start=bool(cfg['IRL']['initialization_job'])
if cfg['algorithm'] == 'IRL' else False,
irl_kwargs=irl_kwargs,
temperature_max=cfg['IRL']['temperature_max']
if cfg['algorithm'] == 'IRL' else 0,
temperature_min=cfg['IRL']['temperature_min']
if cfg['algorithm'] == 'IRL' else 0,
plot_keys=cfg['plot_keys'],
run_no=run_no if multiple_runs else None,
env_kwargs=env_kwargs,
fixed_evaluation_init_states=cfg['fixed_evaluation_init_states'])
print("time taken = %f" % (timer.time() - ts))
e = GymEnv('mjrl_point_mass-v0')
policy = MLP(e.spec, hidden_sizes=(32, 32), seed=SEED)
baseline = MLPBaseline(e.spec,
reg_coef=1e-3,
batch_size=64,
epochs=10,
learn_rate=1e-3)
agent = NPG(e,
policy,
baseline,
normalized_step_size=0.05,
seed=SEED,
save_logs=True)
ts = timer.time()
train_agent(
job_name='point_mass_exp1',
agent=agent,
seed=SEED,
niter=50,
gamma=0.95,
gae_lambda=0.97,
num_cpu=1,
sample_mode='trajectories',
num_traj=40, # samples = 40*25 = 1000
save_freq=5,
evaluation_rollouts=None,
plot_keys=['stoc_pol_mean', 'running_score'])
print("time taken = %f" % (timer.time() - ts))
from mjrl.algos.npg_cg import NPG
from mjrl.utils.train_agent import train_agent
import mjrl.envs
import time as timer
SEED = 500
e = GymEnv('mjrl_point_mass-v0')
policy = MLP(e.spec, hidden_sizes=(32, 32), seed=SEED)
baseline = QuadraticBaseline(e.spec)
agent = NPG(e,
policy,
baseline,
normalized_step_size=0.1,
seed=SEED,
save_logs=True)
ts = timer.time()
train_agent(
job_name='point_mass_exp1',
agent=agent,
seed=SEED,
niter=50,
gamma=0.95,
gae_lambda=0.97,
num_cpu=1,
sample_mode='trajectories',
num_traj=40, # samples = 40*25 = 1000
save_freq=5,
evaluation_rollouts=10)
print("time taken = %f" % (timer.time() - ts))
def experiment(variant):
"""
This is a job script for running NPG/DAPG on hand tasks and other gym envs.
Note that DAPG generalizes PG and BC init + PG finetuning.
With appropriate settings of parameters, we can recover the full family.
"""
import mj_envs
job_data = default_job_data.copy()
job_data.update(variant)
env_params = ENV_PARAMS[variant['env_class']]
job_data.update(env_params)
assert 'algorithm' in job_data.keys()
assert any([job_data['algorithm'] == a for a in ['NPG', 'BCRL', 'DAPG']])
JOB_DIR = logger.get_snapshot_dir()
# ===============================================================================
# Train Loop
# ===============================================================================
seed = int(job_data['seedid'])
e = GymEnv(job_data['env_id'])
policy = MLP(e.spec, hidden_sizes=job_data['policy_size'], seed=seed)
baseline = MLPBaseline(e.spec,
reg_coef=1e-3,
batch_size=job_data['vf_batch_size'],
epochs=job_data['vf_epochs'],
learn_rate=job_data['vf_learn_rate'])
# Get demonstration data if necessary and behavior clone
if job_data['algorithm'] != 'NPG':
print("========================================")
print("Collecting expert demonstrations")
print("========================================")
demo_paths = load_local_or_remote_file(job_data['demo_file'], 'rb')
bc_agent = BC(demo_paths,
policy=policy,
epochs=job_data['bc_epochs'],
batch_size=job_data['bc_batch_size'],
lr=job_data['bc_learn_rate'],
loss_type='MSE',
set_transforms=False)
in_shift, in_scale, out_shift, out_scale = bc_agent.compute_transformations(
)
bc_agent.set_transformations(in_shift, in_scale, out_shift, out_scale)
bc_agent.set_variance_with_data(out_scale)
ts = timer.time()
print("========================================")
print("Running BC with expert demonstrations")
print("========================================")
bc_agent.train()
print("========================================")
print("BC training complete !!!")
print("time taken = %f" % (timer.time() - ts))
print("========================================")
if job_data['eval_rollouts'] >= 1:
score = e.evaluate_policy(policy,
num_episodes=job_data['eval_rollouts'],
mean_action=True)
print("Score with behavior cloning = %f" % score[0][0])
if job_data['algorithm'] != 'DAPG':
# We throw away the demo data when training from scratch or fine-tuning with RL without explicit augmentation
demo_paths = None
# ===============================================================================
# RL Loop
# ===============================================================================
rl_agent = DAPG(e,
policy,
baseline,
demo_paths,
normalized_step_size=job_data['rl_step_size'],
lam_0=job_data['lam_0'],
lam_1=job_data['lam_1'],
seed=seed,
save_logs=True)
print("========================================")
print("Starting reinforcement learning phase")
print("========================================")
ts = timer.time()
train_agent(job_name=JOB_DIR,
agent=rl_agent,
seed=seed,
niter=job_data['rl_num_iter'],
gamma=job_data['rl_gamma'],
gae_lambda=job_data['rl_gae'],
num_cpu=job_data['num_cpu'],
sample_mode='trajectories',
num_traj=job_data['rl_num_traj'],
save_freq=job_data['save_freq'],
evaluation_rollouts=job_data['eval_rollouts'])
print("time taken = %f" % (timer.time() - ts))
