174 lines
6.2 KiB
Python
174 lines
6.2 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.optim as optim
|
|
from torchdiffeq import odeint
|
|
import numpy as np
|
|
import inspect
|
|
import time
|
|
import csv
|
|
import os
|
|
|
|
# Specify directory for storing results
|
|
output_dir = "training/exponential_time_weight"
|
|
controller_output_dir = output_dir + "/controllers"
|
|
os.makedirs(output_dir, exist_ok=True) # Create directory if it doesn't exist
|
|
os.makedirs(controller_output_dir, exist_ok=True) # Create directory if it doesn't exist
|
|
|
|
# Use a previously generated random seed
|
|
random_seed = 4529
|
|
|
|
# Set the seeds for reproducibility
|
|
torch.manual_seed(random_seed)
|
|
np.random.seed(random_seed)
|
|
|
|
# Print the chosen random seed
|
|
print(f"Random seed for torch and numpy: {random_seed}")
|
|
|
|
# Constants
|
|
m = 10.0
|
|
g = 9.81
|
|
R = 1.0
|
|
|
|
# Device setup
|
|
device = torch.device("cpu")
|
|
|
|
# Time grid
|
|
t_start, t_end, t_points = 0, 10, 1000
|
|
t_span = torch.linspace(t_start, t_end, t_points, device=device)
|
|
|
|
# Initial conditions (theta0, omega0, alpha0, desired_theta)
|
|
state_0 = torch.tensor([
|
|
[1/6 * torch.pi, 0.0, 0.0, 0.0],
|
|
[-1/6 * torch.pi, 0.0, 0.0, 0.0],
|
|
[2/3 * torch.pi, 0.0, 0.0, 0.0],
|
|
[-2/3 * torch.pi, 0.0, 0.0, 0.0],
|
|
[0.0, 1/3 * torch.pi, 0.0, 0.0],
|
|
[0.0, -1/3 * torch.pi, 0.0, 0.0],
|
|
[0.0, 2 * torch.pi, 0.0, 0.0],
|
|
[0.0, -2 * torch.pi, 0.0, 0.0],
|
|
[0.0, 0.0, 0.0, 2 * torch.pi],
|
|
[0.0, 0.0, 0.0, -2 * torch.pi],
|
|
[0.0, 0.0, 0.0, 1/2 * torch.pi],
|
|
[0.0, 0.0, 0.0, -1/2 * torch.pi],
|
|
[0.0, 0.0, 0.0, 1/3 * torch.pi],
|
|
[0.0, 0.0, 0.0, -1/3 * torch.pi],
|
|
[1/4 * torch.pi, 1 * torch.pi, 0.0, 0.0],
|
|
[-1/4 * torch.pi, -1 * torch.pi, 0.0, 0.0],
|
|
[1/2 * torch.pi, -1 * torch.pi, 0.0, 1/3 * torch.pi],
|
|
[-1/2 * torch.pi, 1 * torch.pi, 0.0, -1/3 * torch.pi],
|
|
[1/4 * torch.pi, 1 * torch.pi, 0.0, 2 * torch.pi],
|
|
[-1/4 * torch.pi, -1 * torch.pi, 0.0, 2 * torch.pi],
|
|
[1/2 * torch.pi, -1 * torch.pi, 0.0, 4 * torch.pi],
|
|
[-1/2 * torch.pi, 1 * torch.pi, 0.0, -4 * torch.pi],
|
|
], dtype=torch.float32, device=device)
|
|
|
|
class PendulumController(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.net = nn.Sequential(
|
|
nn.Linear(4, 64),
|
|
nn.ReLU(),
|
|
nn.Linear(64, 64),
|
|
nn.ReLU(),
|
|
nn.Linear(64, 1)
|
|
)
|
|
|
|
def forward(self, x):
|
|
raw_torque = self.net(x)
|
|
return torch.clamp(raw_torque, -250, 250)
|
|
|
|
class PendulumDynamics(nn.Module):
|
|
def __init__(self, controller):
|
|
super().__init__()
|
|
self.controller = controller
|
|
|
|
def forward(self, t, state):
|
|
theta, omega, alpha, desired_theta = state[:, 0], state[:, 1], state[:, 2], state[:, 3]
|
|
input = torch.stack([theta, omega, alpha, desired_theta], dim=1)
|
|
tau = self.controller(input).squeeze(-1)
|
|
alpha_desired = (g / R) * torch.sin(theta) + tau / (m * R**2)
|
|
return torch.stack([omega, alpha, alpha_desired - alpha, torch.zeros_like(desired_theta)], dim=1)
|
|
|
|
def loss_fn(state_traj, t_span):
|
|
theta = state_traj[:, :, 0]
|
|
desired_theta = state_traj[:, :, 3]
|
|
|
|
# Make the time weights broadcastable for theta [len(t_span), len(batches)]
|
|
lamda_exp = 1
|
|
time_weights = torch.exp(lamda_exp * t_span).view(-1, 1)
|
|
|
|
return 1e3 * torch.mean(time_weights * (theta - desired_theta) ** 2)
|
|
|
|
# Initialize controller and dynamics
|
|
controller = PendulumController().to(device)
|
|
pendulum_dynamics = PendulumDynamics(controller).to(device)
|
|
|
|
# Optimizer setup
|
|
learning_rate = 1e-2
|
|
weight_decay = 1e-4
|
|
optimizer = optim.Adam(controller.parameters(), lr=learning_rate, weight_decay=weight_decay)
|
|
|
|
# Training parameters
|
|
num_epochs = 5_001 # 5000 + 1 for 5000 total weight updates
|
|
save_every = 1 # How often to save the controller.pth
|
|
|
|
# File paths
|
|
config_file = os.path.join(output_dir, "training_config.txt")
|
|
log_file = os.path.join(output_dir, "training_log.csv")
|
|
model_file = "" # Placeholder for the model file, updated in training loop
|
|
|
|
# Save configuration details
|
|
with open(config_file, "w") as f:
|
|
f.write(f"Random Seed: {random_seed}\n")
|
|
f.write(f"Time Span: {t_start} to {t_end}, Points: {t_points}\n")
|
|
f.write(f"Learning Rate: {learning_rate}\n")
|
|
f.write(f"Weight Decay: {weight_decay}\n")
|
|
f.write("\nLoss Function:\n")
|
|
f.write(inspect.getsource(loss_fn)) # Extract and write loss function source code
|
|
f.write("\nTraining Cases:\n")
|
|
f.write("[theta0, omega0, alpha0, desired_theta]\n")
|
|
for case in state_0.cpu().numpy():
|
|
f.write(f"{case.tolist()}\n")
|
|
|
|
|
|
# Overwrite the log file at the start
|
|
with open(log_file, "w", newline="") as csvfile:
|
|
csv_writer = csv.writer(csvfile)
|
|
csv_writer.writerow(["Epoch", "Loss", "Elapsed Time (s)"])
|
|
|
|
# Training loop with real-time logging and NaN tracking
|
|
start_time = time.time()
|
|
with open(log_file, "a", newline="") as csvfile:
|
|
csv_writer = csv.writer(csvfile)
|
|
|
|
for epoch in range(num_epochs):
|
|
epoch_start_time = time.time()
|
|
|
|
optimizer.zero_grad()
|
|
state_traj = odeint(pendulum_dynamics, state_0, t_span, method='rk4')
|
|
loss = loss_fn(state_traj, t_span)
|
|
|
|
elapsed_time = time.time() - epoch_start_time
|
|
|
|
if torch.isnan(loss):
|
|
print(f"NaN detected at epoch {epoch}. Skipping step.")
|
|
csv_writer.writerow([epoch, "NaN detected", elapsed_time])
|
|
csvfile.flush() # Ensure real-time writing
|
|
optimizer.zero_grad()
|
|
continue
|
|
|
|
loss.backward()
|
|
optimizer.step()
|
|
|
|
# Log normal loss
|
|
csv_writer.writerow([epoch, loss.item(), elapsed_time])
|
|
csvfile.flush() # Ensure real-time writing
|
|
|
|
if epoch % save_every == 0:
|
|
print(f"Epoch {epoch}/{num_epochs} | Loss: {loss.item():.6f} | Time: {elapsed_time:.4f} sec")
|
|
model_file = os.path.join(controller_output_dir, f"controller_{epoch}.pth")
|
|
torch.save(controller.state_dict(), model_file)
|
|
|
|
# Final save
|
|
torch.save(controller.state_dict(), model_file)
|
|
print(f"No time weight training complete. Model and files saved in directory '{output_dir}'. Model saved as '{model_file}'. Logs saved in '{log_file}' and configuration in '{config_file}'.") |