saymrwulf/stable-baselines3
1
0
Fork 0
mirror of https://github.com/saymrwulf/stable-baselines3.git synced 2026-09-16 22:20:26 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Formatted all files

Browse source
This commit is contained in:
Noah Dormann 2019-11-28 15:38:04 +01:00
parent 751ccf85e7
commit e95858784a
6 changed files with 13 additions and 31 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
tests/test_save_load.py
View file

@ -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

3
torchy_baselines/cem_rl/cem_rl.py
View file

@ -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

12
torchy_baselines/common/base_class.py
View file

@ -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)
self._save_to_file_zip(path, data=data, params=params_to_save, opt_params=opt_params_to_save)

5
torchy_baselines/ppo/ppo.py
View file

@ -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)

6
torchy_baselines/sac/sac.py
View file

@ -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()}

13
torchy_baselines/td3/td3.py
View file

@ -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)
self.policy.load_state_dict(load_dict)