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

explanations

Browse source
This commit is contained in:
Stelios Tymvios 2020-07-04 15:22:13 +03:00
parent 376d3c9295
commit 1709ebea79
7 changed files with 61 additions and 68 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
docs/misc/changelog.rst
View file

@ -19,6 +19,7 @@ New Features:
when ``psutil`` is available
- Saving models now automatically creates the necessary folders and raises appropriate warnings (@PartiallyTyped)
- Refactored opening paths for saving and loading to use strings, pathlib or io.BufferedIOBase (@PartiallyTyped)
- Implemented N-Step ReplayBuffer for off-policy methods (@PartiallyTyped)
Bug Fixes:
^^^^^^^^^^

34
stable_baselines3/common/buffers.py
View file

@ -380,7 +380,7 @@ class RolloutBuffer(BaseBuffer):
start_idx = 0
while start_idx < self.buffer_size * self.n_envs:
yield self._get_samples(indices[start_idx : start_idx + batch_size])
yield self._get_samples(indices[start_idx: start_idx + batch_size])
start_idx += batch_size
def _get_samples(self, batch_inds: np.ndarray, env: Optional[VecNormalize] = None) -> RolloutBufferSamples:
@ -414,34 +414,37 @@ class NstepReplayBuffer(ReplayBuffer):
assert 0 <= gamma <= 1
def _get_samples(self, batch_inds: np.ndarray, env: Optional[VecNormalize] = None) -> ReplayBufferSamples:
assert not np.any(batch_inds==self.pos)
assert not np.any(batch_inds == self.pos)
actions = self.actions[batch_inds, 0, :]
gamma = self.gamma
# two dimensional matrix contains all indices
# Broadcasting turns 1dim arange matrix to 2 dimensional matrix that contains all
# the indices, % buffersize keeps us in buffer range
indices = (np.arange(self.n_step) + batch_inds.reshape(-1, 1)) % self.buffer_size
# two dim matrix of not dones. If done then subsequent dones are 0
not_dones = np.multiply.accumulate(1-self.dones[indices], 1).reshape(-1, self.n_step)
# two dim matrix of gammas to the power, will be used with the rewards
gammas = gamma**np.arange(self.n_step)
# two dim matrix of rewards
# two dim matrix of not dones. If done is true, then subsequent dones are turned to 0
# using accumulate. This ensures that we don't use invalid transitions
not_dones = np.multiply.accumulate(1 - self.dones[indices], 1).reshape(-1, self.n_step)
# vector of the discount factors
gammas = gamma ** np.arange(self.n_step)
# two dim matrix of rewards for the indices
rewards = self.rewards[indices].reshape(not_dones.shape)
# filter to ignore loops around the buffer, see:
# https://github.com/DLR-RM/stable-baselines3/pull/81#issuecomment-653741592
# basically we need to ignore indices that are equal to self.pos
# because the indice at self.pos is older
# not_dones that are from other transitions are filtered out.
filt = not_dones * (indices != self.pos)
filt = np.multiply.accumulate(filt, 1).reshape(filt.shape)
# weighted rewards
rewards = filt* gammas * rewards
rewards = filt * gammas * rewards
rewards = np.add.reduce(rewards, 1)
next_obs_indices = np.add.reduce(filt, 1).astype(np.int).reshape(batch_inds.shape)
next_obs_indices = (next_obs_indices + batch_inds - 1) % self.buffer_size
obs = self._normalize_obs(self.observations[batch_inds, 0, :], env)
@ -450,6 +453,7 @@ class NstepReplayBuffer(ReplayBuffer):
else:
next_obs = self._normalize_obs(self.next_observations[next_obs_indices, 0, :], env)
dones = 1- (not_dones[np.arange(len(batch_inds)), next_obs_indices]).reshape(len(batch_inds), 1)
dones = 1 - (not_dones[np.arange(len(batch_inds)), next_obs_indices]).reshape(len(batch_inds), 1)
data = (obs, actions, next_obs, dones.reshape(len(batch_inds), 1), rewards.reshape(len(batch_inds), 1))
return ReplayBufferSamples(*tuple(map(self.to_torch, data)))

43
stable_baselines3/common/off_policy_algorithm.py
View file

@ -2,7 +2,7 @@ import time
import warnings
import pathlib
import io
from typing import Union, Type, Optional, Dict, Any, Callable, List, Tuple, Mapping
from typing import Union, Type, Optional, Dict, Any, Callable, List, Tuple, Mapping, Generic, TypeVar
import gym
import torch as th
@ -19,8 +19,9 @@ from stable_baselines3.common.noise import ActionNoise
from stable_baselines3.common.buffers import ReplayBuffer
from stable_baselines3.common.save_util import save_to_pkl, load_from_pkl
_BufferType = TypeVar("_BufferType", bound=ReplayBuffer)
class OffPolicyAlgorithm(BaseAlgorithm):
class OffPolicyAlgorithm(BaseAlgorithm, Generic[_BufferType]):
"""
The base for Off-Policy algorithms (ex: SAC/TD3)
@ -73,6 +74,8 @@ class OffPolicyAlgorithm(BaseAlgorithm):
env: Union[GymEnv, str],
policy_base: Type[BasePolicy],
learning_rate: Union[float, Callable],
replay_buffer_cls: Type[_BufferType] = None,
replay_buffer_kwargs: Optional[Mapping[str, Any]] = None,
buffer_size: int = int(1e6),
learning_starts: int = 100,
batch_size: int = 256,
@ -95,8 +98,6 @@ class OffPolicyAlgorithm(BaseAlgorithm):
sde_sample_freq: int = -1,
use_sde_at_warmup: bool = False,
sde_support: bool = True,
replay_buffer_cls: Type[ReplayBuffer] = ReplayBuffer,
replay_buffer_kwargs: Mapping[str, Any] = None,
):
super(OffPolicyAlgorithm, self).__init__(
@ -137,7 +138,6 @@ class OffPolicyAlgorithm(BaseAlgorithm):
)
self.actor = None # type: Optional[th.nn.Module]
self.replay_buffer = None # type: Optional[ReplayBuffer]
# Update policy keyword arguments
if sde_support:
self.policy_kwargs["use_sde"] = self.use_sde
@ -145,14 +145,14 @@ class OffPolicyAlgorithm(BaseAlgorithm):
# For gSDE only
self.use_sde_at_warmup = use_sde_at_warmup
assert issubclass(replay_buffer_cls, ReplayBuffer)
self.replay_buffer_cls = replay_buffer_cls
self.replay_buffer_kwargs = replay_buffer_kwargs or {}
self.replay_buffer = None
self.replay_buffer_cls = replay_buffer_cls if replay_buffer_cls is not None else ReplayBuffer
self.replay_buffer_kwargs = dict(replay_buffer_kwargs or {})
def _setup_model(self) -> None:
self._setup_lr_schedule()
self.set_random_seed(self.seed)
self.replay_buffer = ReplayBuffer(
self.replay_buffer = self.replay_buffer_cls(
self.buffer_size,
self.observation_space,
self.action_space,
@ -183,18 +183,19 @@ class OffPolicyAlgorithm(BaseAlgorithm):
:param path: (Union[str, pathlib.Path, io.BufferedIOBase]) Path to the pickled replay buffer.
"""
self.replay_buffer = load_from_pkl(path, self.verbose)
assert isinstance(self.replay_buffer, ReplayBuffer), 'The replay buffer must inherit from ReplayBuffer class'
assert isinstance(self.replay_buffer, ReplayBuffer), "The replay buffer must inherit from ReplayBuffer class"
def _setup_learn(self,
total_timesteps: int,
eval_env: Optional[GymEnv],
callback: Union[None, Callable, List[BaseCallback], BaseCallback] = None,
eval_freq: int = 10000,
n_eval_episodes: int = 5,
log_path: Optional[str] = None,
reset_num_timesteps: bool = True,
tb_log_name: str = 'run',
) -> Tuple[int, BaseCallback]:
def _setup_learn(
self,
total_timesteps: int,
eval_env: Optional[GymEnv],
callback: Union[None, Callable, List[BaseCallback], BaseCallback] = None,
eval_freq: int = 10000,
n_eval_episodes: int = 5,
log_path: Optional[str] = None,
reset_num_timesteps: bool = True,
tb_log_name: str = "run",
) -> Tuple[int, BaseCallback]:
"""
cf `BaseAlgorithm`.
"""
@ -207,7 +208,7 @@ class OffPolicyAlgorithm(BaseAlgorithm):
and self.replay_buffer is not None
and (self.replay_buffer.full or self.replay_buffer.pos > 0)
)
assert self.replay_buffer is not None
if truncate_last_traj:
warnings.warn(
"The last trajectory in the replay buffer will be truncated, "

4
stable_baselines3/dqn/dqn.py
View file

@ -9,7 +9,7 @@ from stable_baselines3.common.off_policy_algorithm import OffPolicyAlgorithm
from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback
from stable_baselines3.common.utils import get_linear_fn
from stable_baselines3.dqn.policies import DQNPolicy
from stable_baselines3.common.buffers import ReplayBuffer
class DQN(OffPolicyAlgorithm):
"""
@ -78,7 +78,7 @@ class DQN(OffPolicyAlgorithm):
_init_setup_model: bool = True):
super(DQN, self).__init__(policy, env, DQNPolicy, learning_rate,
buffer_size, learning_starts, batch_size,
ReplayBuffer, None, buffer_size, learning_starts, batch_size,
tau, gamma, train_freq, gradient_steps,
n_episodes_rollout, action_noise=None, # No action noise
policy_kwargs=policy_kwargs,

2
stable_baselines3/sac/sac.py
View file

@ -6,6 +6,7 @@ import numpy as np
from stable_baselines3.common import logger
from stable_baselines3.common.off_policy_algorithm import OffPolicyAlgorithm
from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback
from stable_baselines3.common.buffers import ReplayBuffer
from stable_baselines3.common.noise import ActionNoise
from stable_baselines3.sac.policies import SACPolicy
@ -93,6 +94,7 @@ class SAC(OffPolicyAlgorithm):
_init_setup_model: bool = True):
super(SAC, self).__init__(policy, env, SACPolicy, learning_rate,
ReplayBuffer, None,
buffer_size, learning_starts, batch_size,
tau, gamma, train_freq, gradient_steps,
n_episodes_rollout, action_noise,

2
stable_baselines3/td3/td3.py
View file

@ -6,6 +6,7 @@ from stable_baselines3.common import logger
from stable_baselines3.common.off_policy_algorithm import OffPolicyAlgorithm
from stable_baselines3.common.noise import ActionNoise
from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback
from stable_baselines3.common.buffers import ReplayBuffer
from stable_baselines3.td3.policies import TD3Policy
@ -77,6 +78,7 @@ class TD3(OffPolicyAlgorithm):
_init_setup_model: bool = True):
super(TD3, self).__init__(policy, env, TD3Policy, learning_rate,
ReplayBuffer, None,
buffer_size, learning_starts, batch_size,
tau, gamma, train_freq, gradient_steps,
n_episodes_rollout, action_noise=action_noise,

43
tests/test_buffers.py
View file

@ -3,29 +3,24 @@ import pytest
import numpy as np
from gym import spaces
from stable_baselines3.common.buffers import NstepReplayBuffer, ReplayBuffer
from stable_baselines3.common.buffers import NstepReplayBuffer
def test_nsteps():
buffer = NstepReplayBuffer(
5,
spaces.Discrete(5),
spaces.Discrete(5),
n_step=5,
gamma=1
)
buffer = NstepReplayBuffer(5, spaces.Discrete(5), spaces.Discrete(5), n_step=5, gamma=1)
buffer.add(0, 1, 10, 1, 0)
buffer.add(1, 2, 11, 1, 0)
buffer.add(2, 3, 12, 1, 0)
buffer.add(3, 4, 13, 1, 0)
buffer.add(4, 5, 14, 1, 0)
obs, act, next_obs, dones, rewards = buffer._get_samples(np.array([1,2,3,4]))
obs, act, next_obs, dones, rewards = buffer._get_samples(np.array([1, 2, 3, 4]))
assert obs.shape == (4, 1)
assert act.shape == (4, 1)
assert next_obs.shape == (4, 1)
assert dones.shape == (4, 1)
assert rewards.shape == (4, 1)
assert np.allclose(dones, np.zeros_like(dones))
assert np.allclose(next_obs, np.array([[5],[5],[5],[5]]))
assert np.allclose(next_obs, np.array([[5], [5], [5], [5]]))
assert np.allclose(rewards, np.array([4, 3, 2, 1]).reshape(4, 1))
assert np.allclose(act, np.array([11, 12, 13, 14]).reshape(4, 1))
@ -33,50 +28,38 @@ def test_nsteps():
with pytest.raises(AssertionError):
buffer._get_samples(np.array([0, 1, 2, 3]))
buffer = NstepReplayBuffer(
5,
spaces.Discrete(5),
spaces.Discrete(5),
n_step=5,
gamma=0.9
)
buffer = NstepReplayBuffer(5, spaces.Discrete(5), spaces.Discrete(5), n_step=5, gamma=0.9)
buffer.add(0, 1, 10, 1, 0)
buffer.add(1, 2, 11, 1, 0)
buffer.add(2, 3, 12, 1, 0)
buffer.add(3, 4, 13, 1, 0)
buffer.add(4, 5, 14, 1, 0)
obs, act, next_obs, dones, rewards = buffer._get_samples(np.array([1,2,3,4]))
obs, act, next_obs, dones, rewards = buffer._get_samples(np.array([1, 2, 3, 4]))
assert obs.shape == (4, 1)
assert act.shape == (4, 1)
assert next_obs.shape == (4, 1)
assert dones.shape == (4, 1)
assert rewards.shape == (4, 1)
assert np.allclose(dones, np.zeros_like(dones))
assert np.allclose(next_obs, np.array([[5],[5],[5],[5]]))
assert np.allclose(rewards, np.array([1+0.9+0.9**2+0.9**3, 1+0.9+0.9**2, 1+0.9, 1]).reshape(4, 1))
assert np.allclose(next_obs, np.array([[5], [5], [5], [5]]))
assert np.allclose(rewards, np.array([1 + 0.9 + 0.9 ** 2 + 0.9 ** 3, 1 + 0.9 + 0.9 ** 2, 1 + 0.9, 1]).reshape(4, 1))
assert np.allclose(act, np.array([11, 12, 13, 14]).reshape(4, 1))
buffer = NstepReplayBuffer(
10,
spaces.Discrete(5),
spaces.Discrete(5),
n_step=5,
gamma=0.9
)
buffer = NstepReplayBuffer(10, spaces.Discrete(5), spaces.Discrete(5), n_step=5, gamma=0.9)
buffer.add(0, 1, 10, 1, 0)
buffer.add(1, 2, 11, 1, 0)
buffer.add(2, 3, 12, 1, 0)
buffer.add(3, 4, 13, 1, 0)
buffer.add(4, 5, 14, 1, 0)
obs, act, next_obs, dones, rewards = buffer._get_samples(np.array([1,2,3,4]))
obs, act, next_obs, dones, rewards = buffer._get_samples(np.array([1, 2, 3, 4]))
assert obs.shape == (4, 1)
assert act.shape == (4, 1)
assert next_obs.shape == (4, 1)
assert dones.shape == (4, 1)
assert rewards.shape == (4, 1)
assert np.allclose(dones, np.zeros_like(dones))
assert np.allclose(next_obs, np.array([[5],[5],[5],[5]]))
assert np.allclose(rewards, np.array([1+0.9+0.9**2+0.9**3, 1+0.9+0.9**2, 1+0.9, 1]).reshape(4, 1))
assert np.allclose(next_obs, np.array([[5], [5], [5], [5]]))
assert np.allclose(rewards, np.array([1 + 0.9 + 0.9 ** 2 + 0.9 ** 3, 1 + 0.9 + 0.9 ** 2, 1 + 0.9, 1]).reshape(4, 1))
assert np.allclose(act, np.array([11, 12, 13, 14]).reshape(4, 1))
# shouldn't be able to get batch with indice 5 because the pointer is at 5