| from __future__ import annotations |
|
|
| import torch |
| from torch import nn |
|
|
|
|
| class NeuralVectorField(nn.Module): |
| def __init__(self) -> None: |
| super().__init__() |
| self.network = nn.Sequential( |
| nn.Linear(2, 32), |
| nn.Tanh(), |
| nn.Linear(32, 32), |
| nn.Tanh(), |
| nn.Linear(32, 2), |
| ) |
|
|
| def forward(self, state: torch.Tensor) -> torch.Tensor: |
| return self.network(state) |
|
|
| def step(self, state: torch.Tensor, delta_time: torch.Tensor) -> torch.Tensor: |
| k1 = self(state) |
| k2 = self(state + 0.5 * delta_time * k1) |
| k3 = self(state + 0.5 * delta_time * k2) |
| k4 = self(state + delta_time * k3) |
| return state + delta_time * (k1 + 2 * k2 + 2 * k3 + k4) / 6 |
|
|
|
|
| class CalibratedDiscreteTransition(nn.Module): |
| def __init__(self) -> None: |
| super().__init__() |
| self.network = nn.Sequential( |
| nn.Linear(3, 32), |
| nn.Tanh(), |
| nn.Linear(32, 31), |
| nn.Tanh(), |
| nn.Linear(31, 2), |
| ) |
| self.log_scale = nn.Parameter(torch.zeros(2)) |
| self.global_log_scale = nn.Parameter(torch.zeros(1)) |
|
|
| def forward(self, state: torch.Tensor, delta_time: torch.Tensor) -> torch.Tensor: |
| inputs = torch.cat([state, delta_time], dim=1) |
| delta = self.network(inputs) * (self.log_scale + self.global_log_scale).exp() |
| return state + delta |
|
|
|
|
| def parameter_count(model: nn.Module) -> int: |
| return sum(parameter.numel() for parameter in model.parameters()) |
|
|
