diff --git a/sb3_trl/misc/distTools.py b/sb3_trl/misc/distTools.py new file mode 100644 index 0000000..0526850 --- /dev/null +++ b/sb3_trl/misc/distTools.py @@ -0,0 +1,46 @@ +import torch as th + +from stable_baselines3.common.distributions import Distribution as SB3_Distribution + + +def get_mean_and_chol(p): + if isinstance(p, th.distributions.Normal): + return p.mean, p.stddev + elif isinstance(p, th.distributions.MultivariateNormal): + return p.mean, p.scale_tril + elif isinstance(p, SB3_Distribution): + return get_mean_and_chol(p.distribution) + else: + raise Exception('Dist-Type not implemented') + + +def get_cov(p): + if isinstance(p, th.distributions.Normal): + return th.diag(p.variance) + elif isinstance(p, th.distributions.MultivariateNormal): + return p.covariance_matrix + elif isinstance(p, SB3_Distribution): + return get_cov(p.distribution) + else: + raise Exception('Dist-Type not implemented') + + +def new_dist_like(orig_p, mean, chol): + if isinstance(orig_p, th.distributions.Normal): + return th.distributions.Normal(mean, chol) + elif isinstance(orig_p, th.distributions.MultivariateNormal): + return th.distributions.MultivariateNormal(mean, scale_tril=chol) + elif isinstance(orig_p, SB3_Distribution): + p = orig_p.distribution + if isinstance(p, th.distributions.Normal): + p_out = orig_p.__class__(orig_p.action_dim) + p_out.distribution = th.distributions.Normal(mean, chol) + elif isinstance(p, th.distributions.MultivariateNormal): + p_out = orig_p.__class__(orig_p.action_dim) + p_out.distribution = th.distributions.MultivariateNormal( + mean, scale_tril=chol) + else: + raise Exception('Dist-Type not implemented (of sb3 dist)') + return p_out + else: + raise Exception('Dist-Type not implemented') diff --git a/sb3_trl/misc/norm.py b/sb3_trl/misc/norm.py new file mode 100644 index 0000000..4a9cbf2 --- /dev/null +++ b/sb3_trl/misc/norm.py @@ -0,0 +1,12 @@ +import torch as th +from torch.distributions.multivariate_normal import _batch_mahalanobis + + +def mahalanobis_blub(u, v, std): + delta = u - v + return th.triangular_solve(delta, std, upper=False)[0].pow(2).sum([-2, -1]) + + +def mahalanobis(u, v, cov): + delta = u - v + return _batch_mahalanobis(cov, delta) diff --git a/sb3_trl/misc/rollout_buffer.py b/sb3_trl/misc/rollout_buffer.py new file mode 100644 index 0000000..fbd49f2 --- /dev/null +++ b/sb3_trl/misc/rollout_buffer.py @@ -0,0 +1,102 @@ +from typing import Any, Dict, Optional, Type, Union, NamedTuple + +import numpy as np +import torch as th +from gym import spaces + +from stable_baselines3.common.buffers import RolloutBuffer +from stable_baselines3.common.vec_env import VecNormalize + + +class GaussianRolloutBufferSamples(NamedTuple): + observations: th.Tensor + actions: th.Tensor + old_values: th.Tensor + old_log_prob: th.Tensor + advantages: th.Tensor + returns: th.Tensor + means: th.Tensor + stds: th.Tensor + + +class GaussianRolloutBuffer(RolloutBuffer): + def __init__( + self, + buffer_size: int, + observation_space: spaces.Space, + action_space: spaces.Space, + device: Union[th.device, str] = "cpu", + gae_lambda: float = 1, + gamma: float = 0.99, + n_envs: int = 1, + ): + + super().__init__(buffer_size, observation_space, action_space, + device, n_envs=n_envs, gae_lambda=gae_lambda, gamma=gamma) + self.means, self.stds = None, None + + def reset(self) -> None: + self.means = np.zeros( + (self.buffer_size, self.n_envs) + self.action_space.shape, dtype=np.float32) + self.stds = np.zeros( + # (self.buffer_size, self.n_envs) + self.action_space.shape + self.action_space.shape, dtype=np.float32) + (self.buffer_size, self.n_envs) + self.action_space.shape, dtype=np.float32) + super().reset() + + def add( + self, + obs: np.ndarray, + action: np.ndarray, + reward: np.ndarray, + episode_start: np.ndarray, + value: th.Tensor, + log_prob: th.Tensor, + mean: th.Tensor, + std: th.Tensor, + ) -> None: + """ + :param obs: Observation + :param action: Action + :param reward: + :param episode_start: Start of episode signal. + :param value: estimated value of the current state + following the current policy. + :param log_prob: log probability of the action + following the current policy. + :param mean: Foo + :param std: Bar + """ + + if len(log_prob.shape) == 0: + # Reshape 0-d tensor to avoid error + log_prob = log_prob.reshape(-1, 1) + + # Reshape needed when using multiple envs with discrete observations + # as numpy cannot broadcast (n_discrete,) to (n_discrete, 1) + if isinstance(self.observation_space, spaces.Discrete): + obs = obs.reshape((self.n_envs,) + self.obs_shape) + + self.observations[self.pos] = np.array(obs).copy() + self.actions[self.pos] = np.array(action).copy() + self.rewards[self.pos] = np.array(reward).copy() + self.episode_starts[self.pos] = np.array(episode_start).copy() + self.values[self.pos] = value.clone().cpu().numpy().flatten() + self.log_probs[self.pos] = log_prob.clone().cpu().numpy() + self.means[self.pos] = mean.clone().cpu().numpy() + self.stds[self.pos] = std.clone().cpu().numpy() + self.pos += 1 + if self.pos == self.buffer_size: + self.full = True + + def _get_samples(self, batch_inds: np.ndarray, env: Optional[VecNormalize] = None) -> GaussianRolloutBufferSamples: + data = ( + self.observations[batch_inds], + self.actions[batch_inds], + self.values[batch_inds].flatten(), + self.log_probs[batch_inds].flatten(), + self.advantages[batch_inds].flatten(), + self.returns[batch_inds].flatten(), + self.means[batch_inds].reshape((len(batch_inds), -1)), + self.stds[batch_inds].reshape((len(batch_inds), -1)), + ) + return GaussianRolloutBufferSamples(*tuple(map(self.to_torch, data)))