| import torch |
| import torch.nn as nn |
|
|
|
|
| class BatchNormNormalizer(nn.Module): |
| def __init__(self, insize, epsilon=1e-05, per_channel=False, norm_only=False): |
| super().__init__() |
| assert len(insize) == 1, "BatchNormNormalizer only supports 1D observation spaces" |
| self._normalizer = nn.SyncBatchNorm(num_features=insize[0], affine=False) |
|
|
| @property |
| def num_features(self): |
| return self._normalizer.num_features |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| input_shape = x.shape |
| if len(x.shape) == 3: |
| x = x.reshape(-1, x.shape[-1]) |
|
|
| x = self._normalizer(x) |
| x = x.view(input_shape) |
| return x |
|
|
| def update(self, input: torch.Tensor): |
| """Update running stats from input. No-op in eval mode. |
| |
| Calls SyncBatchNorm.forward() for its side effect of updating |
| running_mean/running_var (and multi-GPU sync). Output is discarded. |
| """ |
| if not self.training: |
| return |
| if len(input.shape) == 3: |
| input = input.reshape(-1, input.shape[-1]) |
| with torch.no_grad(): |
| self._normalizer(input) |
|
|
| def normalize(self, input: torch.Tensor) -> torch.Tensor: |
| """Normalize using current running stats without updating them.""" |
| input_shape = input.shape |
| if len(input.shape) == 3: |
| input = input.reshape(-1, input_shape[-1]) |
| y = (input - self._normalizer.running_mean) / torch.sqrt( |
| self._normalizer.running_var + self._normalizer.eps |
| ) |
| y = torch.clamp(y, min=-5.0, max=5.0) |
| if len(input_shape) == 3: |
| y = y.view(input_shape) |
| return y |
|
|
