| """H2 adaptation of residual TD3 networks from sim commit 4fa48dd. |
| |
| Layer layout and zero actor initialization are preserved; dimensions come |
| from the H2 contract. Normalizer checkpoints retain full running statistics. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import numpy as np |
| import torch |
| from torch import nn |
|
|
| from .data import ACTOR_INPUT_DIM, RESIDUAL_ACTION_DIM |
|
|
|
|
| class Actor(nn.Module): |
| def __init__(self) -> None: |
| super().__init__() |
| self.net = nn.Sequential( |
| nn.Linear(ACTOR_INPUT_DIM, 256), |
| nn.LayerNorm(256), |
| nn.ReLU(), |
| nn.Linear(256, 256), |
| nn.LayerNorm(256), |
| nn.ReLU(), |
| nn.Linear(256, RESIDUAL_ACTION_DIM), |
| nn.Tanh(), |
| ) |
| output = self.net[-2] |
| nn.init.zeros_(output.weight) |
| nn.init.zeros_(output.bias) |
|
|
| def forward(self, inputs: torch.Tensor) -> torch.Tensor: |
| return self.net(inputs) |
|
|
|
|
| class Critic(nn.Module): |
| def __init__(self) -> None: |
| super().__init__() |
| self.net = nn.Sequential( |
| nn.Linear(ACTOR_INPUT_DIM, 256), |
| nn.LayerNorm(256), |
| nn.ReLU(), |
| nn.Linear(256, 256), |
| nn.LayerNorm(256), |
| nn.ReLU(), |
| nn.Linear(256, 1), |
| ) |
|
|
| def forward(self, inputs: torch.Tensor) -> torch.Tensor: |
| return self.net(inputs) |
|
|
|
|
| def parameter_count(module: nn.Module) -> int: |
| return sum(parameter.numel() for parameter in module.parameters()) |
|
|
|
|
| def actor_parameter_count() -> int: |
| with torch.random.fork_rng(devices=[]): |
| return parameter_count(Actor()) |
|
|
|
|
| def flatten_parameters(module: nn.Module) -> np.ndarray: |
| values = [parameter.detach().cpu().reshape(-1) for parameter in module.parameters()] |
| return torch.cat(values).numpy().astype(np.float32, copy=True) |
|
|
|
|
| def load_flat_parameters(module: nn.Module, values: np.ndarray) -> None: |
| flat = np.asarray(values, dtype=np.float32).reshape(-1) |
| if flat.size != parameter_count(module) or not np.isfinite(flat).all(): |
| raise ValueError("actor parameter payload is invalid") |
| offset = 0 |
| with torch.no_grad(): |
| for parameter in module.parameters(): |
| count = parameter.numel() |
| source = torch.from_numpy(flat[offset : offset + count]).reshape(parameter.shape) |
| parameter.copy_(source.to(device=parameter.device, dtype=parameter.dtype)) |
| offset += count |
|
|
|
|
| class RunningNormalizer: |
| def __init__(self, size: int, *, min_std: float, clip: float) -> None: |
| self.size = int(size) |
| self.min_std = float(min_std) |
| self.clip = float(clip) |
| self.count = 0 |
| self.mean = np.zeros(self.size, dtype=np.float64) |
| self.m2 = np.zeros(self.size, dtype=np.float64) |
| self.frozen = False |
|
|
| def update(self, values: np.ndarray) -> None: |
| if self.frozen: |
| return |
| value = np.asarray(values, dtype=np.float64).reshape(self.size) |
| self.count += 1 |
| delta = value - self.mean |
| self.mean += delta / self.count |
| self.m2 += delta * (value - self.mean) |
|
|
| def freeze(self) -> None: |
| self.frozen = True |
|
|
| @property |
| def std(self) -> np.ndarray: |
| if self.count < 2: |
| variance = np.ones(self.size, dtype=np.float64) |
| else: |
| variance = self.m2 / self.count |
| return np.maximum(np.sqrt(variance), self.min_std).astype(np.float32) |
|
|
| def normalize_tensor(self, values: torch.Tensor) -> torch.Tensor: |
| mean = torch.as_tensor(self.mean, dtype=values.dtype, device=values.device) |
| std = torch.as_tensor(self.std, dtype=values.dtype, device=values.device) |
| return torch.clamp((values - mean) / std, -self.clip, self.clip) |
|
|
| def state_dict(self) -> dict: |
| return { |
| "size": self.size, |
| "min_std": self.min_std, |
| "clip": self.clip, |
| "count": self.count, |
| "mean": self.mean.copy(), |
| "std": self.std, |
| "m2": self.m2.copy(), |
| "frozen": self.frozen, |
| } |
|
|
