176 lines
5.9 KiB
Python
176 lines
5.9 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.optim as optim
|
|
from torchdiffeq import odeint
|
|
import numpy as np
|
|
import matplotlib.pyplot as plt
|
|
|
|
# ----------------------------------------------------------------
|
|
# 1) 3D Controller: [theta, omega, alpha] -> torque
|
|
# ----------------------------------------------------------------
|
|
class PendulumController3D(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.net = nn.Sequential(
|
|
nn.Linear(3, 64),
|
|
nn.ReLU(),
|
|
nn.Linear(64, 64),
|
|
nn.ReLU(),
|
|
nn.Linear(64, 1)
|
|
)
|
|
|
|
def forward(self, x_3d):
|
|
"""
|
|
x_4d: shape (batch_size, 4) => [theta, cos(theta), omega, alpha].
|
|
Returns shape: (batch_size, 1) => torque.
|
|
"""
|
|
raw_torque = self.net(x_3d)
|
|
clamped_torque = torch.clamp(raw_torque, -250, 250) # Clamp torque within [-250, 250]
|
|
return clamped_torque
|
|
|
|
|
|
# ----------------------------------------------------------------
|
|
# 2) Define ODE System Using `odeint`
|
|
# ----------------------------------------------------------------
|
|
m = 10.0
|
|
g = 9.81
|
|
R = 1.0
|
|
|
|
class PendulumDynamics3D(nn.Module):
|
|
"""
|
|
Defines the ODE system for [theta, omega, alpha] with torque tracking.
|
|
"""
|
|
|
|
def __init__(self, controller):
|
|
super().__init__()
|
|
self.controller = controller
|
|
|
|
def forward(self, t, state):
|
|
"""
|
|
state: (batch_size, 4) => [theta, omega, alpha, tau_prev]
|
|
Returns: (batch_size, 4) => [dtheta/dt, domega/dt, dalpha/dt, dtau/dt]
|
|
"""
|
|
|
|
theta = state[:, 0]
|
|
omega = state[:, 1]
|
|
alpha = state[:, 2]
|
|
tau_prev = state[:, 3]
|
|
|
|
|
|
# Create tensor input for controller: [theta, omega, alpha]
|
|
input_3d = torch.stack([theta, omega, alpha], dim=1) # shape (batch_size, 3)
|
|
|
|
# Compute torque using the controller
|
|
tau = self.controller(input_3d).squeeze(-1) # shape (batch_size,)
|
|
|
|
# Compute desired alpha
|
|
alpha_desired = (g / R) * torch.sin(theta) + tau / (m * R**2)
|
|
|
|
# Define ODE system
|
|
dtheta = omega
|
|
domega = alpha
|
|
dalpha = alpha_desired - alpha # Relaxation term
|
|
dtau = tau - tau_prev # Keep track of torque evolution
|
|
|
|
return torch.stack([dtheta, domega, dalpha, dtau], dim=1) # (batch_size, 4)
|
|
|
|
# ----------------------------------------------------------------
|
|
# 3) Loss Function
|
|
# ----------------------------------------------------------------
|
|
def loss_fn(state_traj, t_span):
|
|
"""
|
|
Computes loss based on the trajectory with inverse time weighting (1/t) for theta and omega.
|
|
|
|
Args:
|
|
state_traj: Tensor of shape (time_steps, batch_size, 4).
|
|
t_span: Tensor of time steps (time_steps,).
|
|
|
|
Returns:
|
|
total_loss, (loss_theta, loss_omega)
|
|
"""
|
|
theta = state_traj[:, :, 0] # (time_steps, batch_size)
|
|
omega = state_traj[:, :, 1] # (time_steps, batch_size)
|
|
torque = state_traj[:, :, 3]
|
|
|
|
# Inverse time weights w(t) = 1 / t
|
|
# Add a small epsilon to avoid division by zero
|
|
epsilon = 1e-6
|
|
inverse_time_weights = 1.0 / (t_span + epsilon).unsqueeze(1) # Shape: (time_steps, 1)
|
|
linear_time_weights = t_span.unsqueeze(1)
|
|
|
|
# Apply inverse time weighting for theta and omega
|
|
loss_theta = 1e-1 * torch.mean(inverse_time_weights * theta**2) # Weighted theta loss
|
|
loss_omega = 1e-2 * torch.mean(inverse_time_weights * omega**2) # Weighted omega loss
|
|
loss_torque = 1e-2 * torch.mean(linear_time_weights * torque**2)
|
|
|
|
# Combine the losses
|
|
total_loss = loss_theta #+ loss_torque
|
|
|
|
return total_loss, (loss_theta, loss_omega, loss_torque)
|
|
|
|
# ----------------------------------------------------------------
|
|
# 4) Training Setup
|
|
# ----------------------------------------------------------------
|
|
device = torch.device("cpu" if torch.cuda.is_available() else "cpu")
|
|
|
|
# Create the controller and pendulum dynamics model
|
|
controller = PendulumController3D().to(device)
|
|
pendulum_dynamics = PendulumDynamics3D(controller).to(device)
|
|
|
|
# Define optimizer
|
|
optimizer = optim.Adam(controller.parameters(), lr=1e-1)
|
|
|
|
# Initial conditions: [theta, omega, alpha, tau_prev]
|
|
initial_conditions = [
|
|
[0.1, 0.0, 0.0, 0.0], # Small perturbation
|
|
[-0.5, 0.0, 0.0, 0.0],
|
|
[6.28, 6.28, 0.0, 0.0],
|
|
[1.57, 0.5, 0.0, 0.0],
|
|
[0.0, -6.28, 0.0, 0.0],
|
|
[1.57, -6.28, 0.0, 0.0],
|
|
]
|
|
|
|
# Convert to torch tensor (batch_size, 4)
|
|
state_0 = torch.tensor(initial_conditions, dtype=torch.float32, device=device)
|
|
|
|
# Time grid
|
|
t_span = torch.linspace(0, 10, 1000, device=device)
|
|
|
|
num_epochs = 100_000
|
|
print_every = 25
|
|
|
|
# ----------------------------------------------------------------
|
|
# 5) Training Loop
|
|
# ----------------------------------------------------------------
|
|
for epoch in range(num_epochs):
|
|
optimizer.zero_grad()
|
|
|
|
# Integrate the ODE
|
|
state_traj = odeint(pendulum_dynamics, state_0, t_span, method='rk4')
|
|
# state_traj shape: (time_steps, batch_size, 4)
|
|
|
|
# Compute loss
|
|
total_loss, (l_theta, l_omega, l_torque) = loss_fn(state_traj, t_span)
|
|
|
|
# Check for NaN values
|
|
if torch.isnan(total_loss):
|
|
print(f"NaN detected at epoch {epoch}. Skipping step.")
|
|
optimizer.zero_grad()
|
|
continue # Skip this iteration
|
|
|
|
# Backprop
|
|
total_loss.backward()
|
|
optimizer.step()
|
|
|
|
|
|
# Print progress
|
|
if epoch % print_every == 0:
|
|
print(f"Epoch {epoch:4d}/{num_epochs} | "
|
|
f"Total: {total_loss.item():.6f} | "
|
|
f"Theta: {l_theta.item():.6f} | "
|
|
f"Omega: {l_omega.item():.6f} | "
|
|
f"Torque: {l_torque.item():.6f}")
|
|
|
|
torch.save(controller.state_dict(), "controller_cpu_clamped_inverse_time_punish.pth")
|
|
print("Model saved as 'controller_cpu_clamped_inverse_time_punish.pth'.")
|