huangfu/InteractionImitation
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add policy saving

Browse Source
This commit is contained in:
Arec Jamgochian
2022-02-25 17:00:46 -08:00
parent 59083ebce3
commit 2dfd7e3c2b
1 changed files with 7 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
sgail-ppo-options-setobs2.py
View File

@@ -93,6 +93,11 @@ def training_function(config):
disc_mean_reward_per_episode=info['disc/mean_reward_per_episode'], disc_mean_reward_per_episode=info['disc/mean_reward_per_episode'],
mean_episode_length=info['gen/mean_episode_length']) mean_episode_length=info['gen/mean_episode_length'])
# save model checkpoints
ep = info['epoch'] + 1
if (ep % 25 == 0):
torch.save(info['policy'].state_dict(), f'policy_epoch{ep}.pt')
value, policy = gail_ppo( value, policy = gail_ppo(
env_fn=env_fn, env_fn=env_fn,
expert_data=expert_data, expert_data=expert_data,
@@ -116,7 +121,8 @@ def training_function(config):
lr_schedulers=[pi_lr_scheduler], lr_schedulers=[pi_lr_scheduler],
) )
# save value, policy # save model
torch.save(policy.state_dict(), 'policy_final.pt')
analysis = tune.run( analysis = tune.run(