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

Type A2C and PPO init

Browse source
This commit is contained in:
Antonin Raffin 2020-03-10 18:17:47 +01:00
parent 35d0d2b320
commit 7e3736ed56
2 changed files with 53 additions and 18 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

35
torchy_baselines/a2c/a2c.py
View file

@ -1,10 +1,15 @@
from typing import List, Tuple, Type, Union, Callable, Optional, Dict, Any
from gym import spaces
import torch as th
import torch.nn.functional as F
from torchy_baselines.common.utils import explained_variance
from torchy_baselines.ppo.ppo import PPO
from torchy_baselines.common import logger
from torchy_baselines.common.type_aliases import GymEnv
from torchy_baselines.ppo.ppo import PPO
from torchy_baselines.ppo.policies import PPOPolicy
class A2C(PPO):
@ -46,13 +51,27 @@ class A2C(PPO):
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=7e-4,
n_steps=5, gamma=0.99, gae_lambda=1.0,
ent_coef=0.0, vf_coef=0.5, max_grad_norm=0.5,
rms_prop_eps=1e-5, use_rms_prop=True, use_sde=False, sde_sample_freq=-1,
normalize_advantage=False, tensorboard_log=None, create_eval_env=False,
policy_kwargs=None, verbose=0, seed=None, device='auto',
_init_setup_model=True):
def __init__(self, policy: Union[str, Type[PPOPolicy]],
env: Union[GymEnv, str],
learning_rate: Union[float, Callable] = 7e-4,
n_steps: int = 5,
gamma: float = 0.99,
gae_lambda: float = 1.0,
ent_coef: float = 0.0,
vf_coef: float = 0.5,
max_grad_norm: float = 0.5,
rms_prop_eps: float = 1e-5,
use_rms_prop: bool = True,
use_sde: bool = False,
sde_sample_freq: int = -1,
normalize_advantage: bool = False,
tensorboard_log: Optional[str] = None,
create_eval_env: bool = False,
policy_kwargs: Optional[Dict[str, Any]] = None,
verbose: int = 0,
seed: Optional[int] = None,
device: Union[th.device, str] = 'auto',
_init_setup_model: bool = True):
super(A2C, self).__init__(policy, env, learning_rate=learning_rate,
n_steps=n_steps, batch_size=None, n_epochs=1,

36
torchy_baselines/ppo/ppo.py
View file

@ -1,6 +1,6 @@
import os
import time
from typing import Optional, Tuple, List
from typing import List, Tuple, Type, Union, Callable, Optional, Dict, Any
import gym
from gym import spaces
@ -14,12 +14,13 @@ except ImportError:
SummaryWriter = None
import numpy as np
from torchy_baselines.common import logger
from torchy_baselines.common.base_class import BaseRLModel
from torchy_baselines.common.type_aliases import GymEnv
from torchy_baselines.common.buffers import RolloutBuffer
from torchy_baselines.common.utils import explained_variance, get_schedule_fn
from torchy_baselines.common.vec_env import VecEnv
from torchy_baselines.common.callbacks import BaseCallback
from torchy_baselines.common import logger
from torchy_baselines.ppo.policies import PPOPolicy
@ -73,14 +74,29 @@ class PPO(BaseRLModel):
: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,
n_steps=2048, batch_size=64, n_epochs=10,
gamma=0.99, gae_lambda=0.95, clip_range=0.2, clip_range_vf=None,
ent_coef=0.0, vf_coef=0.5, max_grad_norm=0.5,
use_sde=False, sde_sample_freq=-1,
target_kl=None, tensorboard_log=None, create_eval_env=False,
policy_kwargs=None, verbose=0, seed=None, device='auto',
_init_setup_model=True):
def __init__(self, policy: Union[str, Type[PPOPolicy]],
env: Union[GymEnv, str],
learning_rate: Union[float, Callable] = 3e-4,
n_steps: int = 2048,
batch_size: Optional[int] = 64,
n_epochs: int = 10,
gamma: float = 0.99,
gae_lambda: float = 0.95,
clip_range: float = 0.2,
clip_range_vf: Optional[float] = None,
ent_coef: float = 0.0,
vf_coef: float = 0.5,
max_grad_norm: float = 0.5,
use_sde: bool = False,
sde_sample_freq: int = -1,
target_kl: Optional[float] = None,
tensorboard_log: Optional[str] = None,
create_eval_env: bool = False,
policy_kwargs: Optional[Dict[str, Any]] = None,
verbose: int = 0,
seed: Optional[int] = None,
device: Union[th.device, str] = 'auto',
_init_setup_model: bool = True):
super(PPO, self).__init__(policy, env, PPOPolicy, policy_kwargs=policy_kwargs,
verbose=verbose, device=device, use_sde=use_sde, sde_sample_freq=sde_sample_freq,