diff --git a/tests/test_save_load.py b/tests/test_save_load.py index dec98b4..f3f8d10 100644 --- a/tests/test_save_load.py +++ b/tests/test_save_load.py @@ -10,6 +10,7 @@ from torchy_baselines.common.vec_env import DummyVecEnv from torchy_baselines.common.identity_env import IdentityEnvBox MODEL_LIST = [ + CEMRL, PPO, A2C, TD3, @@ -37,7 +38,6 @@ def test_save_load(model_class): observations = np.array([env.step(env.action_space.sample())[0] for _ in range(10)]) observations = np.squeeze(observations) - # Get dictionary of current parameters params = deepcopy(model.get_policy_parameters()) opt_params = deepcopy(model.get_opt_parameters()) @@ -55,8 +55,7 @@ def test_save_load(model_class): params = new_params - - #get selected actions + # get selected actions selected_actions = [model.predict(observation, deterministic=True) for observation in observations] # Check diff --git a/torchy_baselines/cem_rl/cem_rl.py b/torchy_baselines/cem_rl/cem_rl.py index f35fa43..3305cb5 100644 --- a/torchy_baselines/cem_rl/cem_rl.py +++ b/torchy_baselines/cem_rl/cem_rl.py @@ -78,7 +78,8 @@ class CEMRL(TD3): # set params self.actor.load_from_vector(self.es_params[i]) self.actor_target.load_from_vector(self.es_params[i]) - self.actor.optimizer = th.optim.Adam(self.actor.parameters(), lr=self.learning_rate(self._current_progress)) + self.actor.optimizer = th.optim.Adam(self.actor.parameters(), + lr=self.learning_rate(self._current_progress)) # In the paper: 2 * actor_steps // self.n_grad # In the original implementation: actor_steps // self.n_grad diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py index 94c1cc8..8c363c9 100644 --- a/torchy_baselines/common/base_class.py +++ b/torchy_baselines/common/base_class.py @@ -292,8 +292,6 @@ class BaseRLModel(object): model.__dict__.update(kwargs) model.set_env(env) model.load_parameters(params, opt_params) - # resetup modul after load - # model._setup_model() return model @staticmethod @@ -519,14 +517,6 @@ class BaseRLModel(object): """ return ["replay_buffer"] - def _resetup_model(self): - """ - Function will be called at the end of load and should resetup anything that might not have been saved - warning: this function should always be in compliance with excluded_save_params - :return: - """ - pass - def save(self, path, include=None): """ saves all the params from init and pytorch params in a file for continuous learning @@ -548,4 +538,4 @@ class BaseRLModel(object): params_to_save = self.get_policy_parameters() opt_params_to_save = self.get_opt_parameters() - self._save_to_file_zip(path, data=data, params=params_to_save, opt_params=opt_params_to_save) \ No newline at end of file + self._save_to_file_zip(path, data=data, params=params_to_save, opt_params=opt_params_to_save) diff --git a/torchy_baselines/ppo/ppo.py b/torchy_baselines/ppo/ppo.py index a0c2231..2938c12 100644 --- a/torchy_baselines/ppo/ppo.py +++ b/torchy_baselines/ppo/ppo.py @@ -120,8 +120,8 @@ class PPO(BaseRLModel): self.rollout_buffer = RolloutBuffer(self.n_steps, state_dim, action_dim, self.device, gamma=self.gamma, gae_lambda=self.gae_lambda, n_envs=self.n_envs) self.policy = self.policy_class(self.observation_space, self.action_space, - self.learning_rate, use_sde=self.use_sde, device=self.device, - **self.policy_kwargs) + self.learning_rate, use_sde=self.use_sde, device=self.device, + **self.policy_kwargs) self.policy = self.policy.to(self.device) self.clip_range = get_schedule_fn(self.clip_range) @@ -227,7 +227,6 @@ class PPO(BaseRLModel): # Value loss using the TD(gae_lambda) target value_loss = F.mse_loss(return_batch, values_pred) - # Entropy loss favor exploration entropy_loss = -th.mean(entropy) diff --git a/torchy_baselines/sac/sac.py b/torchy_baselines/sac/sac.py index cc539d5..06e107f 100644 --- a/torchy_baselines/sac/sac.py +++ b/torchy_baselines/sac/sac.py @@ -52,6 +52,7 @@ class SAC(BaseRLModel): Setting it to auto, the code will be run on the GPU if possible. :param _init_setup_model: (bool) Whether or not to build the network at the creation of the instance """ + def __init__(self, policy, env, learning_rate=3e-4, buffer_size=int(1e6), learning_starts=100, batch_size=64, tau=0.005, ent_coef='auto', target_update_interval=1, @@ -124,7 +125,7 @@ class SAC(BaseRLModel): self.replay_buffer = ReplayBuffer(self.buffer_size, obs_dim, action_dim, self.device) self.policy = self.policy_class(self.observation_space, self.action_space, - self.learning_rate, device=self.device, **self.policy_kwargs) + self.learning_rate, device=self.device, **self.policy_kwargs) self.policy = self.policy.to(self.device) self._create_aliases() @@ -281,7 +282,8 @@ class SAC(BaseRLModel): :return: (Dict) of optimizer names and their state_dict """ if self.ent_coef_optimizer is not None: - return {"actor": self.actor.optimizer.state_dict(), "critic": self.critic.optimizer.state_dict(),"ent_coef_optimizer": self.ent_coef_optimizer.state_dict()} + return {"actor": self.actor.optimizer.state_dict(), "critic": self.critic.optimizer.state_dict(), + "ent_coef_optimizer": self.ent_coef_optimizer.state_dict()} else: return {"actor": self.actor.optimizer.state_dict(), "critic": self.critic.optimizer.state_dict()} diff --git a/torchy_baselines/td3/td3.py b/torchy_baselines/td3/td3.py index 1c3639e..5a847ac 100644 --- a/torchy_baselines/td3/td3.py +++ b/torchy_baselines/td3/td3.py @@ -81,19 +81,10 @@ class TD3(BaseRLModel): self.set_random_seed(self.seed) self.replay_buffer = ReplayBuffer(self.buffer_size, obs_dim, action_dim, self.device) self.policy = self.policy_class(self.observation_space, self.action_space, - self.learning_rate, device=self.device, **self.policy_kwargs) + self.learning_rate, device=self.device, **self.policy_kwargs) self.policy = self.policy.to(self.device) self._create_aliases() - def _resetup_model(self): - """ - method used to resetup anything that was not saved - :return: - """ - if self.replay_buffer is None: - obs_dim, action_dim = self.observation_space.shape[0], self.action_space.shape[0] - self.replay_buffer = ReplayBuffer(self.buffer_size, obs_dim, action_dim, self.device) - def _create_aliases(self): self.actor = self.policy.actor self.actor_target = self.policy.actor_target @@ -265,4 +256,4 @@ class TD3(BaseRLModel): """ self.actor.optimizer.load_state_dict(opt_params["actor"]) self.critic.optimizer.load_state_dict(opt_params["critic"]) - self.policy.load_state_dict(load_dict) \ No newline at end of file + self.policy.load_state_dict(load_dict)