Controller across epochs plotter

analysis/controller_across_epochs.py

@@ -0,0 +1,135 @@
+import os
+import numpy as np
+import torch
+import torch.nn as nn
+import matplotlib
+matplotlib.use("Agg")  # Use non-interactive backend
+import matplotlib.pyplot as plt
+from mpl_toolkits.mplot3d import Axes3D
+import multiprocessing
+
+# Define PendulumController class
+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):
+        return self.net(x)  
+
+# ODE solver (RK4 method)
+def pendulum_ode_step(state, dt, desired_theta, controller):
+    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
+
+    # Compute 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
+    return new_state
+
+# Constants
+g = 9.81  # Gravity
+R = 1.0   # Length of the pendulum
+m = 1.0   # Mass
+dt = 0.02  # Time step
+num_steps = 500  # Simulation time steps
+
+# Directory containing controller files
+controller_dir = "/home/judson/Neural-Networks-in-GNC/inverted_pendulum/training/no_time_weight/controllers"
+controller_files = sorted([f for f in os.listdir(controller_dir) if f.startswith("controller_") and f.endswith(".pth")])
+
+# Sorting controllers by epoch
+controller_epochs = [int(f.split('_')[1].split('.')[0]) for f in controller_files]
+sorted_controllers = [x for _, x in sorted(zip(controller_epochs, controller_files))]
+
+# **Granularity Control: Select every Nth controller**
+N = 5  # Change this value to adjust granularity (e.g., every 5th controller)
+selected_controllers = sorted_controllers[::N]  # Take every Nth controller
+
+# Initial condition
+theta0, omega0, alpha0, desired_theta = (-np.pi, -np.pi, 0.0, np.pi / 6)  # Example initial condition
+
+# Function to run a single controller simulation (for multiprocessing)
+def run_simulation(controller_file):
+    epoch = int(controller_file.split('_')[1].split('.')[0])
+    
+    # Load controller
+    controller = PendulumController()
+    controller.load_state_dict(torch.load(os.path.join(controller_dir, controller_file)))
+    controller.eval()
+    
+    # Run simulation
+    state = np.array([theta0, omega0, alpha0])
+    theta_vals = []
+
+    for _ in range(num_steps):
+        theta_vals.append(state[0])
+        state = pendulum_ode_step(state, dt, desired_theta, controller)
+
+    return epoch, theta_vals
+
+# Parallel processing
+if __name__ == "__main__":
+    num_workers = min(multiprocessing.cpu_count(), 16)  # Limit to 16 workers max
+    print(f"Using {num_workers} parallel workers...")
+    print(f"Processing every {N}th controller, total controllers used: {len(selected_controllers)}")
+
+    with multiprocessing.Pool(processes=num_workers) as pool:
+        results = pool.map(run_simulation, selected_controllers)
+
+    # Sort results by epoch
+    results.sort(key=lambda x: x[0])
+    epochs, theta_over_epochs = zip(*results)
+
+    # Convert results to NumPy arrays
+    theta_over_epochs = np.array(theta_over_epochs)
+
+    # Create 3D plot
+    fig = plt.figure(figsize=(10, 7))
+    ax = fig.add_subplot(111, projection='3d')
+
+    # Meshgrid for 3D plotting
+    E, T = np.meshgrid(epochs, np.arange(num_steps) * dt)
+
+    # Plot surface
+    ax.plot_surface(E, T, theta_over_epochs.T, cmap="viridis")
+
+    # Labels
+    ax.set_xlabel("Epoch")
+    ax.set_ylabel("Time (s)")
+    ax.set_zlabel("Theta (rad)")
+    ax.set_title(f"Pendulum Angle Evolution Over Training Epochs (Granularity N={N})")
+
+    plt.savefig("pendulum_plot.png", dpi=1000, bbox_inches="tight")
+    print("Saved plot as 'pendulum_plot.png'.")

check_for_cuda.py

@@ -1,8 +0,0 @@
-import torch
-
-if torch.cuda.is_available():
-    print("CUDA is available")
-    print("Number of CUDA devices:", torch.cuda.device_count())
-    print("Device name:", torch.cuda.get_device_name(0))
-else:
-    print("CUDA is not available")

clamped_inverse_time_penalty/trainer.py

@@ -1,175 +0,0 @@
-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'.")

clamped_inverse_time_penalty/validator.py

@@ -1,116 +0,0 @@
-import torch
-import torch.nn as nn
-import numpy as np
-from scipy.integrate import solve_ivp
-import matplotlib.pyplot as plt
-
-# ----------------------------------------------------------------
-# 1) 3D Controller: [theta, omega, alpha] -> torque
-# ----------------------------------------------------------------
-class PendulumController3D(nn.Module):
-    def __init__(self):
-        super(PendulumController3D, self).__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):
-        return self.net(x_3d)
-
-# Load the trained 3D model
-controller = PendulumController3D()
-controller.load_state_dict(torch.load("controller_cpu_clamped_inverse_time_penalty.pth"))
-# controller.load_state_dict(torch.load("controller_cpu_clamped.pth"))
-controller.eval()
-print("3D Controller loaded.")
-
-# ----------------------------------------------------------------
-# 2) ODE: State = [theta, omega, alpha].
-# ----------------------------------------------------------------
-m = 10.0
-g = 9.81
-R = 1.0
-
-def pendulum_ode_3d(t, state):
-    theta, omega, alpha = state
-
-    # Evaluate NN -> torque
-    inp = torch.tensor([[theta, omega, alpha]], dtype=torch.float32)
-    with torch.no_grad():
-        torque = controller(inp).item()
-        # Clamp torque to ±250 for consistency with training
-        torque = np.clip(torque, -250, 250)
-
-    alpha_des = (g/R)*np.sin(theta) + torque/(m*(R**2))
-
-    dtheta = omega
-    domega = alpha
-    dalpha = alpha_des - alpha
-    return [dtheta, domega, dalpha]
-
-# ----------------------------------------------------------------
-# 3) Validate for multiple initial conditions
-# ----------------------------------------------------------------
-initial_conditions_3d = [
-    (0.1,  0.0, 0.0),
-    (0.5,  0.0, 0.0),
-    (1.0,  0.0, 0.0),
-    (1.57, 0.5, 0.0),
-    (0.0, -6.28, 0.0),
-    (6.28, 6.28, 0.0),
-]
-
-t_span = (0, 20)
-t_eval = np.linspace(0, 20, 2000)
-
-for idx, (theta0, omega0, alpha0) in enumerate(initial_conditions_3d):
-    sol = solve_ivp(
-        pendulum_ode_3d,
-        t_span,
-        [theta0, omega0, alpha0],
-        t_eval=t_eval,
-        method='RK45'
-    )
-
-    t         = sol.t
-    theta     = sol.y[0]
-    omega     = sol.y[1]
-    alpha_arr = sol.y[2]
-
-    # Recompute torque over time
-    torques = []
-    alpha_des_vals = []
-    for (th, om, al) in zip(theta, omega, alpha_arr):
-        with torch.no_grad():
-            torque_val = controller(torch.tensor([[th, om, al]], dtype=torch.float32)).item()
-            torque_val = np.clip(torque_val, -250, 250)
-        torques.append(torque_val)
-        alpha_des_vals.append( (g/R)*np.sin(th) + torque_val/(m*(R**2)) )
-    torques = np.array(torques)
-
-    # Plot
-    fig, ax1 = plt.subplots(figsize=(10,6))
-    ax1.plot(t, theta,     label="theta",      color="blue")
-    ax1.plot(t, omega,     label="omega",      color="green")
-    ax1.plot(t, alpha_arr, label="alpha",      color="red")
-    # optional: ax1.plot(t, alpha_des_vals, label="alpha_des", color="red", linestyle="--")
-
-    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, torques, label="torque", color="purple", linestyle="--")
-    ax2.set_ylabel("Torque [Nm]")
-    ax2.legend(loc="upper right")
-
-    plt.title(f"IC (theta={theta0}, omega={omega0}, alpha={alpha0})")
-    plt.tight_layout()
-    plt.savefig(f"{idx+1}_validation.png")
-    plt.close()
-    print(f"Saved {idx+1}_validation.png")

clamped_no_time_penalty/trainer.py

@@ -1,161 +0,0 @@
-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)
-
-    loss_theta = 1e3 * torch.mean(theta**2)
-
-    # Combine the losses
-    total_loss = loss_theta
-
-    return total_loss, (loss_theta)
-
-# ----------------------------------------------------------------
-# 4) Training Setup
-# ----------------------------------------------------------------
-device = torch.device("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, 200, 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) = 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}")
-
-        torch.save(controller.state_dict(), "controller_cpu_clamped_no_time_penalty.pth")
-        print("Model saved as 'controller_cpu_clamped_no_time_penalty.pth'.")

clamped_no_time_penalty/validator.py

@@ -1,118 +0,0 @@
-import torch
-import torch.nn as nn
-import numpy as np
-from scipy.integrate import solve_ivp
-import matplotlib.pyplot as plt
-
-
-controller_file_name = "controller_cpu_clamped_no_time_penalty.pth"
-
-# ----------------------------------------------------------------
-# 1) 3D Controller: [theta, omega, alpha] -> torque
-# ----------------------------------------------------------------
-class PendulumController3D(nn.Module):
-    def __init__(self):
-        super(PendulumController3D, self).__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):
-        return self.net(x_3d)
-
-# Load the trained 3D model
-controller = PendulumController3D()
-controller.load_state_dict(torch.load(controller_file_name))
-controller.eval()
-print(f"{controller_file_name} loaded.")
-
-# ----------------------------------------------------------------
-# 2) ODE: State = [theta, omega, alpha].
-# ----------------------------------------------------------------
-m = 10.0
-g = 9.81
-R = 1.0
-
-def pendulum_ode_3d(t, state):
-    theta, omega, alpha = state
-
-    # Evaluate NN -> torque
-    inp = torch.tensor([[theta, omega, alpha]], dtype=torch.float32)
-    with torch.no_grad():
-        torque = controller(inp).item()
-        # Clamp torque to ±250 for consistency with training
-        torque = np.clip(torque, -250, 250)
-
-    alpha_des = (g/R)*np.sin(theta) + torque/(m*(R**2))
-
-    dtheta = omega
-    domega = alpha
-    dalpha = alpha_des - alpha
-    return [dtheta, domega, dalpha]
-
-# ----------------------------------------------------------------
-# 3) Validate for multiple initial conditions
-# ----------------------------------------------------------------
-initial_conditions_3d = [
-    (0.1,  0.0, 0.0),
-    (0.5,  0.0, 0.0),
-    (1.0,  0.0, 0.0),
-    (1.57, 0.5, 0.0),
-    (0.0, -6.28, 0.0),
-    (6.28, 6.28, 0.0),
-]
-
-t_span = (0, 20)
-t_eval = np.linspace(0, 20, 2000)
-
-for idx, (theta0, omega0, alpha0) in enumerate(initial_conditions_3d):
-    sol = solve_ivp(
-        pendulum_ode_3d,
-        t_span,
-        [theta0, omega0, alpha0],
-        t_eval=t_eval,
-        method='RK45'
-    )
-
-    t         = sol.t
-    theta     = sol.y[0]
-    omega     = sol.y[1]
-    alpha_arr = sol.y[2]
-
-    # Recompute torque over time
-    torques = []
-    alpha_des_vals = []
-    for (th, om, al) in zip(theta, omega, alpha_arr):
-        with torch.no_grad():
-            torque_val = controller(torch.tensor([[th, om, al]], dtype=torch.float32)).item()
-            torque_val = np.clip(torque_val, -250, 250)
-        torques.append(torque_val)
-        alpha_des_vals.append( (g/R)*np.sin(th) + torque_val/(m*(R**2)) )
-    torques = np.array(torques)
-
-    # Plot
-    fig, ax1 = plt.subplots(figsize=(10,6))
-    ax1.plot(t, theta,     label="theta",      color="blue")
-    ax1.plot(t, omega,     label="omega",      color="green")
-    ax1.plot(t, alpha_arr, label="alpha",      color="red")
-    # optional: ax1.plot(t, alpha_des_vals, label="alpha_des", color="red", linestyle="--")
-
-    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, torques, label="torque", color="purple", linestyle="--")
-    ax2.set_ylabel("Torque [Nm]")
-    ax2.legend(loc="upper right")
-
-    plt.title(f"IC (theta={theta0}, omega={omega0}, alpha={alpha0})")
-    plt.tight_layout()
-    plt.savefig(f"{idx+1}_validation.png")
-    plt.close()
-    print(f"Saved {idx+1}_validation.png")

clamped_quadratic_time_penalty/trainer.py

@@ -1,180 +0,0 @@
-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_3d: shape (batch_size, 3) => [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: exponentially increasing theta^2 penalty over time.
-    
-    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, loss_torque)
-    """
-    theta = state_traj[:, :, 0]  # (time_steps, batch_size)
-    omega = state_traj[:, :, 1]
-    torque = state_traj[:, :, 3]  # tau_prev is stored in state
-
-    # Quadratic weight factor lambda * t**2
-    lambda_factor = 0.5  # Increase for stronger late-time punishment
-    time_weights = (lambda_factor * t_span**2).unsqueeze(1)  # Shape: (time_steps, 1)
-
-    # Apply increasing penalty over time
-    loss_theta = 1e2 * torch.mean(time_weights * (torch.cos(theta) - 1)**2)
-    #loss_theta = 1e-1 * torch.mean(time_weights * theta**2)
-    loss_omega = 1e-1 * torch.mean(omega**2)
-    loss_torque = 1e-5 * torch.mean(torque**2)
-
-    # Extract the final theta value from the trajectory
-    final_theta = state_traj[-1, :, 0]  # (batch_size,)
-    
-    # Compute the loss as the squared error from the target theta
-    loss_final_theta = torch.mean(final_theta ** 2)  # Mean squared error
-
-    total_loss = loss_theta #+ loss_omega + loss_torque
-    return total_loss, (loss_theta, loss_omega, loss_torque, loss_final_theta)
-
-
-# ----------------------------------------------------------------
-# 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-2)
-
-# 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, 500, device=device)  # 10 seconds, 500 steps
-
-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, l_final_theta) = 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()
-    #torch.nn.utils.clip_grad_norm_(controller.parameters(), max_norm=1.0)  # Fix NaNs
-    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} | "
-              f"Final Theta: {l_final_theta.item():.6f}")
-
-        torch.save(controller.state_dict(), "controller_cpu_clamped_quadratic_time_punish.pth")
-        print("Model saved as 'controller_cpu_clamped_quadratic_time_punish.pth'.")

clamped_quadratic_time_penalty/validator.py

@@ -1,116 +0,0 @@
-import torch
-import torch.nn as nn
-import numpy as np
-from scipy.integrate import solve_ivp
-import matplotlib.pyplot as plt
-
-# ----------------------------------------------------------------
-# 1) 3D Controller: [theta, omega, alpha] -> torque
-# ----------------------------------------------------------------
-class PendulumController3D(nn.Module):
-    def __init__(self):
-        super(PendulumController3D, self).__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):
-        return self.net(x_3d)
-
-# Load the trained 3D model
-controller = PendulumController3D()
-controller.load_state_dict(torch.load("controller_cpu_clamped_quadratic_time_penalty.pth"))
-# controller.load_state_dict(torch.load("controller_cpu_clamped.pth"))
-controller.eval()
-print("3D Controller loaded.")
-
-# ----------------------------------------------------------------
-# 2) ODE: State = [theta, omega, alpha].
-# ----------------------------------------------------------------
-m = 10.0
-g = 9.81
-R = 1.0
-
-def pendulum_ode_3d(t, state):
-    theta, omega, alpha = state
-
-    # Evaluate NN -> torque
-    inp = torch.tensor([[theta, omega, alpha]], dtype=torch.float32)
-    with torch.no_grad():
-        torque = controller(inp).item()
-        # Clamp torque to ±250 for consistency with training
-        torque = np.clip(torque, -250, 250)
-
-    alpha_des = (g/R)*np.sin(theta) + torque/(m*(R**2))
-
-    dtheta = omega
-    domega = alpha
-    dalpha = alpha_des - alpha
-    return [dtheta, domega, dalpha]
-
-# ----------------------------------------------------------------
-# 3) Validate for multiple initial conditions
-# ----------------------------------------------------------------
-initial_conditions_3d = [
-    (0.1,  0.0, 0.0),
-    (0.5,  0.0, 0.0),
-    (1.0,  0.0, 0.0),
-    (1.57, 0.5, 0.0),
-    (0.0, -6.28, 0.0),
-    (6.28, 6.28, 0.0),
-]
-
-t_span = (0, 20)
-t_eval = np.linspace(0, 20, 2000)
-
-for idx, (theta0, omega0, alpha0) in enumerate(initial_conditions_3d):
-    sol = solve_ivp(
-        pendulum_ode_3d,
-        t_span,
-        [theta0, omega0, alpha0],
-        t_eval=t_eval,
-        method='RK45'
-    )
-
-    t         = sol.t
-    theta     = sol.y[0]
-    omega     = sol.y[1]
-    alpha_arr = sol.y[2]
-
-    # Recompute torque over time
-    torques = []
-    alpha_des_vals = []
-    for (th, om, al) in zip(theta, omega, alpha_arr):
-        with torch.no_grad():
-            torque_val = controller(torch.tensor([[th, om, al]], dtype=torch.float32)).item()
-            torque_val = np.clip(torque_val, -250, 250)
-        torques.append(torque_val)
-        alpha_des_vals.append( (g/R)*np.sin(th) + torque_val/(m*(R**2)) )
-    torques = np.array(torques)
-
-    # Plot
-    fig, ax1 = plt.subplots(figsize=(10,6))
-    ax1.plot(t, theta,     label="theta",      color="blue")
-    ax1.plot(t, omega,     label="omega",      color="green")
-    ax1.plot(t, alpha_arr, label="alpha",      color="red")
-    # optional: ax1.plot(t, alpha_des_vals, label="alpha_des", color="red", linestyle="--")
-
-    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, torques, label="torque", color="purple", linestyle="--")
-    ax2.set_ylabel("Torque [Nm]")
-    ax2.legend(loc="upper right")
-
-    plt.title(f"IC (theta={theta0}, omega={omega0}, alpha={alpha0})")
-    plt.tight_layout()
-    plt.savefig(f"{idx+1}_validation.png")
-    plt.close()
-    print(f"Saved {idx+1}_validation.png")