def _setup_model(self) -> None:
self._setup_lr_schedule()
self.set_random_seed(self.seed)
self.replay_buffer = ReplayBuffer(
self.buffer_size,
self.observation_space,
self.action_space,
self.device,
optimize_memory_usage=self.optimize_memory_usage,
)
self.policy = self.policy_class(
self.observation_space,
self.action_space,
self.support_dim,
self.lr_schedule,
**self.policy_kwargs, # pytype:disable=not-instantiable
)
self.policy = self.policy.to(self.device)
# Convert train freq parameter to TrainFreq object
self._convert_train_freq()
self._create_aliases()
# Target entropy is used when learning the entropy coefficient
if self.target_entropy == "auto":
# automatically set target entropy if needed
self.target_entropy = -np.prod(self.env.action_space.shape).astype(
np.float32)
else:
# Force conversion
# this will also throw an error for unexpected string
self.target_entropy = float(self.target_entropy)
# The entropy coefficient or entropy can be learned automatically
# see Automating Entropy Adjustment for Maximum Entropy RL section
# of https://arxiv.org/abs/1812.05905
if isinstance(self.ent_coef, str) and self.ent_coef.startswith("auto"):
# Default initial value of ent_coef when learned
init_value = 1.0
if "_" in self.ent_coef:
init_value = float(self.ent_coef.split("_")[1])
assert init_value > 0.0, "The initial value of ent_coef must be greater than 0"
# Note: we optimize the log of the entropy coeff which is slightly different from the paper
# as discussed in https://github.com/rail-berkeley/softlearning/issues/37
self.log_ent_coef = th.log(
th.ones(1, device=self.device) *
init_value).requires_grad_(True)
self.ent_coef_optimizer = th.optim.Adam([self.log_ent_coef],
lr=self.lr_schedule(1))
else:
# Force conversion to float
# this will throw an error if a malformed string (different from 'auto')
# is passed
self.ent_coef_tensor = th.tensor(float(self.ent_coef)).to(
self.device)
def _store_transition(
self,
replay_buffer: ReplayBuffer,
buffer_action: np.ndarray,
new_obs: np.ndarray,
reward: np.ndarray,
done: np.ndarray,
infos: List[Dict[str, Any]],
) -> None:
"""
Store transition in the replay buffer.
We store the normalized action and the unnormalized observation.
It also handles terminal observations (because VecEnv resets automatically).
:param replay_buffer: Replay buffer object where to store the transition.
:param buffer_action: normalized action
:param new_obs: next observation in the current episode
or first observation of the episode (when done is True)
:param reward: reward for the current transition
:param done: Termination signal
:param infos: List of additional information about the transition.
It may contain the terminal observations and information about timeout.
"""
# Store only the unnormalized version
if self._vec_normalize_env is not None:
new_obs_ = self._vec_normalize_env.get_original_obs()
reward_ = self._vec_normalize_env.get_original_reward()
else:
# Avoid changing the original ones
self._last_original_obs, new_obs_, reward_ = self._last_obs, new_obs, reward
# As the VecEnv resets automatically, new_obs is already the
# first observation of the next episode
if done and infos[0].get("terminal_observation") is not None:
next_obs = infos[0]["terminal_observation"]
# VecNormalize normalizes the terminal observation
if self._vec_normalize_env is not None:
next_obs = self._vec_normalize_env.unnormalize_obs(next_obs)
else:
next_obs = new_obs_
replay_buffer.add(
self._last_original_obs,
next_obs,
buffer_action,
reward_,
done,
infos,
)
self._last_obs = new_obs
# Save the unnormalized observation
if self._vec_normalize_env is not None:
self._last_original_obs = new_obs_
def _store_transition_not_done_if_truncated(
self,
replay_buffer: ReplayBuffer,
buffer_action: np.ndarray,
new_obs: np.ndarray,
reward: np.ndarray,
done: np.ndarray,
infos: List[Dict[str, Any]],
) -> None:
"""
Store transition in the replay buffer.
We store the normalized action and the unnormalized observation.
It also handles terminal observations (because VecEnv resets automatically).
:param replay_buffer: Replay buffer object where to store the transition.
:param buffer_action: normalized action
:param new_obs: next observation in the current episode
or first observation of the episode (when done is True)
:param reward: reward for the current transition
:param done: Termination signal
:param infos: List of additional information about the transition.
It contains the terminal observations.
"""
# Store only the unnormalized version
if self._vec_normalize_env is not None:
new_obs_ = self._vec_normalize_env.get_original_obs()
reward_ = self._vec_normalize_env.get_original_reward()
else:
# Avoid changing the original ones
self._last_original_obs, new_obs_, reward_ = self._last_obs, new_obs, reward
# As the VecEnv resets automatically, new_obs is already the
# first observation of the next episode
if done and infos[0].get("terminal_observation") is not None:
next_obs = infos[0]["terminal_observation"]
# VecNormalize normalizes the terminal observation
if self._vec_normalize_env is not None:
next_obs = self._vec_normalize_env.unnormalize_obs(next_obs)
else:
next_obs = new_obs_
# NOTE: The monkey patch is inside the following block of code
done_ = np.array([False]) if infos[0].get("TimeLimit.truncated",
False) else done
replay_buffer.add(self._last_original_obs, next_obs, buffer_action,
reward_, done_)
# replay_buffer.add(self._last_original_obs, next_obs, buffer_action, reward_, done)
self._last_obs = new_obs
# Save the unnormalized observation
if self._vec_normalize_env is not None:
self._last_original_obs = new_obs_
def _setup_model(self):
self._setup_lr_schedule()
self.set_random_seed(self.seed)
self.replay_buffer = ReplayBuffer(self.buffer_size, self.observation_space,
self.action_space, self.device)
self.policy = self.policy_class(self.observation_space, self.action_space,
self.lr_schedule, **self.policy_kwargs)
self.policy = self.policy.to(self.device)
def learn(self, initial_models):
mesa_algo = TD3(
"MlpPolicy", self.env, verbose=1, learning_starts=1
) # Note: Unecessarily initializes parameters (could speed up a bit by fixing)'
mesa_algo.set_parameters(to_torch(initial_models), exact_match=False)
LOG_DIR = "/home/jet/catkin_ws/src/marsha/marsha_ai/training/logs/"
MODEL_DIR = "/home/jet/catkin_ws/src/marsha/marsha_ai/training/models/"
callback_list = []
callback_list.append(TensorboardCallback())
callback_list.append(
StopTrainingOnMaxEpisodes(max_episodes=5, verbose=1))
"""callback_list.append(EvalCallback(self.env, best_model_save_path=MODEL_DIR, log_path=LOG_DIR,
deterministic=True,
eval_freq=5,
n_eval_episodes=1))"""
mesa_algo.learn(total_timesteps=1000, callback=callback_list
) #rospy.get_param("/hyperparameters/total_timesteps")
print("finished training! Testing mesa network...")
test_buffer = ReplayBuffer(100,
TaskEnv.observation_space,
TaskEnv.action_space,
device="cuda")
test_env = Monitor(self.env)
done = False
ob = test_env.reset()
while not done:
action, state = mesa_algo.predict(ob)
next_ob, reward, done, info = test_env.step(action)
test_buffer.add(ob, next_ob, action, reward, done, [info])
ob = next_ob
meta_buffer = {"test": test_buffer, "train": mesa_algo.replay_buffer}
optimized_mesa_parameters = mesa_algo.get_parameters()
tf_mesa_models = from_torch(optimized_mesa_parameters)
return meta_buffer, tf_mesa_models
def _setup_model(self) -> None:
super()._setup_model()
self._create_aliases()
# Target entropy is used when learning the entropy coefficient
if self.target_entropy == 'auto':
# automatically set target entropy if needed
self.target_entropy = -np.prod(self.env.action_space.shape).astype(
np.float32)
else:
# Force conversion
# this will also throw an error for unexpected string
self.target_entropy = float(self.target_entropy)
# The entropy coefficient or entropy can be learned automatically
# see Automating Entropy Adjustment for Maximum Entropy RL section
# of https://arxiv.org/abs/1812.05905
if isinstance(self.ent_coef, str) and self.ent_coef.startswith('auto'):
# Default initial value of ent_coef when learned
init_value = 1.0
if '_' in self.ent_coef:
init_value = float(self.ent_coef.split('_')[1])
assert init_value > 0., "The initial value of ent_coef must be greater than 0"
# Note: we optimize the log of the entropy coeff which is slightly different from the paper
# as discussed in https://github.com/rail-berkeley/softlearning/issues/37
self.log_ent_coef = th.log(
th.ones(1, device=self.device) *
init_value).requires_grad_(True)
self.ent_coef_optimizer = th.optim.Adam([self.log_ent_coef],
lr=self.lr_schedule(1))
else:
# Force conversion to float
# this will throw an error if a malformed string (different from 'auto')
# is passed
self.ent_coef_tensor = th.tensor(float(self.ent_coef)).to(
self.device)
self.bc_buffer = ReplayBuffer(int(1e4), self.observation_space,
self.action_space, self.device)
def _setup_model(self) -> None:
self._setup_lr_schedule()
self.set_random_seed(self.seed)
self.replay_buffer = ReservoirBuffer(
self.buffer_size,
self.observation_space,
self.action_space,
self.device,
optimize_memory_usage=self.optimize_memory_usage,
)
self.current_experience_buffer = ReplayBuffer(
1,
self.observation_space,
self.action_space,
self.device,
optimize_memory_usage=False,
)
self.policy = self.policy_class(
self.observation_space,
self.action_space,
self.lr_schedule,
**self.policy_kwargs # pytype:disable=not-instantiable
)
self.policy = self.policy.to(self.device)
def _setup_model(self) -> None:
self._setup_lr_schedule()
self.set_random_seed(self.seed)
self.replay_buffer = ReplayBuffer(
self.buffer_size,
self.observation_space,
self.action_space,
self.device,
optimize_memory_usage=self.optimize_memory_usage,
)
self.policy = self.policy_class(
self.observation_space,
self.action_space,
self.lr_schedule,
**self.policy_kwargs, # pytype:disable=not-instantiable
)
self.policy = self.policy.to(self.device)
# Convert train freq parameter to TrainFreq object
self._convert_train_freq()
def setup_buffer(self, num_samples):
assert self.n_envs == 1, "I don't think multiple envs works for offline policies, but you can check and make suitable updates"
self._old_buffer = self.replay_buffer
callback = DoNothingCallback(self)
self.replay_buffer = ReplayBuffer(
num_samples,
self.observation_space,
self.action_space,
self.device,
)
self.env.reset()
train_freq = TrainFreq(num_samples, TrainFrequencyUnit("step"))
self.collect_rollouts(
self.env,
train_freq=train_freq,
action_noise=self.action_noise,
callback=callback,
learning_starts=0,
replay_buffer=self.replay_buffer,
log_interval=10,
)
class AWAC(OffPolicyAlgorithm):
"""
Soft Actor-Critic (SAC)
Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor,
This implementation borrows code from original implementation (https://github.com/haarnoja/sac)
from OpenAI Spinning Up (https://github.com/openai/spinningup), from the softlearning repo
(https://github.com/rail-berkeley/softlearning/)
and from Stable Baselines (https://github.com/hill-a/stable-baselines)
Paper: https://arxiv.org/abs/1801.01290
Introduction to SAC: https://spinningup.openai.com/en/latest/algorithms/sac.html
Note: we use double q target and not value target as discussed
in https://github.com/hill-a/stable-baselines/issues/270
:param policy: (SACPolicy or str) The policy model to use (MlpPolicy, CnnPolicy, ...)
:param env: (GymEnv or str) The environment to learn from (if registered in Gym, can be str)
:param learning_rate: (float or callable) learning rate for adam optimizer,
the same learning rate will be used for all networks (Q-Values, Actor and Value function)
it can be a function of the current progress remaining (from 1 to 0)
:param buffer_size: (int) size of the replay buffer
:param learning_starts: (int) how many steps of the model to collect transitions for before learning starts
:param batch_size: (int) Minibatch size for each gradient update
:param tau: (float) the soft update coefficient ("Polyak update", between 0 and 1)
:param gamma: (float) the discount factor
:param train_freq: (int) Update the model every ``train_freq`` steps.
:param gradient_steps: (int) How many gradient update after each step
:param n_episodes_rollout: (int) Update the model every ``n_episodes_rollout`` episodes.
Note that this cannot be used at the same time as ``train_freq``
:param action_noise: (ActionNoise) the action noise type (None by default), this can help
for hard exploration problem. Cf common.noise for the different action noise type.
:param ent_coef: (str or float) Entropy regularization coefficient. (Equivalent to
inverse of reward scale in the original SAC paper.) Controlling exploration/exploitation trade-off.
Set it to 'auto' to learn it automatically (and 'auto_0.1' for using 0.1 as initial value)
:param target_update_interval: (int) update the target network every ``target_network_update_freq`` steps.
:param target_entropy: (str or float) target entropy when learning ``ent_coef`` (``ent_coef = 'auto'``)
:param create_eval_env: (bool) Whether to create a second environment that will be
used for evaluating the agent periodically. (Only available when passing string for the environment)
:param policy_kwargs: (dict) additional arguments to be passed to the policy on creation
:param verbose: (int) the verbosity level: 0 no output, 1 info, 2 debug
:param seed: (int) Seed for the pseudo random generators
:param device: (str or th.device) Device (cpu, cuda, ...) on which the code should be run.
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: Union[str, Type[AWACPolicy]],
env: Union[GymEnv, str],
learning_rate: Union[float, Callable] = 3e-4,
buffer_size: int = int(1e6),
learning_starts: int = 100,
batch_size: int = 256,
tau: float = 0.005,
gamma: float = 0.99,
train_freq: int = 1,
gradient_steps: int = 1,
n_episodes_rollout: int = -1,
action_noise: Optional[ActionNoise] = None,
ent_coef: Union[str, float] = 'auto',
target_update_interval: int = 1,
target_entropy: Union[str, float] = 'auto',
awr_use_mle_for_vf: bool = True,
beta: int = 50,
tensorboard_log: Optional[str] = None,
create_eval_env: bool = False,
policy_kwargs: Dict[str, Any] = None,
verbose: int = 0,
seed: Optional[int] = None,
device: Union[th.device, str] = 'auto',
_init_setup_model: bool = True):
super().__init__(policy,
env,
AWACPolicy,
learning_rate,
buffer_size,
learning_starts,
batch_size,
policy_kwargs,
tensorboard_log,
verbose,
device,
create_eval_env=create_eval_env,
seed=seed,
use_sde=False,
sde_sample_freq=-1,
use_sde_at_warmup=False)
self.target_entropy = target_entropy
self.log_ent_coef = None # type: Optional[th.Tensor]
self.target_update_interval = target_update_interval
self.tau = tau
# Entropy coefficient / Entropy temperature
# Inverse of the reward scale
self.ent_coef = ent_coef
self.target_update_interval = target_update_interval
self.train_freq = train_freq
self.gradient_steps = gradient_steps
self.n_episodes_rollout = n_episodes_rollout
self.action_noise = action_noise
self.gamma = gamma
self.ent_coef_optimizer = None
self.awr_use_mle_for_vf = awr_use_mle_for_vf
self.beta = beta
self.bc_buffer = None
if _init_setup_model:
self._setup_model()
def _setup_model(self) -> None:
super()._setup_model()
self._create_aliases()
# Target entropy is used when learning the entropy coefficient
if self.target_entropy == 'auto':
# automatically set target entropy if needed
self.target_entropy = -np.prod(self.env.action_space.shape).astype(
np.float32)
else:
# Force conversion
# this will also throw an error for unexpected string
self.target_entropy = float(self.target_entropy)
# The entropy coefficient or entropy can be learned automatically
# see Automating Entropy Adjustment for Maximum Entropy RL section
# of https://arxiv.org/abs/1812.05905
if isinstance(self.ent_coef, str) and self.ent_coef.startswith('auto'):
# Default initial value of ent_coef when learned
init_value = 1.0
if '_' in self.ent_coef:
init_value = float(self.ent_coef.split('_')[1])
assert init_value > 0., "The initial value of ent_coef must be greater than 0"
# Note: we optimize the log of the entropy coeff which is slightly different from the paper
# as discussed in https://github.com/rail-berkeley/softlearning/issues/37
self.log_ent_coef = th.log(
th.ones(1, device=self.device) *
init_value).requires_grad_(True)
self.ent_coef_optimizer = th.optim.Adam([self.log_ent_coef],
lr=self.lr_schedule(1))
else:
# Force conversion to float
# this will throw an error if a malformed string (different from 'auto')
# is passed
self.ent_coef_tensor = th.tensor(float(self.ent_coef)).to(
self.device)
self.bc_buffer = ReplayBuffer(int(1e4), self.observation_space,
self.action_space, self.device)
def _create_aliases(self) -> None:
self.actor = self.policy.actor
self.critic = self.policy.critic
self.critic_target = self.policy.critic_target
def pretrain_bc(self, gradient_steps: int, batch_size: int = 64):
statistics = []
with trange(gradient_steps) as t:
for gradient_step in t:
replay_data = self.bc_buffer.sample(
batch_size, env=self._vec_normalize_env)
dist = self.actor(replay_data.observations)
actions_pi, log_prob = dist.log_prob_and_rsample()
actor_loss = -log_prob.mean()
actor_mse_loss = F.mse_loss(actions_pi.detach(),
replay_data.actions)
self.actor.optimizer.zero_grad()
actor_loss.backward()
self.actor.optimizer.step()
statistics.append((actor_loss.item(), actor_mse_loss.item()))
t.set_postfix(mse_loss=actor_mse_loss.item(),
policy_loss=actor_loss.item())
actor_losses, mse_losses = tuple(zip(*statistics))
logger.record("pretrain/n_updates",
self._n_updates,
exclude='tensorboard')
logger.record("pretrain/actor_loss", np.mean(actor_losses))
logger.record("pretrain/actor_mse_loss", np.mean(mse_losses))
def pretrain_rl(self, gradient_steps: int, batch_size: int = 64) -> None:
statistics = []
with trange(gradient_steps) as t:
for gradient_step in t:
replay_data = self.replay_buffer.sample(
batch_size, env=self._vec_normalize_env)
stats = self.train_batch(replay_data)
statistics.append(stats)
self._n_updates += 1
t.set_postfix(qf_loss=stats[1], policy_loss=stats[0])
actor_losses, critic_losses, ent_coef_losses, ent_coefs = tuple(
zip(*statistics))
logger.record("pretrain/n_updates",
self._n_updates,
exclude='tensorboard')
logger.record("pretrain/ent_coef", np.mean(ent_coefs))
logger.record("pretrain/actor_loss", np.mean(actor_losses))
logger.record("pretrain/critic_loss", np.mean(critic_losses))
logger.record("pretrain/ent_coef_loss", np.mean(ent_coef_losses))
def train(self, gradient_steps: int, batch_size: int = 64) -> None:
statistics = []
for gradient_step in range(gradient_steps):
replay_data = self.replay_buffer.sample(
batch_size, env=self._vec_normalize_env)
stats = self.train_batch(replay_data)
statistics.append(stats)
self._n_updates += 1
actor_losses, critic_losses, ent_coef_losses, ent_coefs = tuple(
zip(*statistics))
logger.record("train/n_updates",
self._n_updates,
exclude='tensorboard')
logger.record("train/ent_coef", np.mean(ent_coefs))
logger.record("train/actor_loss", np.mean(actor_losses))
logger.record("train/critic_loss", np.mean(critic_losses))
logger.record("train/ent_coef_loss", np.mean(ent_coef_losses))
def train_batch(self, replay_data):
# Action by the current actor for the sampled state
dist = self.actor(replay_data.observations)
actions_pi, log_prob = dist.log_prob_and_rsample()
actor_mle = dist.mean
"""ent_coeff loss"""
ent_coef_loss = None
if self.ent_coef_optimizer is not None:
# Important: detach the variable from the graph
# so we don't change it with other losses
# see https://github.com/rail-berkeley/softlearning/issues/60
ent_coef = th.exp(self.log_ent_coef.detach())
ent_coef_loss = -(
self.log_ent_coef *
(log_prob + self.target_entropy).detach()).mean()
else:
ent_coef = self.ent_coef_tensor
"""q loss"""
with th.no_grad():
# Select action according to policy
next_dist = self.actor(replay_data.next_observations)
next_actions, next_log_prob = next_dist.log_prob_and_rsample()
# Compute the target Q value
target_q1, target_q2 = self.critic_target(
replay_data.next_observations, next_actions)
target_q = th.min(target_q1, target_q2)
target_q = replay_data.rewards + (
1 - replay_data.dones) * self.gamma * target_q
# td error + entropy term
# q_backup = target_q - ent_coef * next_log_prob
q_backup = target_q
# Get current Q estimates
# using action from the replay buffer
current_q1, current_q2 = self.critic(replay_data.observations,
replay_data.actions)
# Compute critic loss
critic_loss = 0.5 * (F.mse_loss(current_q1, q_backup) +
F.mse_loss(current_q2, q_backup))
"""action loss"""
# Advantage-weighted regression
# if self.awr_use_mle_for_vf:
# v1_pi,v2_pi = self.critic(replay_data.observations, actor_mle)
# v_pi = th.min(v1_pi, v2_pi)
# else:
# v1_pi,v2_pi = self.critic(replay_data.observations, actions_pi)
# v_pi = th.min(v1_pi, v2_pi)
q_adv = th.min(current_q1, current_q2)
v1_pi, v2_pi = self.critic(replay_data.observations, actor_mle)
v_pi = th.min(v1_pi, v2_pi)
# q_adv = th.min(*self.critic(replay_data.observations, actions_pi))
score = q_adv - v_pi
weights = F.softmax(score / self.beta, dim=0)
# actor_loss = ent_coef * log_prob.mean()
actor_logpp = dist.log_prob(replay_data.actions)
actor_loss = (-actor_logpp * len(weights) * weights.detach()).mean()
"""Updates"""
# Optimize entropy coefficient, also called
# entropy temperature or alpha in the paper
if ent_coef_loss is not None:
self.ent_coef_optimizer.zero_grad()
ent_coef_loss.backward()
self.ent_coef_optimizer.step()
# Optimize the critic
self.critic.optimizer.zero_grad()
critic_loss.backward()
self.critic.optimizer.step()
# Optimize the actor
self.actor.optimizer.zero_grad()
actor_loss.backward()
self.actor.optimizer.step()
# Update target networks
if self._n_updates % self.target_update_interval == 0:
for param, target_param in zip(self.critic.parameters(),
self.critic_target.parameters()):
target_param.data.copy_(self.tau * param.data +
(1 - self.tau) * target_param.data)
if ent_coef_loss is None:
ent_coef_loss = th.tensor([0])
return actor_loss.item(), critic_loss.item(), ent_coef_loss.item(
), ent_coef.item()
def learn(self,
total_timesteps: int,
callback: MaybeCallback = None,
log_interval: int = 4,
eval_env: Optional[GymEnv] = None,
eval_freq: int = -1,
n_eval_episodes: int = 5,
tb_log_name: str = "AWAC",
eval_log_path: Optional[str] = None,
reset_num_timesteps: bool = True) -> OffPolicyAlgorithm:
total_timesteps, callback = self._setup_learn(
total_timesteps, eval_env, callback, eval_freq, n_eval_episodes,
eval_log_path, reset_num_timesteps, tb_log_name)
callback.on_training_start(locals(), globals())
self.pretrain_bc(int(1e3), batch_size=self.batch_size)
observations, actions, next_observations, rewards, dones = self.bc_buffer.observations, self.bc_buffer.actions, self.bc_buffer.next_observations, self.bc_buffer.rewards, self.bc_buffer.dones
for data in zip(observations, next_observations, actions, rewards,
dones):
self.replay_buffer.add(*data)
self.pretrain_rl(int(1e4), batch_size=self.batch_size)
while self.num_timesteps < total_timesteps:
rollout = self.collect_rollouts(
self.env,
n_episodes=self.n_episodes_rollout,
n_steps=self.train_freq,
action_noise=self.action_noise,
callback=callback,
learning_starts=self.learning_starts,
replay_buffer=self.replay_buffer,
log_interval=log_interval)
if rollout.continue_training is False:
break
self._update_current_progress_remaining(self.num_timesteps,
total_timesteps)
if self.num_timesteps > 0 and self.num_timesteps > self.learning_starts:
gradient_steps = self.gradient_steps if self.gradient_steps > 0 else rollout.episode_timesteps
self.train(gradient_steps, batch_size=self.batch_size)
callback.on_training_end()
return self
def excluded_save_params(self) -> List[str]:
"""
Returns the names of the parameters that should be excluded by default
when saving the model.
:return: (List[str]) List of parameters that should be excluded from save
"""
# Exclude aliases
return super().excluded_save_params() + [
"actor", "critic", "critic_target"
]
def get_torch_variables(self) -> Tuple[List[str], List[str]]:
"""
cf base class
"""
state_dicts = ["policy", "actor.optimizer", "critic.optimizer"]
saved_tensors = ['log_ent_coef']
if self.ent_coef_optimizer is not None:
state_dicts.append('ent_coef_optimizer')
else:
saved_tensors.append('ent_coef_tensor')
return state_dicts, saved_tensors
qf1 = QNetwork(envs).to(device)
qf2 = QNetwork(envs).to(device)
qf1_target = QNetwork(envs).to(device)
qf2_target = QNetwork(envs).to(device)
target_actor = Actor(envs).to(device)
target_actor.load_state_dict(actor.state_dict())
qf1_target.load_state_dict(qf1.state_dict())
qf2_target.load_state_dict(qf2.state_dict())
q_optimizer = optim.Adam(list(qf1.parameters()) + list(qf2.parameters()),
lr=args.learning_rate)
actor_optimizer = optim.Adam(list(actor.parameters()),
lr=args.learning_rate)
envs.single_observation_space.dtype = np.float32
rb = ReplayBuffer(args.buffer_size,
envs.single_observation_space,
envs.single_action_space,
device=device)
loss_fn = nn.MSELoss()
# TRY NOT TO MODIFY: start the game
obs = envs.reset()
for global_step in range(args.total_timesteps):
# ALGO LOGIC: put action logic here
if global_step < args.learning_starts:
actions = envs.action_space.sample()
else:
actions = actor.forward(torch.Tensor(obs).to(device))
actions = np.array([(
actions.tolist()[0] +
np.random.normal(0,
max_action * args.exploration_noise,
class CMVCVaRSAC(OffPolicyAlgorithm):
def __init__(
self,
policy: Union[str, Type[CMVC51SACPolicy]],
env: Union[GymEnv, str],
min_v: float = -25,
max_v: float = +25,
support_dim: int = 200,
learning_rate: Union[float, Schedule] = 3e-4,
buffer_size: int = int(5e4),
learning_starts: int = 100,
batch_size: int = 64,
tau: float = 0.005,
gamma: float = 0.99,
train_freq: Union[int, Tuple[int, str]] = 1,
gradient_steps: int = 1,
action_noise: Optional[ActionNoise] = None,
optimize_memory_usage: bool = False,
ent_coef: Union[str, float] = "auto",
target_update_interval: int = 1,
target_entropy: Union[str, float] = "auto",
use_sde: bool = False,
sde_sample_freq: int = -1,
use_sde_at_warmup: bool = False,
tensorboard_log: Optional[str] = None,
create_eval_env: bool = False,
policy_kwargs: Dict[str, Any] = None,
verbose: int = 0,
seed: Optional[int] = None,
device: Union[th.device, str] = "auto",
_init_setup_model: bool = True,
cvar_alpha=0.3,
cmv_beta=1,
):
super(CMVCVaRSAC, self).__init__(
policy,
env,
CMVC51SACPolicy,
learning_rate,
buffer_size,
learning_starts,
batch_size,
tau,
gamma,
train_freq,
gradient_steps,
action_noise,
policy_kwargs=policy_kwargs,
tensorboard_log=tensorboard_log,
verbose=verbose,
device=device,
create_eval_env=create_eval_env,
seed=seed,
use_sde=use_sde,
sde_sample_freq=sde_sample_freq,
use_sde_at_warmup=use_sde_at_warmup,
optimize_memory_usage=optimize_memory_usage,
supported_action_spaces=(gym.spaces.Box),
)
self.target_entropy = target_entropy
self.log_ent_coef = None # type: Optional[th.Tensor]
# Entropy coefficient / Entropy temperature
# Inverse of the reward scale
self.ent_coef = ent_coef
self.target_update_interval = target_update_interval
self.ent_coef_optimizer = None
self.min_v = min_v
self.max_v = max_v
self.support_dim = support_dim
self.interval = (1 / (support_dim - 1)) * (max_v - min_v)
self.supports = th.from_numpy(
np.array([min_v + i * self.interval for i in range(support_dim)],
dtype=np.float32)).to(self.device)
self._total_timesteps = None
self.cvar_alpha = cvar_alpha
self.cmv_beta = cmv_beta
if _init_setup_model:
self._setup_model()
def _setup_model(self) -> None:
self._setup_lr_schedule()
self.set_random_seed(self.seed)
self.replay_buffer = ReplayBuffer(
self.buffer_size,
self.observation_space,
self.action_space,
self.device,
optimize_memory_usage=self.optimize_memory_usage,
)
self.policy = self.policy_class(
self.observation_space,
self.action_space,
self.support_dim,
self.lr_schedule,
**self.policy_kwargs, # pytype:disable=not-instantiable
)
self.policy = self.policy.to(self.device)
# Convert train freq parameter to TrainFreq object
self._convert_train_freq()
self._create_aliases()
# Target entropy is used when learning the entropy coefficient
if self.target_entropy == "auto":
# automatically set target entropy if needed
self.target_entropy = -np.prod(self.env.action_space.shape).astype(
np.float32)
else:
# Force conversion
# this will also throw an error for unexpected string
self.target_entropy = float(self.target_entropy)
# The entropy coefficient or entropy can be learned automatically
# see Automating Entropy Adjustment for Maximum Entropy RL section
# of https://arxiv.org/abs/1812.05905
if isinstance(self.ent_coef, str) and self.ent_coef.startswith("auto"):
# Default initial value of ent_coef when learned
init_value = 1.0
if "_" in self.ent_coef:
init_value = float(self.ent_coef.split("_")[1])
assert init_value > 0.0, "The initial value of ent_coef must be greater than 0"
# Note: we optimize the log of the entropy coeff which is slightly different from the paper
# as discussed in https://github.com/rail-berkeley/softlearning/issues/37
self.log_ent_coef = th.log(
th.ones(1, device=self.device) *
init_value).requires_grad_(True)
self.ent_coef_optimizer = th.optim.Adam([self.log_ent_coef],
lr=self.lr_schedule(1))
else:
# Force conversion to float
# this will throw an error if a malformed string (different from 'auto')
# is passed
self.ent_coef_tensor = th.tensor(float(self.ent_coef)).to(
self.device)
def _create_aliases(self) -> None:
self.actor = self.policy.actor
self.critic = self.policy.critic
self.critic_target = self.policy.critic_target
self.cmv_net = self.policy.cmv_net
self.beta_critic = self.policy.beta_critic
self.beta_critic_target = self.policy.beta_critic_target
def projection(self, support_rows, target_z):
projected_target_z = th.zeros_like(target_z)
support_rows = support_rows.clamp(self.min_v, self.max_v - 1e-3)
p = ((support_rows - self.min_v) % self.interval) / self.interval
idx = ((support_rows - self.min_v) // self.interval).long()
projected_target_z = projected_target_z.scatter_add(
1, idx, target_z * p)
projected_target_z = projected_target_z.scatter_add(
1, idx + 1, target_z * (1 - p))
return projected_target_z
def train(self, gradient_steps: int, batch_size: int = 64) -> None:
# Update optimizers learning rate
optimizers = [self.actor.optimizer, self.critic.optimizer]
if self.ent_coef_optimizer is not None:
optimizers += [self.ent_coef_optimizer]
# Update learning rate according to lr schedule
self._update_learning_rate(optimizers)
ent_coef_losses, ent_coefs = [], []
actor_losses, critic_losses, reward_losses, feature_pred_losses = [], [], [], []
cvars = []
qs = []
critic_beta_losses = []
for gradient_step in range(gradient_steps):
# Sample replay buffer
replay_data = self.replay_buffer.sample(
batch_size, env=self._vec_normalize_env)
# We need to sample because `log_std` may have changed between two gradient steps
if self.use_sde:
self.actor.reset_noise()
# Action by the current actor for the sampled state
actions_pi, log_prob = self.actor.action_log_prob(
replay_data.observations)
log_prob = log_prob.reshape(-1, 1)
ent_coef_loss = None
if self.ent_coef_optimizer is not None:
# Important: detach the variable from the graph
# so we don't change it with other losses
# see https://github.com/rail-berkeley/softlearning/issues/60
ent_coef = th.exp(self.log_ent_coef.detach())
ent_coef_loss = -(
self.log_ent_coef *
(log_prob + self.target_entropy).detach()).mean()
ent_coef_losses.append(ent_coef_loss.item())
else:
ent_coef = self.ent_coef_tensor
ent_coefs.append(ent_coef.item())
# Optimize entropy coefficient, also called
# entropy temperature or alpha in the paper
if ent_coef_loss is not None:
self.ent_coef_optimizer.zero_grad()
ent_coef_loss.backward()
self.ent_coef_optimizer.step()
with th.no_grad():
# Select action according to policy
next_actions, next_log_prob = self.actor.action_log_prob(
replay_data.next_observations)
# Compute the next Q values: min over all critics targets
target_zs = th.cat(self.critic_target(
replay_data.next_observations, next_actions),
dim=1)
# add entropy term
target_supports = self.supports.clone().detach()
target_supports = target_supports - ent_coef * next_log_prob.reshape(
-1, 1)
# td error + entropy term
target_supports = replay_data.rewards + (
1 - replay_data.dones) * self.gamma * target_supports
# Get current Q-values estimates for each critic network
# using action from the replay buffer
current_zs = th.cat(self.critic(replay_data.observations,
replay_data.actions),
dim=1)
target_zs = self.projection(target_supports, target_zs)
# Compute critic loss
critic_loss = -th.mean(th.log(current_zs + 1e-12) * target_zs)
critic_losses.append(critic_loss.item())
# Optimize the critic
self.critic.optimizer.zero_grad()
critic_loss.backward()
self.critic.optimizer.step()
# Compute
# For the CMV Learning
with th.no_grad():
next_z = self.critic_target.features_extractor(
replay_data.next_observations)
# predicted reward, and predicted next observation
r_pred, z_pred = self.cmv_net(replay_data.observations,
replay_data.actions)
mse_r_pred = th.mean(th.square(r_pred - replay_data.rewards))
mse_z_pred = th.mean(th.square(z_pred - next_z))
loss_cmv = mse_r_pred + mse_z_pred
reward_losses.append(mse_r_pred.item())
feature_pred_losses.append(mse_z_pred.item())
# Optimize the CMV Nets
self.cmv_net.optimizer.zero_grad()
loss_cmv.backward()
self.cmv_net.optimizer.step()
with th.no_grad():
# Select action according to policy
# Compute the next Q values: min over all critics targets
next_q_beta_values = th.cat(self.beta_critic_target(
replay_data.next_observations, next_actions),
dim=1)
next_q_beta_values, _ = th.min(next_q_beta_values,
dim=1,
keepdim=True)
# add entropy term
next_q_beta_values = next_q_beta_values - ent_coef * next_log_prob.reshape(
-1, 1)
# td error + entropy term
target_q_beta_values = loss_cmv.detach() + (
1 - replay_data.dones) * (self.gamma**
2) * next_q_beta_values
# Get current Q-beta values estimates for each critic network
# using action from the replay buffer
current_q_beta_values = self.beta_critic(replay_data.observations,
replay_data.actions)
# Compute critic beta loss
critic_beta_loss = \
0.5 * sum(
[F.mse_loss(current_q_beta, target_q_beta_values) for current_q_beta in current_q_beta_values])
critic_beta_losses.append(critic_beta_loss.item())
# Optimize the critic beta
self.beta_critic.optimizer.zero_grad()
critic_beta_loss.backward()
self.beta_critic.optimizer.step()
# Compute actor loss
# Alternative: actor_loss = th.mean(log_prob - qf1_pi - qf1_beta_pi)
# Mean over all critic networks
z_pi = th.cat(self.critic.forward(replay_data.observations,
actions_pi),
dim=1)
z_cdf = th.cumsum(z_pi, dim=-1)
adjust_pdf = th.where(th.le(z_cdf, self.cvar_alpha), z_pi,
th.zeros_like(z_pi))
adjust_pdf = th.div(adjust_pdf,
th.sum(adjust_pdf, dim=-1, keepdim=True))
q_pi = adjust_pdf @ self.supports
cvars.append(th.mean(q_pi).item())
qs.append(th.mean(z_pi @ self.supports).item())
q_beta_values_pi = th.cat(self.beta_critic.forward(
replay_data.observations, actions_pi),
dim=1)
max_qf_beta_pi, _ = th.max(q_beta_values_pi, dim=1, keepdim=True)
actor_loss = (ent_coef * log_prob - q_pi +
self.cmv_beta * next_q_beta_values).mean()
actor_losses.append(actor_loss.item())
# Optimize the actor
self.actor.optimizer.zero_grad()
actor_loss.backward()
self.actor.optimizer.step()
# Update target networks
if gradient_step % self.target_update_interval == 0:
polyak_update(self.critic.parameters(),
self.critic_target.parameters(), self.tau)
polyak_update(self.beta_critic.parameters(),
self.beta_critic_target.parameters(), self.tau)
self._n_updates += gradient_steps
fps = int(self.num_timesteps / (time.time() - self.start_time))
remaining_steps = self._total_timesteps - self.num_timesteps
eta = int(round(remaining_steps / fps))
logger.record("time/eta",
timedelta(seconds=eta),
exclude="tensorboard")
logger.record("train/CVaR Alpha", self.cvar_alpha)
logger.record("train/CMV Beta", self.cmv_beta)
logger.record("train/CVaR", np.mean(cvars))
logger.record("train/Q-value", np.mean(qs))
logger.record("train/n_updates",
self._n_updates,
exclude="tensorboard")
logger.record("train/ent_coef", np.mean(ent_coefs))
logger.record("train/actor_loss", np.mean(actor_losses))
logger.record("train/critic_loss", np.mean(critic_losses))
logger.record("train/reward error", np.mean(reward_losses))
logger.record("train/s_t+1_error", np.mean(feature_pred_losses))
logger.record("train/beta_Q_loss", np.mean(critic_beta_losses))
if len(ent_coef_losses) > 0:
logger.record("train/ent_coef_loss", np.mean(ent_coef_losses))
def learn(
self,
total_timesteps: int,
callback: MaybeCallback = None,
log_interval: int = 4,
eval_env: Optional[GymEnv] = None,
eval_freq: int = -1,
n_eval_episodes: int = 5,
tb_log_name: str = "C51SAC",
eval_log_path: Optional[str] = None,
reset_num_timesteps: bool = True,
) -> OffPolicyAlgorithm:
return super(CMVCVaRSAC, self).learn(
total_timesteps=total_timesteps,
callback=callback,
log_interval=log_interval,
eval_env=eval_env,
eval_freq=eval_freq,
n_eval_episodes=n_eval_episodes,
tb_log_name=tb_log_name,
eval_log_path=eval_log_path,
reset_num_timesteps=reset_num_timesteps,
)
def _excluded_save_params(self) -> List[str]:
return super(CMVCVaRSAC, self)._excluded_save_params() + [
"actor", "critic", "critic_target"
]
def _get_torch_save_params(self) -> Tuple[List[str], List[str]]:
state_dicts = ["policy", "actor.optimizer", "critic.optimizer"]
saved_pytorch_variables = ["log_ent_coef"]
if self.ent_coef_optimizer is not None:
state_dicts.append("ent_coef_optimizer")
else:
saved_pytorch_variables.append("ent_coef_tensor")
return state_dicts, saved_pytorch_variables
device = torch.device("cuda" if torch.cuda.is_available() and args.cuda else "cpu")
envs = VecFrameStack(
DummyVecEnv([make_env(args.gym_id, args.seed, 0)]),
4,
)
assert isinstance(envs.action_space, Discrete), "only discrete action space is supported"
# TRY NOT TO MODIFY: seeding
random.seed(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
torch.backends.cudnn.deterministic = args.torch_deterministic
# ALGO LOGIC: initialize agent here:
rb = ReplayBuffer(args.buffer_size, envs.observation_space, envs.action_space, device=device, optimize_memory_usage=True)
q_network = QNetwork(envs).to(device)
target_network = QNetwork(envs).to(device)
target_network.load_state_dict(q_network.state_dict())
optimizer = optim.Adam(q_network.parameters(), lr=args.learning_rate)
loss_fn = nn.MSELoss()
# TRY NOT TO MODIFY: start the game
obs = envs.reset()
for global_step in range(args.total_timesteps):
epsilon = linear_schedule(args.start_e, args.end_e, args.exploration_fraction * args.total_timesteps, global_step)
if random.random() < epsilon:
actions = [envs.action_space.sample()]
else:
logits = q_network.forward(torch.Tensor(obs).to(device))
actions = torch.argmax(logits, dim=1).cpu().numpy()
def main(args):
if args.env_type == 'atari':
# env = make_atari_env(
# env_name=args.env_name,
# action_repeat=args.action_repeat,
# frame_stack=args.frame_stack
# )
# eval_env = make_atari_env(
# env_name=args.env_name,
# action_repeat=args.action_repeat,
# frame_stack=args.frame_stack
# )
env = make_atari_env(env_id=args.env_name)
eval_env = make_atari_env(env_id=args.env_name)
elif args.env_type == 'dmc_locomotion':
env = make_locomotion_env(
env_name=args.env_name,
seed=args.seed,
episode_length=args.episode_length,
from_pixels=args.pixel_obs,
action_repeat=args.action_repeat,
obs_height=args.obs_height,
obs_width=args.obs_width,
camera_id=args.env_camera_id,
mode=args.mode
)
eval_env = make_locomotion_env(
env_name=args.env_name,
seed=args.seed,
episode_length=args.episode_length,
from_pixels=args.pixel_obs,
action_repeat=args.action_repeat,
obs_height=args.obs_height,
obs_width=args.obs_width,
camera_id=args.env_camera_id,
mode=args.mode
)
# Initialize environment
if args.seed is None:
args.seed = np.random.randint(int(1e9))
utils.set_seed_everywhere(args.seed, env=env, eval_env=eval_env)
utils.make_dir(args.work_dir)
model_dir = utils.make_dir(os.path.join(args.work_dir, 'model'))
video_dir = utils.make_dir(os.path.join(args.work_dir, 'video'))
video = VideoRecorder(video_dir if args.save_video else None, args.env_type,
height=448, width=448, camera_id=args.video_camera_id)
# Prepare agent
assert torch.cuda.is_available(), 'must have cuda enabled'
device = torch.device(args.device)
if args.env_type == 'atari':
# replay_buffer = buffers.FrameStackReplayBuffer(
# obs_space=env.observation_space,
# action_space=env.action_space,
# capacity=args.replay_buffer_capacity,
# frame_stack=args.frame_stack,
# device=device,
# optimize_memory_usage=True,
# )
from stable_baselines3.common.buffers import ReplayBuffer
replay_buffer = ReplayBuffer(
args.replay_buffer_capacity,
env.observation_space,
env.action_space,
device,
optimize_memory_usage=True,
)
# replay_buffer = buffers.ReplayBuffer(
# obs_space=env.observation_space,
# action_space=env.action_space,
# capacity=args.replay_buffer_capacity,
# device=device,
# optimize_memory_usage=True,
# )
elif args.env_type == 'dmc_locomotion':
replay_buffer = buffers.ReplayBuffer(
obs_space=env.observation_space,
action_space=env.action_space,
capacity=args.replay_buffer_capacity,
device=device,
optimize_memory_usage=True,
)
agent = make_agent(
obs_space=env.observation_space,
action_space=env.action_space,
device=device,
args=args
)
logger = Logger(args.work_dir,
log_frequency=args.log_freq,
action_repeat=args.action_repeat,
save_tb=args.save_tb)
episode, episode_reward, episode_step, done = 0, 0, 0, True
obs = env.reset()
start_time = time.time()
for step in range(args.train_steps + 1):
# (chongyi zheng): we can also evaluate and save model when current episode is not finished
# Evaluate agent periodically
if step % args.eval_freq == 0:
print('Evaluating:', args.work_dir)
logger.log('eval/episode', episode, step)
evaluate(eval_env, agent, video, args.num_eval_episodes, logger, step)
# Save agent periodically
if step % args.save_freq == 0 and step > 0:
if args.save_model:
agent.save(model_dir, step)
if done:
if step > 0:
logger.log('train/duration', time.time() - start_time, step)
start_time = time.time()
logger.dump(step, ty='train', save=(step > args.init_steps))
logger.log('train/episode_reward', episode_reward, step)
obs = env.reset()
done = False
episode_reward = 0
episode_step = 0
episode += 1
logger.log('train/episode', episode, step)
# Sample action for data collection
if step < args.init_steps:
action = env.action_space.sample()
else:
# with utils.eval_mode(agent):
action = agent.act(obs, True)
# action, _ = agent.predict(obs, deterministic=False)
if 'dqn' in args.algo:
agent.schedule_exploration_rate(step, logger)
# agent._update_current_progress_remaining(step, args.train_steps)
# agent._on_step()
# Run training update
if step >= args.init_steps and step % args.train_freq == 0:
# TODO (chongyi zheng): Do we need multiple updates after initial data collection?
# num_updates = args.init_steps if step == args.init_steps else 1
# for _ in range(num_updates):
# agent.update(replay_buffer, logger, step)
for _ in range(args.num_train_iters):
agent.update(replay_buffer, logger, step)
# agent.train(batch_size=args.batch_size, gradient_steps=args.num_train_iters)
# Take step
next_obs, reward, done, info = env.step(action)
# agent.replay_buffer.add(obs, next_obs, action, reward, done)
# replay_buffer.add(obs, action, reward, next_obs, done)
replay_buffer.add(np.expand_dims(np.asarray(obs), axis=0),
np.expand_dims(np.asarray(next_obs), axis=0),
np.expand_dims(action, axis=0),
np.expand_dims(np.sign(reward), axis=0),
np.expand_dims(done, axis=0))
episode_reward += reward
obs = next_obs
episode_step += 1
class ReservoirOffPolicyAlgorithm(BaseAlgorithm):
"""
The base for Off-Policy algorithms (ex: SAC/TD3)
:param policy: Policy object
:param env: The environment to learn from
(if registered in Gym, can be str. Can be None for loading trained models)
:param policy_base: The base policy used by this method
:param learning_rate: learning rate for the optimizer,
it can be a function of the current progress remaining (from 1 to 0)
:param buffer_size: size of the replay buffer
:param learning_starts: how many steps of the model to collect transitions for before learning starts
:param batch_size: Minibatch size for each gradient update
:param tau: the soft update coefficient ("Polyak update", between 0 and 1)
:param gamma: the discount factor
:param train_freq: Update the model every ``train_freq`` steps. Set to `-1` to disable.
:param gradient_steps: How many gradient steps to do after each rollout
(see ``train_freq`` and ``n_episodes_rollout``)
Set to ``-1`` means to do as many gradient steps as steps done in the environment
during the rollout.
:param n_episodes_rollout: Update the model every ``n_episodes_rollout`` episodes.
Note that this cannot be used at the same time as ``train_freq``. Set to `-1` to disable.
:param action_noise: the action noise type (None by default), this can help
for hard exploration problem. Cf common.noise for the different action noise type.
:param optimize_memory_usage: Enable a memory efficient variant of the replay buffer
at a cost of more complexity.
See https://github.com/DLR-RM/stable-baselines3/issues/37#issuecomment-637501195
:param policy_kwargs: Additional arguments to be passed to the policy on creation
:param tensorboard_log: the log location for tensorboard (if None, no logging)
:param verbose: The verbosity level: 0 none, 1 training information, 2 debug
:param device: Device on which the code should run.
By default, it will try to use a Cuda compatible device and fallback to cpu
if it is not possible.
:param support_multi_env: Whether the algorithm supports training
with multiple environments (as in A2C)
:param create_eval_env: Whether to create a second environment that will be
used for evaluating the agent periodically. (Only available when passing string for the environment)
:param monitor_wrapper: When creating an environment, whether to wrap it
or not in a Monitor wrapper.
:param seed: Seed for the pseudo random generators
:param use_sde: Whether to use State Dependent Exploration (SDE)
instead of action noise exploration (default: False)
:param sde_sample_freq: Sample a new noise matrix every n steps when using gSDE
Default: -1 (only sample at the beginning of the rollout)
:param use_sde_at_warmup: Whether to use gSDE instead of uniform sampling
during the warm up phase (before learning starts)
:param sde_support: Whether the model support gSDE or not
:param remove_time_limit_termination: Remove terminations (dones) that are due to time limit.
See https://github.com/hill-a/stable-baselines/issues/863
"""
def __init__(
self,
policy: Type[BasePolicy],
env: Union[GymEnv, str],
policy_base: Type[BasePolicy],
learning_rate: Union[float, Callable],
buffer_size: int = int(1e6),
learning_starts: int = 100,
batch_size: int = 256,
tau: float = 0.005,
gamma: float = 0.99,
train_freq: int = 1,
gradient_steps: int = 1,
n_episodes_rollout: int = -1,
action_noise: Optional[ActionNoise] = None,
optimize_memory_usage: bool = False,
policy_kwargs: Dict[str, Any] = None,
tensorboard_log: Optional[str] = None,
verbose: int = 0,
device: Union[th.device, str] = "auto",
support_multi_env: bool = False,
create_eval_env: bool = False,
monitor_wrapper: bool = True,
seed: Optional[int] = None,
use_sde: bool = False,
sde_sample_freq: int = -1,
use_sde_at_warmup: bool = False,
sde_support: bool = True,
remove_time_limit_termination: bool = False,
):
super(ReservoirOffPolicyAlgorithm, self).__init__(
policy=policy,
env=env,
policy_base=policy_base,
learning_rate=learning_rate,
policy_kwargs=policy_kwargs,
tensorboard_log=tensorboard_log,
verbose=verbose,
device=device,
support_multi_env=support_multi_env,
create_eval_env=create_eval_env,
monitor_wrapper=monitor_wrapper,
seed=seed,
use_sde=use_sde,
sde_sample_freq=sde_sample_freq,
)
self.buffer_size = buffer_size
self.batch_size = batch_size
self.learning_starts = learning_starts
self.tau = tau
self.gamma = gamma
self.train_freq = train_freq
self.gradient_steps = gradient_steps
self.n_episodes_rollout = n_episodes_rollout
self.action_noise = action_noise
self.optimize_memory_usage = optimize_memory_usage
# Remove terminations (dones) that are due to time limit
# see https://github.com/hill-a/stable-baselines/issues/863
self.remove_time_limit_termination = remove_time_limit_termination
if train_freq > 0 and n_episodes_rollout > 0:
warnings.warn(
"You passed a positive value for `train_freq` and `n_episodes_rollout`."
"Please make sure this is intended. "
"The agent will collect data by stepping in the environment "
"until both conditions are true: "
"`number of steps in the env` >= `train_freq` and "
"`number of episodes` > `n_episodes_rollout`"
)
self.actor = None # type: Optional[th.nn.Module]
self.replay_buffer = None # type: Optional[ReservoirBuffer]
self.current_experience_buffer = None # type: Optional[ReplayBuffer]
# Update policy keyword arguments
if sde_support:
self.policy_kwargs["use_sde"] = self.use_sde
# For gSDE only
self.use_sde_at_warmup = use_sde_at_warmup
self.update_env(env, support_multi_env=support_multi_env, create_eval_env=create_eval_env,
monitor_wrapper=monitor_wrapper, is_reservoir_replay=True)
def _setup_model(self) -> None:
self._setup_lr_schedule()
self.set_random_seed(self.seed)
self.replay_buffer = ReservoirBuffer(
self.buffer_size,
self.observation_space,
self.action_space,
self.device,
optimize_memory_usage=self.optimize_memory_usage,
)
self.current_experience_buffer = ReplayBuffer(
1,
self.observation_space,
self.action_space,
self.device,
optimize_memory_usage=False,
)
self.policy = self.policy_class(
self.observation_space,
self.action_space,
self.lr_schedule,
**self.policy_kwargs # pytype:disable=not-instantiable
)
self.policy = self.policy.to(self.device)
def save_replay_buffer(self, path: Union[str, pathlib.Path, io.BufferedIOBase]) -> None:
"""
Save the replay buffer as a pickle file.
:param path: Path to the file where the replay buffer should be saved.
if path is a str or pathlib.Path, the path is automatically created if necessary.
"""
assert self.replay_buffer is not None, "The replay buffer is not defined"
save_to_pkl(path, self.replay_buffer, self.verbose)
def load_replay_buffer(self, path: Union[str, pathlib.Path, io.BufferedIOBase]) -> None:
"""
Load a replay buffer from a pickle file.
:param path: 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"
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`.
"""
# Prevent continuity issue by truncating trajectory
# when using memory efficient replay buffer
# see https://github.com/DLR-RM/stable-baselines3/issues/46
truncate_last_traj = (
self.optimize_memory_usage
and reset_num_timesteps
and self.replay_buffer is not None
and (self.replay_buffer.full or self.replay_buffer.pos > 0)
)
if truncate_last_traj:
warnings.warn(
"The last trajectory in the replay buffer will be truncated, "
"see https://github.com/DLR-RM/stable-baselines3/issues/46."
"You should use `reset_num_timesteps=False` or `optimize_memory_usage=False`"
"to avoid that issue."
)
# Go to the previous index
pos = (self.replay_buffer.pos - 1) % self.replay_buffer.buffer_size
self.replay_buffer.dones[pos] = True
return super()._setup_learn(
total_timesteps, eval_env, callback, eval_freq, n_eval_episodes, log_path, reset_num_timesteps, tb_log_name
)
def learn(
self,
total_timesteps: int,
callback: MaybeCallback = None,
log_interval: int = 4,
eval_env: Optional[GymEnv] = None,
eval_freq: int = -1,
n_eval_episodes: int = 5,
tb_log_name: str = "run",
eval_log_path: Optional[str] = None,
reset_num_timesteps: bool = True,
) -> "ReservoirOffPolicyAlgorithm":
total_timesteps, callback = self._setup_learn(
total_timesteps, eval_env, callback, eval_freq, n_eval_episodes, eval_log_path, reset_num_timesteps,
tb_log_name
)
callback.on_training_start(locals(), globals())
while self.num_timesteps < total_timesteps:
rollout = self.collect_rollouts(
self.env,
n_episodes=self.n_episodes_rollout,
n_steps=self.train_freq,
action_noise=self.action_noise,
callback=callback,
learning_starts=self.learning_starts,
replay_buffer=self.replay_buffer,
log_interval=log_interval,
)
if rollout.continue_training is False:
break
if self.num_timesteps > 0 and self.num_timesteps > self.learning_starts:
# If no `gradient_steps` is specified,
# do as many gradients steps as steps performed during the rollout
gradient_steps = self.gradient_steps if self.gradient_steps > 0 else rollout.episode_timesteps
self.train(batch_size=self.batch_size, gradient_steps=gradient_steps)
callback.on_training_end()
return self
def train(self, gradient_steps: int, batch_size: int) -> None:
"""
Sample the replay buffer and do the updates
(gradient descent and update target networks)
"""
raise NotImplementedError()
def _sample_action(
self, learning_starts: int, action_noise: Optional[ActionNoise] = None
) -> Tuple[np.ndarray, np.ndarray]:
"""
Sample an action according to the exploration policy.
This is either done by sampling the probability distribution of the policy,
or sampling a random action (from a uniform distribution over the action space)
or by adding noise to the deterministic output.
:param action_noise: Action noise that will be used for exploration
Required for deterministic policy (e.g. TD3). This can also be used
in addition to the stochastic policy for SAC.
:param learning_starts: Number of steps before learning for the warm-up phase.
:return: action to take in the environment
and scaled action that will be stored in the replay buffer.
The two differs when the action space is not normalized (bounds are not [-1, 1]).
"""
# Select action randomly or according to policy
if self.num_timesteps < learning_starts and not (self.use_sde and self.use_sde_at_warmup):
# Warmup phase
unscaled_action = np.array([self.action_space.sample()])
else:
# Note: when using continuous actions,
# we assume that the policy uses tanh to scale the action
# We use non-deterministic action in the case of SAC, for TD3, it does not matter
unscaled_action, _ = self.predict(self._last_obs, deterministic=False)
# Rescale the action from [low, high] to [-1, 1]
if isinstance(self.action_space, gym.spaces.Box):
scaled_action = self.policy.scale_action(unscaled_action)
# Add noise to the action (improve exploration)
if action_noise is not None:
scaled_action = np.clip(scaled_action + action_noise(), -1, 1)
# We store the scaled action in the buffer
buffer_action = scaled_action
action = self.policy.unscale_action(scaled_action)
else:
# Discrete case, no need to normalize or clip
buffer_action = unscaled_action
action = buffer_action
return action, buffer_action
def _dump_logs(self) -> None:
"""
Write log.
"""
try:
fps = int(self.num_timesteps / (time.time() - self.start_time))
except ZeroDivisionError:
warnings.warn("fps dump had zero division somehow, storing 0 instead.")
fps = 0
logger.record("time/episodes", self._episode_num, exclude="tensorboard")
if len(self.ep_info_buffer) > 0 and len(self.ep_info_buffer[0]) > 0:
logger.record("rollout/ep_rew_mean", safe_mean([ep_info["r"] for ep_info in self.ep_info_buffer]))
logger.record("rollout/ep_len_mean", safe_mean([ep_info["l"] for ep_info in self.ep_info_buffer]))
logger.record("time/fps", fps)
logger.record("time/time_elapsed", int(time.time() - self.start_time), exclude="tensorboard")
logger.record("time/total timesteps", self.num_timesteps, exclude="tensorboard")
if self.use_sde:
logger.record("train/std", (self.actor.get_std()).mean().item())
if len(self.ep_success_buffer) > 0:
logger.record("rollout/success rate", safe_mean(self.ep_success_buffer))
# Pass the number of timesteps for tensorboard
logger.dump(step=self.num_timesteps)
def _on_step(self) -> None:
"""
Method called after each step in the environment.
It is meant to trigger DQN target network update
but can be used for other purposes
"""
pass
def collect_rollouts(
self,
env: VecEnv,
callback: BaseCallback,
n_episodes: int = 1,
n_steps: int = -1,
action_noise: Optional[ActionNoise] = None,
learning_starts: int = 0,
replay_buffer: Optional[ReservoirBuffer] = None,
log_interval: Optional[int] = None,
) -> RolloutReturn:
"""
Collect experiences and store them into a ReplayBuffer.
:param env: The training environment
:param callback: Callback that will be called at each step
(and at the beginning and end of the rollout)
:param n_episodes: Number of episodes to use to collect rollout data
You can also specify a ``n_steps`` instead
:param n_steps: Number of steps to use to collect rollout data
You can also specify a ``n_episodes`` instead.
:param action_noise: Action noise that will be used for exploration
Required for deterministic policy (e.g. TD3). This can also be used
in addition to the stochastic policy for SAC.
:param learning_starts: Number of steps before learning for the warm-up phase.
:param replay_buffer:
:param log_interval: Log data every ``log_interval`` episodes
:return:
"""
episode_rewards, total_timesteps = [], []
total_steps, total_episodes = 0, 0
assert isinstance(env, VecEnv), "You must pass a VecEnv"
assert env.num_envs == 1, "OffPolicyAlgorithm only support single environment"
if self.use_sde:
self.actor.reset_noise()
callback.on_rollout_start()
continue_training = True
while total_steps < n_steps or total_episodes < n_episodes:
done = False
episode_reward, episode_timesteps = 0.0, 0
while not done:
if self.use_sde and self.sde_sample_freq > 0 and total_steps % self.sde_sample_freq == 0:
# Sample a new noise matrix
self.actor.reset_noise()
# Select action randomly or according to policy
action, buffer_action = self._sample_action(learning_starts, action_noise)
# Rescale and perform action
new_obs, reward, done, infos = env.step(action)
self.num_timesteps += 1
episode_timesteps += 1
total_steps += 1
# Give access to local variables
callback.update_locals(locals())
# Only stop training if return value is False, not when it is None.
if callback.on_step() is False:
return RolloutReturn(0.0, total_steps, total_episodes, continue_training=False)
episode_reward += reward
# Retrieve reward and episode length if using Monitor wrapper
self._update_info_buffer(infos, done)
# Store data in replay buffer
if replay_buffer is not None:
# Store only the unnormalized version
if self._vec_normalize_env is not None:
new_obs_ = self._vec_normalize_env.get_original_obs()
reward_ = self._vec_normalize_env.get_original_reward()
else:
# Avoid changing the original ones
self._last_original_obs, new_obs_, reward_ = self._last_obs, new_obs, reward
replay_buffer.add(self._last_original_obs, new_obs_, buffer_action, reward_, done,
self.num_timesteps - 1)
self.current_experience_buffer.add(self._last_original_obs, new_obs_, buffer_action, reward_, done)
self._last_obs = new_obs
# Save the unnormalized observation
if self._vec_normalize_env is not None:
self._last_original_obs = new_obs_
self._update_current_progress_remaining(self.num_timesteps, self._total_timesteps)
# For DQN, check if the target network should be updated
# and update the exploration schedule
# For SAC/TD3, the update is done as the same time as the gradient update
# see https://github.com/hill-a/stable-baselines/issues/900
self._on_step()
if 0 < n_steps <= total_steps:
break
if done:
total_episodes += 1
self._episode_num += 1
episode_rewards.append(episode_reward)
total_timesteps.append(episode_timesteps)
if action_noise is not None:
action_noise.reset()
# Log training infos
if log_interval is not None and self._episode_num % log_interval == 0:
self._dump_logs()
mean_reward = np.mean(episode_rewards) if total_episodes > 0 else 0.0
callback.on_rollout_end()
return RolloutReturn(mean_reward, total_steps, total_episodes, continue_training)
def update_env(self, env, support_multi_env: bool = False, eval_env: Optional[GymEnv] = None,
monitor_wrapper: bool = True, is_reservoir_replay: bool = True, **kwargs):
"""
Replace current env with new env.
:param env: Gym environment (activated, not a string).
:param support_multi_env: Whether the algorithm supports training
with multiple environments (as in A2C)
:param eval_env: Environment to use for evaluation.
:param monitor_wrapper: When creating an environment, whether to wrap it
or not in a Monitor wrapper.
:param is_reservoir_replay: Whether experience replay is normal or reservoir
:param kwargs: Does nothing, just so more arguments can pass without method failing
:return:
"""
if self.replay_buffer is not None:
if is_reservoir_replay:
self.replay_buffer.is_reservoir = True
else:
self.replay_buffer.is_reservoir = False
if env is not None:
if eval_env is not None:
self.eval_env = eval_env
if monitor_wrapper:
self.eval_env = Monitor(self.eval_env, filename=None)
if monitor_wrapper:
env = Monitor(env, filename=None)
env = self._wrap_env(env, self.verbose)
self.observation_space = env.observation_space
self.action_space = env.action_space
self.n_envs = env.num_envs
self.env = env
if not support_multi_env and self.n_envs > 1:
raise ValueError(
"Error: the model does not support multiple envs; it requires " "a single vectorized environment."
)
def add_memories_from_another_replay_mem(self, another_replay_mem: ReplayBuffer):
for i in range(another_replay_mem.buffer_size):
self.replay_buffer.add(
obs=another_replay_mem.observations[i],
next_obs=another_replay_mem.next_observations[i],
action=another_replay_mem.actions[i],
reward=another_replay_mem.rewards[i],
done=another_replay_mem.dones[i],
experience_index=0, # doesn't matter
)
