| 56 | """Holder for all normalizers in dict or non-dict obs spaces.""" |
| 57 | |
| 58 | def __init__(self, obs_space: spaces.Dict, cfg: ObsNormalizerConfig): |
| 59 | super().__init__() |
| 60 | self.cfg: ObsNormalizerConfig = cfg |
| 61 | if isinstance(cfg.normalizers, dict): |
| 62 | if not cfg.allow_mismatching_keys: |
| 63 | if set(obs_space.keys()) != set(cfg.normalizers.keys()): |
| 64 | raise ValueError( |
| 65 | f"ObsNormalizerConfig keys {set(cfg.normalizers.keys())} do not match observation space keys {set(obs_space.keys())}. " |
| 66 | "Set allow_mismatching_keys=True to ignore this check." |
| 67 | ) |
| 68 | self._normalizers = nn.ModuleDict({key: cfg.normalizers[key].build(obs_space[key]) for key in cfg.normalizers.keys()}) |
| 69 | else: |
| 70 | self._normalizers = cfg.normalizers.build(obs_space) |
| 71 | |
| 72 | def forward(self, x: dict[str, torch.Tensor] | torch.Tensor) -> dict[str, torch.Tensor] | torch.Tensor: |
| 73 | # TODO is this is-instance check bad for performance? |