saymrwulf/stable-baselines3
1
0
Fork 0
mirror of https://github.com/saymrwulf/stable-baselines3.git synced 2026-07-12 17:58:00 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Reformat

Browse source
This commit is contained in:
Antonin RAFFIN 2019-09-21 17:17:09 +02:00
parent 3ececcd3a9
commit 2469ff3859
11 changed files with 22 additions and 23 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
tests/test_run.py
View file

@ -1,9 +1,8 @@
import os
import gym
from torchy_baselines import TD3, CEMRL, PPO
def test_td3():
model = TD3('MlpPolicy', 'Pendulum-v0', policy_kwargs=dict(net_arch=[64, 64]),
start_timesteps=100, verbose=1, create_eval_env=True)
@ -12,14 +11,16 @@ def test_td3():
model.load("test_save")
os.remove("test_save.pth")
def test_cemrl():
model = CEMRL('MlpPolicy', 'Pendulum-v0', policy_kwargs=dict(net_arch=[16]), pop_size=2, n_grad=1,
start_timesteps=100, verbose=1, create_eval_env=True)
start_timesteps=100, verbose=1, create_eval_env=True)
model.learn(total_timesteps=20000, eval_freq=1000)
model.save("test_save")
model.load("test_save")
os.remove("test_save.pth")
def test_ppo():
model = PPO('MlpPolicy', 'Pendulum-v0', policy_kwargs=dict(net_arch=[16]), verbose=1, create_eval_env=True)
model.learn(total_timesteps=1000, eval_freq=500)

4
tests/test_vec_normalize.py
View file

@ -11,8 +11,8 @@ ENV_ID = 'Pendulum-v0'
def test_runningmeanstd():
"""Test RunningMeanStd object"""
for (x_1, x_2, x_3) in [
(np.random.randn(3), np.random.randn(4), np.random.randn(5)),
(np.random.randn(3, 2), np.random.randn(4, 2), np.random.randn(5, 2))]:
(np.random.randn(3), np.random.randn(4), np.random.randn(5)),
(np.random.randn(3, 2), np.random.randn(4, 2), np.random.randn(5, 2))]:
rms = RunningMeanStd(epsilon=0.0, shape=x_1.shape[1:])
x_cat = np.concatenate([x_1, x_2, x_3], axis=0)

10
torchy_baselines/cem_rl/cem_rl.py
View file

@ -1,12 +1,10 @@
import time
import torch as th
import torch.nn.functional as F
import numpy as np
from torchy_baselines.td3.td3 import TD3
from torchy_baselines.common.evaluation import evaluate_policy
from torchy_baselines.cem_rl.cem import CEM
from torchy_baselines.common.evaluation import evaluate_policy
from torchy_baselines.td3.td3 import TD3
class CEMRL(TD3):
@ -16,6 +14,7 @@ class CEMRL(TD3):
Paper: https://arxiv.org/abs/1810.01222
Code: https://github.com/apourchot/CEM-RL
"""
def __init__(self, policy, env, policy_kwargs=None, verbose=0,
sigma_init=1e-3, pop_size=10, damp=1e-3, damp_limit=1e-5,
elitism=False, n_grad=5, policy_freq=2, batch_size=100,
@ -100,7 +99,7 @@ class CEMRL(TD3):
else:
# scales with a bigger population
# but less training steps per agent
n_training_steps == 2 * (actor_steps // self.n_grad)
n_training_steps = 2 * (actor_steps // self.n_grad)
for it in range(n_training_steps):
# Sample replay buffer
replay_data = self.replay_buffer.sample(self.batch_size)
@ -148,7 +147,6 @@ class CEMRL(TD3):
print("Total T: {} Episode Num: {} Episode T: {} Reward: {}".format(
self.num_timesteps, episode_num, episode_timesteps, episode_reward))
self.es.tell(self.es_params, self.fitnesses)
timesteps_since_eval += actor_steps
return self

2
torchy_baselines/common/base_class.py
View file

@ -1,6 +1,5 @@
from abc import ABCMeta, abstractmethod
import gym
import torch as th
import numpy as np
@ -216,7 +215,6 @@ class BaseRLModel(object):
if self.eval_env is not None:
self.eval_env.seed(seed)
def collect_rollouts(self, env, n_episodes=1, action_noise_std=0.0,
deterministic=False, callback=None, remove_timelimits=True,
start_timesteps=0, num_timesteps=0, replay_buffer=None):

4
torchy_baselines/common/buffers.py
View file

@ -52,10 +52,12 @@ class BaseBuffer(object):
def _get_samples(self, batch_inds):
raise NotImplementedError()
class ReplayBuffer(BaseBuffer):
"""
Taken from https://github.com/apourchot/CEM-RL
"""
def __init__(self, buffer_size, state_dim, action_dim, device='cpu', n_envs=1):
super(ReplayBuffer, self).__init__(buffer_size, state_dim, action_dim, device, n_envs=n_envs)
@ -89,7 +91,7 @@ class ReplayBuffer(BaseBuffer):
class RolloutBuffer(BaseBuffer):
def __init__(self, buffer_size, state_dim, action_dim, device='cpu',
lambda_=1, gamma=0.99, n_envs=1):
lambda_=1, gamma=0.99, n_envs=1):
super(RolloutBuffer, self).__init__(buffer_size, state_dim, action_dim, device, n_envs=n_envs)
# TODO: try the buffer on the gpu?
self.lambda_ = lambda_

1
torchy_baselines/common/distributions.py
View file

@ -1,6 +1,7 @@
import torch as th
from torch.distributions import Normal
class Distribution(object):
def __init__(self):
super(Distribution, self).__init__()

3
torchy_baselines/common/utils.py
View file

@ -1,8 +1,7 @@
import random
import scipy.signal
import torch as th
import numpy as np
import torch as th
def set_random_seed(seed, using_cuda=False):

4
torchy_baselines/common/vec_env/__init__.py
View file

@ -1,6 +1,6 @@
# flake8: noqa F401
from torchy_baselines.common.vec_env.base_vec_env import AlreadySteppingError, NotSteppingError, VecEnv, VecEnvWrapper, \
CloudpickleWrapper
from torchy_baselines.common.vec_env.base_vec_env import AlreadySteppingError, NotSteppingError,\
VecEnv, VecEnvWrapper, CloudpickleWrapper
from torchy_baselines.common.vec_env.dummy_vec_env import DummyVecEnv
from torchy_baselines.common.vec_env.subproc_vec_env import SubprocVecEnv
from torchy_baselines.common.vec_env.vec_frame_stack import VecFrameStack

2
torchy_baselines/ppo/policies.py
View file

@ -44,7 +44,6 @@ class PPOPolicy(BasePolicy):
self.log_std = nn.Parameter(th.zeros(self.action_dim))
# Init weights: use orthogonal initialization
for module in [self.shared_net, self.actor_net, self.value_net]:
gain = 0.01 if module == self.actor_net else 1.0
# Values from stable-baselines check why
gain = {
self.shared_net: np.sqrt(2),
@ -98,6 +97,7 @@ class PPOPolicy(BasePolicy):
def value_forward(self):
pass
MlpPolicy = PPOPolicy
register_policy("MlpPolicy", MlpPolicy)

1
torchy_baselines/td3/policies.py
View file

@ -79,6 +79,7 @@ class TD3Policy(BasePolicy):
def make_critic(self):
return Critic(**self.net_args).to(self.device)
MlpPolicy = TD3Policy
register_policy("MlpPolicy", MlpPolicy)

7
torchy_baselines/td3/td3.py
View file

@ -21,7 +21,7 @@ class TD3(BaseRLModel):
buffer_size=int(1e6), learning_rate=1e-3, seed=0, device='auto',
action_noise_std=0.1, start_timesteps=100, policy_freq=2,
batch_size=100, create_eval_env=False,
_init_setup_model=True):
_init_setup_model=True):
super(TD3, self).__init__(policy, env, TD3Policy, policy_kwargs, verbose, device,
create_eval_env=create_eval_env)
@ -135,8 +135,7 @@ class TD3(BaseRLModel):
for param, target_param in zip(self.actor.parameters(), self.actor_target.parameters()):
target_param.data.copy_(tau_actor * param.data + (1 - tau_actor) * target_param.data)
def train(self, n_iterations, batch_size=100, discount=0.99,
tau=0.005, policy_noise=0.2, noise_clip=0.5, policy_freq=2):
def train(self, n_iterations, batch_size=100, policy_freq=2):
for it in range(n_iterations):
@ -177,7 +176,7 @@ class TD3(BaseRLModel):
if self.num_timesteps > 0:
if self.verbose > 1:
print("Total T: {} Episode Num: {} Episode T: {} Reward: {}".format(
self.num_timesteps, episode_num, episode_timesteps, episode_reward))
self.num_timesteps, episode_num, episode_timesteps, episode_reward))
self.train(episode_timesteps, batch_size=self.batch_size, policy_freq=self.policy_freq)
# Evaluate episode