Fix issues where the initialized controllers could be different. Created a controller_base.pth that is used for all controller initialization

This commit is contained in:
judsonupchurch 2025-02-24 02:28:16 +00:00
parent 8be7ad97a8
commit c89998da28
40 changed files with 156 additions and 21141 deletions

File diff suppressed because one or more lines are too long

View File

@ -14,8 +14,8 @@ initial_conditions = {
"extreme_perturbation": (4*np.pi, 0.0, 0.0, 0), "extreme_perturbation": (4*np.pi, 0.0, 0.0, 0),
} }
loss_functions = ["constant", "linear", "quadratic", "cubic", "inverse", "inverse_squared", "inverse_cubed"] loss_functions = ["constant", "linear", "quadratic", "cubic", "inverse", "inverse_squared", "inverse_cubed"]
epoch_range = (0, 1000) # Start and end of epoch range epoch_range = (0, 3) # Start and end of epoch range
epoch_step = 5 # Interval between epochs epoch_step = 1 # Interval between epochs
dt = 0.02 # Time step for simulation dt = 0.02 # Time step for simulation
num_steps = 500 # Number of steps in each simulation num_steps = 500 # Number of steps in each simulation
@ -26,11 +26,11 @@ if __name__ == "__main__":
for condition_name, initial_condition in initial_conditions.items(): for condition_name, initial_condition in initial_conditions.items():
condition_text = f"IC_{'_'.join(map(lambda x: str(round(x, 2)), initial_condition))}" condition_text = f"IC_{'_'.join(map(lambda x: str(round(x, 2)), initial_condition))}"
desired_theta = initial_condition[-1] desired_theta = initial_condition[-1]
condition_path = f"/home/judson/Neural-Networks-in-GNC/inverted_pendulum/analysis/average_normalized/{condition_name}" condition_path = f"/home/judson/Neural-Networks-in-GNC/inverted_pendulum/analysis/max_normalized/{condition_name}"
os.makedirs(condition_path, exist_ok=True) # Create directory if it does not exist os.makedirs(condition_path, exist_ok=True) # Create directory if it does not exist
for loss_function in loss_functions: for loss_function in loss_functions:
# Construct the path to the controller directory # Construct the path to the controller directory
directory = f"/home/judson/Neural-Networks-in-GNC/inverted_pendulum/training/average_normalized/{loss_function}/controllers" directory = f"/home/judson/Neural-Networks-in-GNC/inverted_pendulum/training/max_normalized/{loss_function}/controllers"
# Fetch the controller files according to the specified range and interval # Fetch the controller files according to the specified range and interval
controllers = get_controller_files(directory, epoch_range, epoch_step) controllers = get_controller_files(directory, epoch_range, epoch_step)
# Pack parameters for parallel processing # Pack parameters for parallel processing
@ -53,6 +53,7 @@ if __name__ == "__main__":
all_results[loss_function] = {} all_results[loss_function] = {}
all_results[loss_function][condition_name] = (epochs, theta_over_epochs) all_results[loss_function][condition_name] = (epochs, theta_over_epochs)
# continue
# Plotting the 3D epoch evolution # Plotting the 3D epoch evolution
print(f"Plotting the 3d epoch evolution for {loss_function} under {condition_text}") print(f"Plotting the 3d epoch evolution for {loss_function} under {condition_text}")
title = f"Pendulum Angle Evolution for {loss_function} and {condition_text}" title = f"Pendulum Angle Evolution for {loss_function} and {condition_text}"
@ -62,6 +63,8 @@ if __name__ == "__main__":
print("") print("")
# Plot the theta as a function of epoch for all loss functions # Plot the theta as a function of epoch for all loss functions
continue
specific_theta_index = num_steps // 2 specific_theta_index = num_steps // 2
save_path = os.path.join(condition_path, f"theta_at_5sec_across_epochs.png") save_path = os.path.join(condition_path, f"theta_at_5sec_across_epochs.png")
plot_theta_vs_epoch(all_results, condition_name, desired_theta, save_path, f"Theta at 5 Seconds across Epochs for {condition_text}", specific_theta_index) plot_theta_vs_epoch(all_results, condition_name, desired_theta, save_path, f"Theta at 5 Seconds across Epochs for {condition_text}", specific_theta_index)
@ -70,4 +73,9 @@ if __name__ == "__main__":
save_path = os.path.join(condition_path, f"final_theta_across_epochs.png") save_path = os.path.join(condition_path, f"final_theta_across_epochs.png")
plot_theta_vs_epoch(all_results, condition_name, desired_theta, save_path, f"Final Theta across Epochs for {condition_text}", specific_theta_index) plot_theta_vs_epoch(all_results, condition_name, desired_theta, save_path, f"Final Theta across Epochs for {condition_text}", specific_theta_index)
print(f"Completed plotting for all loss functions under {condition_name} condition.\n") print(f"Completed plotting for all loss functions under {condition_name} condition.\n")
import json
with open("all_results.json", 'w') as file:
json.dump(all_results, file)

Binary file not shown.

After

Width:  |  Height:  |  Size: 280 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 255 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 253 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 264 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 259 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 260 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 262 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 268 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 248 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 256 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 272 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 248 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 301 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 289 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 299 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 310 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 307 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 299 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 292 KiB

View File

@ -5,7 +5,7 @@ import matplotlib.pyplot as plt
import pandas as pd import pandas as pd
# Load Controller # Load Controller
controller_file_name = "controller_with_desired_theta.pth" controller_file_name = f"/home/judson/Neural-Networks-in-GNC/inverted_pendulum/training/controller_base.pth"
class PendulumController(nn.Module): class PendulumController(nn.Module):
def __init__(self): def __init__(self):
@ -93,31 +93,32 @@ in_sample_cases = [
(2/3 * 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), (-2/3 * np.pi, 0.0, 0.0, 0.0),
# Omega perturbations # # Omega perturbations
(0.0, 1/3 * np.pi, 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, -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, -2 * np.pi, 0.0, 0.0), # (0.0, -2 * np.pi, 0.0, 0.0),
# Return to non-zero theta # # Return to non-zero theta
(0.0, 0.0, 0.0, 2*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, -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/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),
(0.0, 0.0, 0.0, -1/3 *np.pi), # (0.0, 0.0, 0.0, -1/3 *np.pi),
# Mix cases # # Mix cases
(1/4 * np.pi, 1 * np.pi, 0.0, 0.0), # (1/4 * np.pi, 1 * np.pi, 0.0, 0.0),
(-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/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/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),
(-1/2 * np.pi, 1 * np.pi, 0.0, -4 *np.pi), # (-1/2 * np.pi, 1 * np.pi, 0.0, -4 *np.pi),
] ]
# Validation in-sample cases # Validation in-sample cases
print("Performing in-sample validation") print("Performing in-sample validation")
@ -212,25 +213,25 @@ print("\nPerforming out-of-sample validation")
# Out of sample cases previously generated by numpy # Out of sample cases previously generated by numpy
out_sample_cases = [ out_sample_cases = [
(-2.198958, -4.428501, 0.450833, 0.000000), (-2.198958, -4.428501, 0.450833, 0.000000),
(1.714196, -0.769896, 0.202738, 0.000000), # (1.714196, -0.769896, 0.202738, 0.000000),
(0.241195, -5.493715, 0.438996, 0.000000), # (0.241195, -5.493715, 0.438996, 0.000000),
(0.030605, 4.901513, -0.479243, 0.000000), # (0.030605, 4.901513, -0.479243, 0.000000),
(1.930445, -1.301926, -0.454050, 0.000000), # (1.930445, -1.301926, -0.454050, 0.000000),
(-0.676063, 4.246865, 0.036303, 0.000000), # (-0.676063, 4.246865, 0.036303, 0.000000),
(0.734920, -5.925202, 0.047097, 0.000000), # (0.734920, -5.925202, 0.047097, 0.000000),
(-3.074471, -3.535424, 0.315438, 0.000000), # (-3.074471, -3.535424, 0.315438, 0.000000),
(-0.094486, 6.111091, 0.150525, 0.000000), # (-0.094486, 6.111091, 0.150525, 0.000000),
(-1.647671, 5.720526, 0.334181, 0.000000), # (-1.647671, 5.720526, 0.334181, 0.000000),
(-2.611260, 5.087704, 0.045460, -3.610785), # (-2.611260, 5.087704, 0.045460, -3.610785),
(1.654137, 0.982081, -0.192725, 1.003872), # (1.654137, 0.982081, -0.192725, 1.003872),
(-2.394899, 3.550547, -0.430938, 3.261897), # (-2.394899, 3.550547, -0.430938, 3.261897),
(0.474917, 0.555166, -0.285173, 1.866752), # (0.474917, 0.555166, -0.285173, 1.866752),
(-0.640369, -4.678490, -0.340663, 3.150098), # (-0.640369, -4.678490, -0.340663, 3.150098),
(1.747517, -3.248204, -0.001520, 1.221787), # (1.747517, -3.248204, -0.001520, 1.221787),
(2.505283, -2.875006, -0.065617, -3.690269), # (2.505283, -2.875006, -0.065617, -3.690269),
(1.337244, 2.221707, 0.044979, -2.459730), # (1.337244, 2.221707, 0.044979, -2.459730),
(1.531012, 2.230981, -0.291206, -1.924535), # (1.531012, 2.230981, -0.291206, -1.924535),
(-1.065792, 4.320740, 0.075405, -1.550644), # (-1.065792, 4.320740, 0.075405, -1.550644),
] ]

View File

@ -12,14 +12,12 @@ from PendulumDynamics import PendulumDynamics
# Device setup # Device setup
device = torch.device("cpu") device = torch.device("cpu")
base_controller_path = f"/home/judson/Neural-Networks-in-GNC/inverted_pendulum/training/controller_base.pth"
# Initial conditions (theta0, omega0, alpha0, desired_theta) # Initial conditions (theta0, omega0, alpha0, desired_theta)
from initial_conditions import initial_conditions from initial_conditions import initial_conditions
state_0 = torch.tensor(initial_conditions, dtype=torch.float32, device=device) state_0 = torch.tensor(initial_conditions, dtype=torch.float32, device=device)
# Device setup
device = torch.device("cpu")
# Constants # Constants
m = 10.0 m = 10.0
g = 9.81 g = 9.81
@ -30,30 +28,15 @@ t_start, t_end, t_points = 0, 10, 1000
t_span = torch.linspace(t_start, t_end, t_points, device=device) t_span = torch.linspace(t_start, t_end, t_points, device=device)
# Specify directory for storing results # Specify directory for storing results
output_dir = "average_normalized" output_dir = "max_normalized"
os.makedirs(output_dir, exist_ok=True) os.makedirs(output_dir, exist_ok=True)
# Use a previously generated random seed # Optimizer values
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}")
# Initialize controller and dynamics
controller = PendulumController().to(device)
pendulum_dynamics = PendulumDynamics(controller, m, R, g).to(device)
# Optimizer setup
learning_rate = 1e-1 learning_rate = 1e-1
weight_decay = 1e-4 weight_decay = 1e-4
optimizer = optim.Adam(controller.parameters(), lr=learning_rate, weight_decay=weight_decay)
# Training parameters # Training parameters
num_epochs = 1001 num_epochs = 1000
# Define loss functions # Define loss functions
def make_loss_fn(weight_fn): def make_loss_fn(weight_fn):
@ -89,7 +72,7 @@ weight_functions = {
'description': 'Quadratic weight: Weights increase cubically from 0 to 1, normalized by the average weight' 'description': 'Quadratic weight: Weights increase cubically from 0 to 1, normalized by the average weight'
}, },
'inverse': { 'inverse': {
'function': lambda t: ((t+1)**-1 / ((t+1)**-1).max()) / ((t+1)**-2 / ((t+1)**-1).max()).mean(), 'function': lambda t: ((t+1)**-1 / ((t+1)**-1).max()) / ((t+1)**-1 / ((t+1)**-1).max()).mean(),
'description': 'Inverse weight: Weights decrease inversely, normalized by the average weight' 'description': 'Inverse weight: Weights decrease inversely, normalized by the average weight'
}, },
'inverse_squared': { 'inverse_squared': {
@ -105,7 +88,10 @@ weight_functions = {
# Training loop for each weight function # Training loop for each weight function
for name, weight_info in weight_functions.items(): for name, weight_info in weight_functions.items():
controller = PendulumController().to(device) controller = PendulumController().to(device)
controller.load_state_dict(torch.load(base_controller_path))
pendulum_dynamics = PendulumDynamics(controller, m, R, g).to(device) pendulum_dynamics = PendulumDynamics(controller, m, R, g).to(device)
print(f"Loaded {base_controller_path} as base controller")
optimizer = optim.Adam(controller.parameters(), lr=learning_rate, weight_decay=weight_decay) optimizer = optim.Adam(controller.parameters(), lr=learning_rate, weight_decay=weight_decay)
loss_fn = make_loss_fn(weight_info['function']) loss_fn = make_loss_fn(weight_info['function'])
@ -123,7 +109,7 @@ for name, weight_info in weight_functions.items():
# Overwrite configuration and log files # Overwrite configuration and log files
with open(config_file, "w") as f: with open(config_file, "w") as f:
f.write(f"Random Seed: {random_seed}\n") f.write(f"Base controller path: {base_controller_path}\n")
f.write(f"Time Span: {t_start} to {t_end}, Points: {t_points}\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"Learning Rate: {learning_rate}\n")
f.write(f"Weight Decay: {weight_decay}\n") f.write(f"Weight Decay: {weight_decay}\n")
@ -138,13 +124,20 @@ for name, weight_info in weight_functions.items():
with open(log_file, "w", newline="") as csvfile: with open(log_file, "w", newline="") as csvfile:
csv_writer = csv.writer(csvfile) csv_writer = csv.writer(csvfile)
csv_writer.writerow(["Epoch", "Loss"]) csv_writer.writerow(["Epoch", "Loss"])
# Training loop # Training loop
for epoch in range(num_epochs): for epoch in range(0, num_epochs+1):
optimizer.zero_grad() optimizer.zero_grad()
state_traj = odeint(pendulum_dynamics, state_0, t_span, method='rk4') state_traj = odeint(pendulum_dynamics, state_0, t_span, method='rk4')
loss = loss_fn(state_traj, t_span) loss = loss_fn(state_traj, t_span)
loss.backward() loss.backward()
# Save the model before training happens
model_file = os.path.join(controllers_dir, f"controller_{epoch}.pth")
torch.save(controller.state_dict(), model_file)
print(f"{model_file} saved with loss: {loss}")
# Update the weights and biases
optimizer.step() optimizer.step()
# Logging # Logging
@ -152,9 +145,4 @@ for name, weight_info in weight_functions.items():
csv_writer = csv.writer(csvfile) csv_writer = csv.writer(csvfile)
csv_writer.writerow([epoch, loss.item()]) csv_writer.writerow([epoch, loss.item()])
# Save the model
model_file = os.path.join(controllers_dir, f"controller_{epoch}.pth")
torch.save(controller.state_dict(), model_file)
print(f"{model_file} saved with loss: {loss}")
print("Training complete. Models and logs are saved under respective directories for each loss function.") print("Training complete. Models and logs are saved under respective directories for each loss function.")

View File

@ -0,0 +1,18 @@
import torch
import numpy as np
from PendulumController import PendulumController
device = torch.device("cpu")
controller = PendulumController().to(device)
# Use a previously generated random seed
random_seed = 4529
# Set the seeds for reproducibility
torch.manual_seed(random_seed)
np.random.seed(random_seed)
controller = PendulumController().to(device)
model_file = "controller_base.pth"
torch.save(controller.state_dict(), model_file)

View File

@ -1,4 +1,3 @@
Random Seed: 4529
Time Span: 0 to 10, Points: 1000 Time Span: 0 to 10, Points: 1000
Learning Rate: 0.1 Learning Rate: 0.1
Weight Decay: 0.0001 Weight Decay: 0.0001

File diff suppressed because it is too large Load Diff

View File

@ -1,4 +1,3 @@
Random Seed: 4529
Time Span: 0 to 10, Points: 1000 Time Span: 0 to 10, Points: 1000
Learning Rate: 0.1 Learning Rate: 0.1
Weight Decay: 0.0001 Weight Decay: 0.0001

File diff suppressed because it is too large Load Diff

View File

@ -1,4 +1,3 @@
Random Seed: 4529
Time Span: 0 to 10, Points: 1000 Time Span: 0 to 10, Points: 1000
Learning Rate: 0.1 Learning Rate: 0.1
Weight Decay: 0.0001 Weight Decay: 0.0001

File diff suppressed because it is too large Load Diff

View File

@ -1,4 +1,3 @@
Random Seed: 4529
Time Span: 0 to 10, Points: 1000 Time Span: 0 to 10, Points: 1000
Learning Rate: 0.1 Learning Rate: 0.1
Weight Decay: 0.0001 Weight Decay: 0.0001

File diff suppressed because it is too large Load Diff

View File

@ -1,4 +1,3 @@
Random Seed: 4529
Time Span: 0 to 10, Points: 1000 Time Span: 0 to 10, Points: 1000
Learning Rate: 0.1 Learning Rate: 0.1
Weight Decay: 0.0001 Weight Decay: 0.0001

File diff suppressed because it is too large Load Diff

View File

@ -1,4 +1,3 @@
Random Seed: 4529
Time Span: 0 to 10, Points: 1000 Time Span: 0 to 10, Points: 1000
Learning Rate: 0.1 Learning Rate: 0.1
Weight Decay: 0.0001 Weight Decay: 0.0001

File diff suppressed because it is too large Load Diff

View File

@ -1,4 +1,3 @@
Random Seed: 4529
Time Span: 0 to 10, Points: 1000 Time Span: 0 to 10, Points: 1000
Learning Rate: 0.1 Learning Rate: 0.1
Weight Decay: 0.0001 Weight Decay: 0.0001

File diff suppressed because it is too large Load Diff

View File

@ -12,14 +12,12 @@ from PendulumDynamics import PendulumDynamics
# Device setup # Device setup
device = torch.device("cpu") device = torch.device("cpu")
base_controller_path = f"/home/judson/Neural-Networks-in-GNC/inverted_pendulum/training/controller_base.pth"
# Initial conditions (theta0, omega0, alpha0, desired_theta) # Initial conditions (theta0, omega0, alpha0, desired_theta)
from initial_conditions import initial_conditions from initial_conditions import initial_conditions
state_0 = torch.tensor(initial_conditions, dtype=torch.float32, device=device) state_0 = torch.tensor(initial_conditions, dtype=torch.float32, device=device)
# Device setup
device = torch.device("cpu")
# Constants # Constants
m = 10.0 m = 10.0
g = 9.81 g = 9.81
@ -33,27 +31,12 @@ t_span = torch.linspace(t_start, t_end, t_points, device=device)
output_dir = "max_normalized" output_dir = "max_normalized"
os.makedirs(output_dir, exist_ok=True) os.makedirs(output_dir, exist_ok=True)
# Use a previously generated random seed # Optimizer values
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}")
# Initialize controller and dynamics
controller = PendulumController().to(device)
pendulum_dynamics = PendulumDynamics(controller, m, R, g).to(device)
# Optimizer setup
learning_rate = 1e-1 learning_rate = 1e-1
weight_decay = 1e-4 weight_decay = 1e-4
optimizer = optim.Adam(controller.parameters(), lr=learning_rate, weight_decay=weight_decay)
# Training parameters # Training parameters
num_epochs = 1001 num_epochs = 1000
# Define loss functions # Define loss functions
def make_loss_fn(weight_fn): def make_loss_fn(weight_fn):
@ -93,11 +76,11 @@ weight_functions = {
'description': 'Inverse weight: Weights decrease inversely, normalized by max' 'description': 'Inverse weight: Weights decrease inversely, normalized by max'
}, },
'inverse_squared': { 'inverse_squared': {
'function': lambda t: (t+1)**-2 / ((t+1)**-1).max(), 'function': lambda t: (t+1)**-2 / ((t+1)**-2).max(),
'description': 'Inverse squared weight: Weights decrease inversely squared, normalized by max' 'description': 'Inverse squared weight: Weights decrease inversely squared, normalized by max'
}, },
'inverse_cubed': { 'inverse_cubed': {
'function': lambda t: (t+1)**-3 / ((t+1)**-1).max(), 'function': lambda t: (t+1)**-3 / ((t+1)**-3).max(),
'description': 'Inverse cubed weight: Weights decrease inversely cubed, normalized by max' 'description': 'Inverse cubed weight: Weights decrease inversely cubed, normalized by max'
} }
} }
@ -105,7 +88,9 @@ weight_functions = {
# Training loop for each weight function # Training loop for each weight function
for name, weight_info in weight_functions.items(): for name, weight_info in weight_functions.items():
controller = PendulumController().to(device) controller = PendulumController().to(device)
controller.load_state_dict(torch.load(base_controller_path))
pendulum_dynamics = PendulumDynamics(controller, m, R, g).to(device) pendulum_dynamics = PendulumDynamics(controller, m, R, g).to(device)
print(f"Loaded {base_controller_path} as base controller")
optimizer = optim.Adam(controller.parameters(), lr=learning_rate, weight_decay=weight_decay) optimizer = optim.Adam(controller.parameters(), lr=learning_rate, weight_decay=weight_decay)
loss_fn = make_loss_fn(weight_info['function']) loss_fn = make_loss_fn(weight_info['function'])
@ -124,7 +109,7 @@ for name, weight_info in weight_functions.items():
# Overwrite configuration and log files # Overwrite configuration and log files
with open(config_file, "w") as f: with open(config_file, "w") as f:
f.write(f"Random Seed: {random_seed}\n") f.write(f"Base controller path: {base_controller_path}\n")
f.write(f"Time Span: {t_start} to {t_end}, Points: {t_points}\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"Learning Rate: {learning_rate}\n")
f.write(f"Weight Decay: {weight_decay}\n") f.write(f"Weight Decay: {weight_decay}\n")
@ -139,13 +124,21 @@ for name, weight_info in weight_functions.items():
with open(log_file, "w", newline="") as csvfile: with open(log_file, "w", newline="") as csvfile:
csv_writer = csv.writer(csvfile) csv_writer = csv.writer(csvfile)
csv_writer.writerow(["Epoch", "Loss"]) csv_writer.writerow(["Epoch", "Loss"])
# Training loop # Training loop
for epoch in range(num_epochs): for epoch in range(0, num_epochs+1):
optimizer.zero_grad() optimizer.zero_grad()
state_traj = odeint(pendulum_dynamics, state_0, t_span, method='rk4') state_traj = odeint(pendulum_dynamics, state_0, t_span, method='rk4')
loss = loss_fn(state_traj, t_span) loss = loss_fn(state_traj, t_span)
loss.backward() loss.backward()
# Save the model before training on this epoch
# Therefore, controller_epoch represents the controller after {epoch} training iterations
model_file = os.path.join(controllers_dir, f"controller_{epoch}.pth")
torch.save(controller.state_dict(), model_file)
print(f"{model_file} saved with loss: {loss}")
# Update the weights and biases
optimizer.step() optimizer.step()
# Logging # Logging
@ -153,9 +146,4 @@ for name, weight_info in weight_functions.items():
csv_writer = csv.writer(csvfile) csv_writer = csv.writer(csvfile)
csv_writer.writerow([epoch, loss.item()]) csv_writer.writerow([epoch, loss.item()])
# Save the model
model_file = os.path.join(controllers_dir, f"controller_{epoch}.pth")
torch.save(controller.state_dict(), model_file)
print(f"{model_file} saved with loss: {loss}")
print("Training complete. Models and logs are saved under respective directories for each loss function.") print("Training complete. Models and logs are saved under respective directories for each loss function.")

File diff suppressed because it is too large Load Diff