| """State and delta normalization for the flow expert. |
| |
| State layout per object: [px, py, vx, vy, angle, angvel]. |
| States are affine-normalized; deltas (flow targets) are scaled to ~unit magnitude. |
| """ |
|
|
| import math |
|
|
| import numpy as np |
| import torch |
|
|
|
|
| def wrap_angle(a): |
| """Wrap angle(s) to [-pi, pi). Works on numpy arrays and torch tensors.""" |
| if isinstance(a, torch.Tensor): |
| return torch.remainder(a + math.pi, 2 * math.pi) - math.pi |
| return np.mod(a + math.pi, 2 * math.pi) - math.pi |
|
|
|
|
| class StateNormalizer: |
| def __init__(self, state_offset, state_scale, delta_scale): |
| self.state_offset = np.asarray(state_offset, dtype=np.float32) |
| self.state_scale = np.asarray(state_scale, dtype=np.float32) |
| self.delta_scale = np.asarray(delta_scale, dtype=np.float32) |
|
|
| def _cast(self, ref, arr): |
| if isinstance(ref, torch.Tensor): |
| return torch.as_tensor(arr, dtype=ref.dtype, device=ref.device) |
| return arr |
|
|
| def norm_state(self, s): |
| return (s - self._cast(s, self.state_offset)) / self._cast(s, self.state_scale) |
|
|
| def denorm_state(self, s): |
| return s * self._cast(s, self.state_scale) + self._cast(s, self.state_offset) |
|
|
| def norm_delta(self, d): |
| return d / self._cast(d, self.delta_scale) |
|
|
| def denorm_delta(self, d): |
| return d * self._cast(d, self.delta_scale) |
|
|
| @staticmethod |
| def compute_delta(cur, nxt): |
| """next - cur with the angle component wrapped to [-pi, pi).""" |
| d = nxt - cur |
| d[..., 4] = wrap_angle(d[..., 4]) |
| return d |
|
|
| def state_dict(self): |
| return { |
| "state_offset": self.state_offset.tolist(), |
| "state_scale": self.state_scale.tolist(), |
| "delta_scale": self.delta_scale.tolist(), |
| } |
|
|
| @classmethod |
| def from_config(cls, cfg): |
| return cls(cfg.state_offset, cfg.state_scale, cfg.delta_scale) |
|
|
