File size: 1,549 Bytes
cc86714 | 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 | 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())
|