train.py

# this file reads the dataaset and creates the nn, then saves the pt file
# the last part evaluates the network on another dataset

import copy

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.optim as optim
import tqdm
from sklearn.model_selection import train_test_split


SCALE = 10000


df = pd.read_csv('orc/A3/single_pendulum_cost_10201.csv')
X = df[["q0", "dq0"]].values
y = df["cost"].values
y = y / SCALE




# train-test split for model evaluation
X_train, X_test, y_train, y_test = train_test_split(X, y, train_size=0.7, shuffle=True)

print(np.shape(X), np.shape(y))
print()

print(np.shape(X_train), np.shape(y_train))
print(np.shape(X_test), np.shape(y_test))
print()


# Convert to 2D PyTorch tensors
X_train = torch.tensor(X_train, dtype=torch.float32)
y_train = torch.tensor(y_train, dtype=torch.float32).reshape(-1, 1)
X_test = torch.tensor(X_test, dtype=torch.float32)
y_test = torch.tensor(y_test, dtype=torch.float32).reshape(-1, 1)

print(np.shape(X_train), np.shape(y_train))
print(np.shape(X_test), np.shape(y_test))


# Define the model
model = nn.Sequential(
    nn.Linear(2, 12),
    nn.ReLU(),
    nn.Linear(12, 12),
    nn.ReLU(),
    nn.Linear(12, 6),
    nn.ReLU(),
    nn.Linear(6, 1)
)

# loss function and optimizer
loss_fn = nn.MSELoss()  # mean square error
optimizer = optim.Adam(model.parameters(), lr=0.0001)

n_epochs = 100   # number of epochs to run
batch_size = 10  # size of each batch
batch_start = torch.arange(0, len(X_train), batch_size)


# Hold the best model
best_mse = np.inf   # init to infinity
best_weights = None
history = []

print("Training")
for epoch in range(n_epochs):
    model.train()
    with tqdm.tqdm(batch_start, unit="batch", mininterval=0, disable=False) as bar:
        bar.set_description(f"Epoch {epoch}")
        for start in bar:
            # take a batch
            X_batch = X_train[start:start+batch_size]
            y_batch = y_train[start:start+batch_size]
            # forward pass
            y_pred = model(X_batch)
            loss = loss_fn(y_pred, y_batch)
            # backward pass
            optimizer.zero_grad()
            loss.backward()
            # update weights
            optimizer.step()
            # print progress
            bar.set_postfix(mse=float(loss))
    # evaluate accuracy at end of each epoch
    model.eval()
    y_pred = model(X_test)
    mse = loss_fn(y_pred, y_test)
    mse = float(mse)
    history.append(mse)
    if mse < best_mse:
        best_mse = mse
        best_weights = copy.deepcopy(model.state_dict())

# restore model and return best accuracy
model.load_state_dict(best_weights)
print("MSE: %.2f" % best_mse)
print("RMSE: %.2f" % np.sqrt(best_mse))
plt.plot(history)
plt.show()

torch.save(model.state_dict(), "orc/A3/weights.pt")
exit()

model.eval()

df_rand = pd.read_csv('orc/A3/single_pendulum_cost_rand.csv')
X_rand = df_rand[["q0", "dq0"]].values
y_rand = df_rand["cost"].values
# y_rand = y_rand / SCALE

pred_rand = model(torch.tensor(X_rand, dtype=torch.float32)).detach().numpy() * SCALE
np.clip(pred_rand, 0.0, None, out=pred_rand)
# print(pred)


pred_rand = pred_rand.reshape(-1)
errors = np.abs((y_rand - pred_rand)) / np.mean(y_rand)
for i in range(len(y_rand)):
    x0 = X_rand[i]
    real_y = y_rand[i]
    predicted_y = pred_rand[i]
    error = errors[i]
    print(x0, "\t", real_y, "\t", predicted_y, "\t", error)

print("max error", max(errors))
errors = np.sort(errors, axis=0)
print(errors[-10:])