File size: 4,107 Bytes
d8290eb | 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 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 | """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,
}
|