| import torch |
| import torch.nn as nn |
|
|
| """ |
| updates statistic from a full data |
| |
| Memory optimization note: If you encounter CUDA out of memory errors, consider setting: |
| export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True |
| This can help reduce memory fragmentation. |
| """ |
|
|
|
|
| class RunningMeanStd(nn.Module): |
|
|
| def __init__(self, insize, epsilon=1e-05, per_channel=False, norm_only=False): |
| super().__init__() |
| print("RunningMeanStd: ", insize) |
| self.insize = insize |
| self.mean_size = insize[0] |
| self.epsilon = epsilon |
|
|
| self.norm_only = norm_only |
| self.per_channel = per_channel |
| if per_channel: |
| if len(self.insize) == 3: |
| self.axis = [0, 2, 3] |
| if len(self.insize) == 2: |
| self.axis = [0, 2] |
| if len(self.insize) == 1: |
| self.axis = [0] |
| in_size = self.insize[0] |
| else: |
| self.axis = [0] |
| in_size = insize |
|
|
| self.register_buffer("running_mean", torch.zeros(in_size, dtype=torch.float32)) |
| self.register_buffer("running_var", torch.ones(in_size, dtype=torch.float32)) |
| self.register_buffer("count", torch.ones((), dtype=torch.float32)) |
|
|
| self.frozen = False |
| self.frozen_partial = False |
|
|
| def freeze(self): |
| self.frozen = True |
|
|
| def unfreeze(self): |
| self.frozen = False |
|
|
| def freeze_partial(self, diff): |
| self.frozen_partial = True |
| self.diff = diff |
|
|
| def sync_across_gpus(self, accelerator): |
| if accelerator.num_processes <= 1: |
| return |
| |
| |
| flat_stats = torch.cat( |
| [ |
| self.running_mean.flatten(), |
| self.running_var.flatten(), |
| ] |
| ) |
|
|
| |
| gathered = accelerator.gather(flat_stats[None]) |
|
|
| |
| world_size = gathered.shape[0] |
| D = self.running_mean.numel() |
|
|
| means_ = gathered[:, :D].reshape(world_size, *self.running_mean.shape).mean(dim=0) |
| vars_ = gathered[:, D : 2 * D].reshape(world_size, *self.running_var.shape).mean(dim=0) |
|
|
| |
| self.running_mean.copy_(means_) |
| self.running_var.copy_(vars_) |
|
|
| def _update_mean_var_count_from_moments( |
| self, mean, var, count, batch_mean, batch_var, batch_count |
| ): |
| delta = batch_mean - mean |
| tot_count = count + batch_count |
|
|
| new_mean = mean + delta * batch_count / tot_count |
| m_a = var * count |
| m_b = batch_var * batch_count |
| M2 = m_a + m_b + delta**2 * count * batch_count / tot_count |
| new_var = M2 / tot_count |
| new_count = tot_count |
| return new_mean, new_var, new_count |
|
|
| def forward(self, input, unnorm=False): |
| |
| input_shape = input.shape |
| if len(input.shape) == 3: |
| input = input.reshape(-1, input_shape[-1]) |
|
|
| if self.per_channel: |
| if len(self.insize) == 3: |
| |
| current_mean = self.running_mean.view([1, self.insize[0], 1, 1]) |
| current_var = self.running_var.view([1, self.insize[0], 1, 1]) |
| elif len(self.insize) == 2: |
| current_mean = self.running_mean.view([1, self.insize[0], 1]) |
| current_var = self.running_var.view([1, self.insize[0], 1]) |
| elif len(self.insize) == 1: |
| current_mean = self.running_mean.view([1, self.insize[0]]) |
| current_var = self.running_var.view([1, self.insize[0]]) |
|
|
| else: |
| current_mean = self.running_mean |
| current_var = self.running_var |
| |
|
|
| if unnorm: |
| y = torch.clamp(input, min=-5.0, max=5.0) |
| y = torch.sqrt(current_var + self.epsilon) * y + current_mean |
| else: |
| if self.norm_only: |
| y = input / torch.sqrt(current_var + self.epsilon) |
| else: |
| |
| y = input - current_mean |
| y = y / torch.sqrt(current_var + self.epsilon) |
| y = torch.clamp(y, min=-5.0, max=5.0) |
|
|
| |
| if self.training and not self.frozen: |
| mean = input.mean(self.axis) |
| var = input.var(self.axis) |
|
|
| new_mean, new_var, new_count = self._update_mean_var_count_from_moments( |
| self.running_mean, self.running_var, self.count, mean, var, input.size()[0] |
| ) |
| if self.frozen_partial: |
| |
| self.running_mean[-self.diff :], self.running_var[-self.diff :], self.count = ( |
| new_mean[-self.diff :], |
| new_var[-self.diff :], |
| new_count, |
| ) |
| else: |
| self.running_mean, self.running_var, self.count = new_mean, new_var, new_count |
|
|
| if len(input_shape) == 3: |
| y = y.view(input_shape) |
|
|
| return y |
|
|
|
|
| class RunningMeanStdObs(nn.Module): |
|
|
| def __init__(self, insize, epsilon=1e-05, per_channel=False, norm_only=False): |
| assert isinstance(insize, dict) |
| super().__init__() |
| self.running_mean_std = nn.ModuleDict( |
| {k: RunningMeanStd(v, epsilon, per_channel, norm_only) for k, v in insize.items()} |
| ) |
|
|
| def forward(self, input, unnorm=False): |
| res = {k: self.running_mean_std[k](v, unnorm) for k, v in input.items()} |
| return res |
|
|
|
|
| from collections.abc import Sequence |
| from copy import deepcopy |
|
|
|
|
| class VecNorm(nn.Module): |
| """Simple running normalization for observations. |
| |
| Keeps track of running mean and variance to normalize observations on-the-fly. |
| |
| Args: |
| obs_keys: List of observation keys to normalize |
| decay: Decay rate for running statistics (default: 0.99) |
| eps: Small constant for numerical stability (default: 1e-4) |
| device: Device to store statistics on |
| """ |
|
|
| def __init__( |
| self, |
| obs_keys: Sequence[str] | None, |
| decay: float = 0.9999, |
| eps: float = 1e-4, |
| ): |
| super().__init__() |
| self.obs_keys = obs_keys |
| self.decay = decay |
| self.eps = eps |
|
|
| |
| self.sum: dict[str, nn.Parameter] = nn.ParameterDict({}) |
| self.ssq: dict[str, nn.Parameter] = nn.ParameterDict({}) |
| self.cnt: dict[str, nn.Parameter] = nn.ParameterDict({}) |
| self.mean: dict[str, nn.Parameter] = nn.ParameterDict({}) |
| self.var: dict[str, nn.Parameter] = nn.ParameterDict({}) |
|
|
| self.initialized = False |
| self.frozen = False |
|
|
| def init_stats(self, obs_dict: dict[str, torch.Tensor]): |
| """Initialize running statistics based on observation shapes.""" |
| if self.obs_keys is None: |
| self.obs_keys = list(obs_dict.keys()) |
| for key in self.obs_keys: |
| if key in obs_dict: |
| v = obs_dict[key] |
| self.sum[key] = nn.Parameter(torch.zeros_like(v[0]), requires_grad=False) |
| self.ssq[key] = nn.Parameter(torch.zeros_like(v[0]), requires_grad=False) |
| self.cnt[key] = nn.Parameter(torch.zeros_like(v[0]), requires_grad=False) |
| self.mean[key] = nn.Parameter(torch.zeros_like(v[0]), requires_grad=False) |
| self.var[key] = nn.Parameter(torch.ones_like(v[0]), requires_grad=False) |
| self.initialized = True |
|
|
| def update(self, obs_dict: dict[str, torch.Tensor]): |
| """Update running statistics with new observations.""" |
| if not self.initialized: |
| self.init_stats(obs_dict) |
|
|
| if self.frozen: |
| return |
|
|
| for key in self.obs_keys: |
| if key not in obs_dict: |
| continue |
|
|
| x = obs_dict[key] |
| sum_x = x.sum(dim=0) |
| ssq_x = (x**2).sum(dim=0) |
| cnt_x = x.shape[0] |
|
|
| self.sum[key] = self.sum[key] * self.decay + sum_x |
| self.ssq[key] = self.ssq[key] * self.decay + ssq_x |
| self.cnt[key] = self.cnt[key] * self.decay + cnt_x |
|
|
| self.mean[key] = self.sum[key] / self.cnt[key] |
| self.var[key] = (self.ssq[key] / self.cnt[key] - self.mean[key] ** 2).clamp( |
| min=self.eps |
| ) |
|
|
| def normalize(self, obs_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: |
| """Normalize observations using running statistics.""" |
| if not self.initialized: |
| self.init_stats(obs_dict) |
| return obs_dict.copy() |
|
|
| normalized = obs_dict.copy() |
| for key in self.obs_keys: |
| normalized[key] = (normalized[key] - self.mean[key]) / self.var[key].sqrt() |
|
|
| return normalized |
|
|
| def denormalize(self, obs_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: |
| """Denormalize observations using running statistics.""" |
| if not self.initialized: |
| self.init_stats(obs_dict) |
| return obs_dict.copy() |
|
|
| denormalized = obs_dict.copy() |
| for key in self.obs_keys: |
| denormalized[key] = denormalized[key] * self.var[key].sqrt() + self.mean[key] |
|
|
| return denormalized |
|
|
| def freeze(self): |
| """Freeze running statistics updates.""" |
| self.frozen = True |
|
|
| def unfreeze(self): |
| """Unfreeze running statistics updates.""" |
| self.frozen = False |
|
|
| def get_stats(self): |
| """Get current running statistics.""" |
| return { |
| "sum": {key: param.data.clone() for key, param in self.sum.items()}, |
| "ssq": {key: param.data.clone() for key, param in self.ssq.items()}, |
| "cnt": {key: param.data.clone() for key, param in self.cnt.items()}, |
| "mean": {key: param.data.clone() for key, param in self.mean.items()}, |
| "var": {key: param.data.clone() for key, param in self.var.items()}, |
| } |
|
|
| def load_stats(self, stats): |
| """Load running statistics.""" |
| for key in stats["sum"]: |
| self.sum[key] = nn.Parameter(deepcopy(stats["sum"][key]), requires_grad=False) |
| self.ssq[key] = nn.Parameter(deepcopy(stats["ssq"][key]), requires_grad=False) |
| self.cnt[key] = nn.Parameter(deepcopy(stats["cnt"][key]), requires_grad=False) |
| self.mean[key] = nn.Parameter(deepcopy(stats["mean"][key]), requires_grad=False) |
| self.var[key] = nn.Parameter(deepcopy(stats["var"][key]), requires_grad=False) |
| self.initialized = True |
|
|
