File size: 1,729 Bytes
0651b21 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 | 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: # do nothing if in evaluation mode
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
|