From 2dfd7e3c2b772ed0dc6c8e74119b31bd96099d74 Mon Sep 17 00:00:00 2001 From: Arec Jamgochian Date: Fri, 25 Feb 2022 17:00:46 -0800 Subject: [PATCH] add policy saving --- sgail-ppo-options-setobs2.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/sgail-ppo-options-setobs2.py b/sgail-ppo-options-setobs2.py index 5d7f569..db15496 100644 --- a/sgail-ppo-options-setobs2.py +++ b/sgail-ppo-options-setobs2.py @@ -92,6 +92,11 @@ def training_function(config): tune.report(gen_mean_reward_per_episode=info['gen/mean_reward_per_episode'], disc_mean_reward_per_episode=info['disc/mean_reward_per_episode'], 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( env_fn=env_fn, @@ -116,7 +121,8 @@ def training_function(config): lr_schedulers=[pi_lr_scheduler], ) - # save value, policy + # save model + torch.save(policy.state_dict(), 'policy_final.pt') analysis = tune.run(