198 lines
8.1 KiB
Python
198 lines
8.1 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
from torchdiffeq import odeint
|
|
import numpy as np
|
|
import os
|
|
|
|
import matplotlib.pyplot as plt
|
|
import pandas as pd
|
|
|
|
import sys
|
|
sys.path.append("/home/judson/Neural-Networks-in-GNC/inverted_pendulum/analysis")
|
|
from PendulumController import PendulumController
|
|
|
|
# List of controller file names to validate.
|
|
# Replace these paths with your actual controller file paths.
|
|
controller_file_names = [
|
|
"/home/judson/Neural-Networks-in-GNC/inverted_pendulum/training/training_files/time_weighting/constant/controllers/controller_1000.pth",
|
|
]
|
|
|
|
# Constants for simulation
|
|
m = 10.0
|
|
g = 9.81
|
|
R = 1.0
|
|
dt = 0.01 # time step
|
|
T = 20.0 # total simulation time
|
|
num_steps = int(T / dt)
|
|
t_eval = np.linspace(0, T, num_steps)
|
|
|
|
# In-sample validation cases: [theta0, omega0, alpha0, desired_theta]
|
|
in_sample_cases = [
|
|
(1/6 * np.pi, 0.0, 0.0, 0.0),
|
|
# (-1/6 * np.pi, 0.0, 0.0, 0.0),
|
|
# (2/3 * np.pi, 0.0, 0.0, 0.0),
|
|
# (-2/3 * np.pi, 0.0, 0.0, 0.0),
|
|
# (0.0, 1/3 * np.pi, 0.0, 0.0),
|
|
# (0.0, -1/3 * np.pi, 0.0, 0.0),
|
|
# (0.0, 2 * np.pi, 0.0, 0.0),
|
|
# (0.0, -2 * np.pi, 0.0, 0.0),
|
|
# (0.0, 0.0, 0.0, 2*np.pi),
|
|
# (0.0, 0.0, 0.0, -2*np.pi),
|
|
# (0.0, 0.0, 0.0, 1/2 * np.pi),
|
|
# (0.0, 0.0, 0.0, -1/2 * np.pi),
|
|
# (0.0, 0.0, 0.0, 1/3 * np.pi),
|
|
# (0.0, 0.0, 0.0, -1/3 * np.pi),
|
|
# (1/4 * np.pi, 1 * np.pi, 0.0, 0.0),
|
|
# (-1/4 * np.pi, -1 * np.pi, 0.0, 0.0),
|
|
# (1/2 * np.pi, -1 * np.pi, 0.0, 1/3 * np.pi),
|
|
# (-1/2 * np.pi, 1 * np.pi, 0.0, -1/3 * np.pi),
|
|
# (1/4 * np.pi, 1 * np.pi, 0.0, 2 * np.pi),
|
|
# (-1/4 * np.pi, -1 * np.pi, 0.0, 2 * np.pi),
|
|
# (1/2 * np.pi, -1 * np.pi, 0.0, 4 * np.pi),
|
|
# (-1/2 * np.pi, 1 * np.pi, 0.0, -4 * np.pi),
|
|
]
|
|
|
|
# Out-of-sample validation cases (generated previously)
|
|
out_sample_cases = [
|
|
# (-2.198958, -4.428501, 0.450833, 0.000000),
|
|
# (1.714196, -0.769896, 0.202738, 0.000000),
|
|
# (0.241195, -5.493715, 0.438996, 0.000000),
|
|
# (0.030605, 4.901513, -0.479243, 0.000000),
|
|
# (1.930445, -1.301926, -0.454050, 0.000000),
|
|
# (-0.676063, 4.246865, 0.036303, 0.000000),
|
|
# (0.734920, -5.925202, 0.047097, 0.000000),
|
|
# (-3.074471, -3.535424, 0.315438, 0.000000),
|
|
# (-0.094486, 6.111091, 0.150525, 0.000000),
|
|
# (-1.647671, 5.720526, 0.334181, 0.000000),
|
|
# (-2.611260, 5.087704, 0.045460, -3.610785),
|
|
# (1.654137, 0.982081, -0.192725, 1.003872),
|
|
# (-2.394899, 3.550547, -0.430938, 3.261897),
|
|
# (0.474917, 0.555166, -0.285173, 1.866752),
|
|
# (-0.640369, -4.678490, -0.340663, 3.150098),
|
|
# (1.747517, -3.248204, -0.001520, 1.221787),
|
|
# (2.505283, -2.875006, -0.065617, -3.690269),
|
|
# (1.337244, 2.221707, 0.044979, -2.459730),
|
|
# (1.531012, 2.230981, -0.291206, -1.924535),
|
|
# (-1.065792, 4.320740, 0.075405, -1.550644),
|
|
]
|
|
|
|
# Define the ODE integration step using RK4, returning the new state and the computed torque.
|
|
def pendulum_ode_step(state, dt, desired_theta):
|
|
theta, omega, alpha = state
|
|
|
|
def compute_torque(th, om, al):
|
|
inp = torch.tensor([[th, om, al, desired_theta]], dtype=torch.float32)
|
|
with torch.no_grad():
|
|
torque = controller(inp)
|
|
torque = torch.clamp(torque, -250, 250)
|
|
return float(torque)
|
|
|
|
def derivatives(state, torque):
|
|
th, om, al = state
|
|
a = (g / R) * np.sin(th) + torque / (m * R**2)
|
|
return np.array([om, a, 0]) # dtheta, domega, dalpha
|
|
|
|
# RK4 Steps
|
|
torque1 = compute_torque(theta, omega, alpha)
|
|
k1 = dt * derivatives(state, torque1)
|
|
|
|
state_k2 = state + 0.5 * k1
|
|
torque2 = compute_torque(state_k2[0], state_k2[1], state_k2[2])
|
|
k2 = dt * derivatives(state_k2, torque2)
|
|
|
|
state_k3 = state + 0.5 * k2
|
|
torque3 = compute_torque(state_k3[0], state_k3[1], state_k3[2])
|
|
k3 = dt * derivatives(state_k3, torque3)
|
|
|
|
state_k4 = state + k3
|
|
torque4 = compute_torque(state_k4[0], state_k4[1], state_k4[2])
|
|
k4 = dt * derivatives(state_k4, torque4)
|
|
|
|
new_state = state + (k1 + 2*k2 + 2*k3 + k4) / 6.0
|
|
final_torque = (torque1 + 2*torque2 + 2*torque3 + torque4) / 6.0
|
|
|
|
return new_state, final_torque
|
|
|
|
# Function to run validation (in-sample or out-of-sample) for a given controller.
|
|
def run_validation(controller, cases, t_eval, dt, num_steps, validation_type, save_dir):
|
|
losses = []
|
|
final_thetas = []
|
|
os.makedirs(save_dir, exist_ok=True)
|
|
for idx, (theta0, omega0, alpha0, desired_theta) in enumerate(cases):
|
|
state = np.array([theta0, omega0, alpha0])
|
|
theta_vals, omega_vals, alpha_vals, torque_vals = [], [], [], []
|
|
for _ in range(num_steps):
|
|
theta_vals.append(state[0])
|
|
omega_vals.append(state[1])
|
|
alpha_vals.append(state[2])
|
|
state, torque = pendulum_ode_step(state, dt, desired_theta)
|
|
torque_vals.append(torque)
|
|
theta_vals = np.array(theta_vals)
|
|
omega_vals = np.array(omega_vals)
|
|
alpha_vals = np.array(alpha_vals)
|
|
torque_vals = np.array(torque_vals)
|
|
final_theta = theta_vals[-1]
|
|
final_thetas.append(final_theta)
|
|
loss = 1e3 * np.mean((theta_vals - desired_theta)**2)
|
|
losses.append(loss)
|
|
|
|
# Plot results
|
|
fig, ax1 = plt.subplots(figsize=(10, 6))
|
|
ax1.plot(t_eval, theta_vals, label="theta", color="blue")
|
|
ax1.plot(t_eval, omega_vals, label="omega", color="green")
|
|
ax1.plot(t_eval, alpha_vals, label="alpha", color="red")
|
|
ax1.axhline(desired_theta, label="Desired Theta", color="black")
|
|
ax1.set_xlabel("Time [s]")
|
|
ax1.set_ylabel("Theta, Omega, Alpha")
|
|
ax1.grid(True)
|
|
ax1.legend(loc="upper left")
|
|
ax2 = ax1.twinx()
|
|
ax2.plot(t_eval, torque_vals, label="Torque", color="purple", linestyle="--")
|
|
ax2.set_ylabel("Torque [Nm]")
|
|
ax2.legend(loc="upper right")
|
|
plt.title(f"Case {idx+1}: IC (theta={theta0:.3f}, omega={omega0:.3f}, alpha={alpha0:.3f})")
|
|
plt.tight_layout()
|
|
filename = f"{idx+1}_theta0_{theta0:.3f}_omega0_{omega0:.3f}_alpha0_{alpha0:.3f}_desired_{desired_theta:.3f}_final_{final_theta:.3f}.png"
|
|
plt.savefig(os.path.join(save_dir, filename))
|
|
plt.close()
|
|
print(f"Saved {validation_type} case {idx+1}")
|
|
|
|
df = pd.DataFrame(cases, columns=["theta0", "omega0", "alpha0", "desired_theta"])
|
|
df["final_theta"] = final_thetas
|
|
df["loss"] = losses
|
|
df.insert(0, "Run #", range(1, len(cases) + 1))
|
|
df.to_csv(os.path.join(save_dir, "summary.csv"), index=False)
|
|
print(f"Average Loss for {validation_type}: {np.mean(losses):.4f}")
|
|
desired_vals = np.array([case[3] for case in cases])
|
|
avg_abs_err = np.mean(np.abs(np.array(final_thetas) - desired_vals))
|
|
print(f"Average abs(Final Theta - Desired Theta) for {validation_type}: {avg_abs_err:.4f}")
|
|
return df
|
|
|
|
# Main validation loop for each controller file
|
|
base_validation_dir = "validation"
|
|
os.makedirs(base_validation_dir, exist_ok=True)
|
|
|
|
for controller_file in controller_file_names:
|
|
controller_id = os.path.splitext(os.path.basename(controller_file))[0]
|
|
controller_dir = os.path.join(base_validation_dir, controller_id)
|
|
in_sample_dir = os.path.join(controller_dir, "in-sample")
|
|
out_sample_dir = os.path.join(controller_dir, "out-of-sample")
|
|
os.makedirs(in_sample_dir, exist_ok=True)
|
|
os.makedirs(out_sample_dir, exist_ok=True)
|
|
|
|
# Load controller
|
|
controller = PendulumController()
|
|
controller.load_state_dict(torch.load(controller_file))
|
|
controller.eval()
|
|
print(f"{controller_file} loaded as {controller_id}")
|
|
|
|
# In-sample validation
|
|
print("Performing in-sample validation")
|
|
df_in = run_validation(controller, in_sample_cases, t_eval, dt, num_steps, "in-sample", in_sample_dir)
|
|
|
|
# Out-of-sample validation
|
|
print("Performing out-of-sample validation")
|
|
df_out = run_validation(controller, out_sample_cases, t_eval, dt, num_steps, "out-of-sample", out_sample_dir)
|
|
|
|
print("Validation complete. Check the 'validation' directory for plots and summaries.")
|