From 5404c759a651ddda0ac57b8a55651f7ec8aaff8c Mon Sep 17 00:00:00 2001 From: MishaLaskin Date: Fri, 21 Feb 2020 07:44:24 -0800 Subject: [PATCH] edits --- curl_sac.py | 5 +++++ dmc2gym | 1 + train.py | 2 +- 3 files changed, 7 insertions(+), 1 deletion(-) create mode 160000 dmc2gym diff --git a/curl_sac.py b/curl_sac.py index 7c16099..8dfc87a 100644 --- a/curl_sac.py +++ b/curl_sac.py @@ -482,6 +482,11 @@ class CurlSacAgent(object): self.critic.state_dict(), '%s/critic_%s.pt' % (model_dir, step) ) + def save_curl(self, model_dir, step): + torch.save( + self.CURL.state_dict(), '%s/curl_%s.pt' % (model_dir, step) + ) + def load(self, model_dir, step): self.actor.load_state_dict( torch.load('%s/actor_%s.pt' % (model_dir, step)) diff --git a/dmc2gym b/dmc2gym new file mode 160000 index 0000000..da8b463 --- /dev/null +++ b/dmc2gym @@ -0,0 +1 @@ +Subproject commit da8b46354ef670d9b82848bd7d1946c9e94dbd2b diff --git a/train.py b/train.py index bbf1f6b..294dfa4 100644 --- a/train.py +++ b/train.py @@ -230,7 +230,7 @@ def main(): L.log('eval/episode', episode, step) evaluate(env, agent, video, args.num_eval_episodes, L, step,args) if args.save_model: - agent.save(model_dir, step) + agent.save_curl(model_dir, step) if args.save_buffer: replay_buffer.save(buffer_dir)