update ik learning script

stack/main/scripts/learn_neural_ik.py

@@ -1,6 +1,7 @@
 import torch
 import torch.nn as nn
 import torch.optim as optim
+from torch.nn.utils import spectral_norm
 from torch.utils.data import Dataset, DataLoader, random_split
 import numpy as np
 import pandas as pd
@@ -22,8 +23,17 @@ def __getitem__(self, idx):
 
 
 # Load data
-us_df = pd.read_csv('../data/trajectories/steady_state/control_inputs_uniform.csv')
-ys_df = pd.read_csv('../data/trajectories/steady_state/observations_steady_state_src_demo_17oct24.csv')
+us_df = pd.read_csv('../data/trajectories/steady_state/control_inputs_beta.csv')
+us_df = pd.concat([us_df, pd.read_csv('../data/trajectories/steady_state/control_inputs_uniform_pt1.csv')])
+us_df = pd.concat([us_df, pd.read_csv('../data/trajectories/steady_state/control_inputs_uniform_pt2.csv')])
+us_df = pd.concat([us_df, pd.read_csv('../data/trajectories/steady_state/control_inputs_uniform_pt3.csv')])
+
+# Observations
+ys_df = pd.read_csv('../data/trajectories/steady_state/observations_steady_state_beta_src_demo.csv')
+ys_df = pd.concat([ys_df, pd.read_csv('../data/trajectories/steady_state/observations_steady_state_uniform_src_demo_pt1.csv')])
+ys_df = pd.concat([ys_df, pd.read_csv('../data/trajectories/steady_state/observations_steady_state_uniform_src_demo_pt2.csv')])
+ys_df = pd.concat([ys_df, pd.read_csv('../data/trajectories/steady_state/observations_steady_state_uniform_src_demo_pt3.csv')])
+
 rest_positions = np.array([0.1005, -0.10698, 0.10445, -0.10302, -0.20407, 0.10933, 0.10581, -0.32308, 0.10566])
 ys_df = ys_df - rest_positions
 N = len(ys_df)
@@ -36,45 +46,79 @@ def __getitem__(self, idx):
 
 # Randomly split dataset into train and test
 dataset = UYPairsDataset(us, ys)
-train_size = int(0.85 * N)
-test_size = N - train_size
-train_dataset, test_dataset = random_split(dataset, [train_size, test_size])
+train_size = int(0.8 * N)
+val_size = int(0.1 * N)
+test_size = N - train_size - val_size
+train_dataset, val_dataset, test_dataset = random_split(dataset, [train_size, val_size, test_size])
 
 # Create data loaders
 train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
+val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False)
 test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)
 
-
-# Neural Network
+    
+# Spectrally-normalized Relu MLP
 class MLP(nn.Module):
-    def __init__(self):
+    def __init__(self,
+                 num_inputs: int = 9,
+                 num_outputs: int = 6,
+                 num_neurons: list = [32, 32],
+                 act: nn.Module = nn.ReLU(),
+                 spectral_normalize: bool = False):
         super(MLP, self).__init__()
-        self.fc1 = nn.Linear(9, 32)
-        self.fc2 = nn.Linear(32, 32)
-        self.fc3 = nn.Linear(32, 6)
+
+        layers = []
+
+        # Input layer
+        if spectral_normalize:
+            input_layer = spectral_norm(nn.Linear(num_inputs, num_neurons[0]))
+        else:
+            input_layer = nn.Linear(num_inputs, num_neurons[0])
+        layers.append(input_layer)
+        layers.append(act)
+
+        # Hidden layers
+        for i in range(len(num_neurons) - 1):
+            if spectral_normalize:
+                hidden_layer = spectral_norm(nn.Linear(num_neurons[i], num_neurons[i + 1]))
+            else:
+                hidden_layer = nn.Linear(num_neurons[i], num_neurons[i + 1])
+            layers.append(hidden_layer)
+            layers.append(act)
+
+        # Output layer
+        if spectral_normalize:
+            output_layer = spectral_norm(nn.Linear(num_neurons[-1], num_outputs))
+        else:
+            output_layer = nn.Linear(num_neurons[-1], num_outputs)
+
+        layers.append(output_layer)
+        self.layers = nn.ModuleList(layers)
 
     def forward(self, y):
-        x = torch.relu(self.fc1(y))
-        x = torch.relu(self.fc2(x))
-        x = self.fc3(x)
-        return x
+        x = y
+        for layer in self.layers:
+            x = layer(x)
+        u = x
+        return u
+
 
 
 # Instantiate the model, loss function, and optimizer
-model = MLP()
+inverse_kinematic_model = MLP()
 criterion = nn.MSELoss()
-optimizer = optim.Adam(model.parameters(), lr=0.0025, weight_decay=1e-6)
+optimizer = optim.Adam(inverse_kinematic_model.parameters(), lr=0.0025, weight_decay=1e-6)
 
 # Training loop
 num_epochs = 10000
 for epoch in range(num_epochs):
-    model.train()
+    inverse_kinematic_model.train()
     running_loss = 0.0
     for ys_batch, us_batch in train_loader:
         optimizer.zero_grad()  # Zero the gradients
 
         # Forward pass
-        outputs = model(ys_batch)
+        outputs = inverse_kinematic_model(ys_batch)
         loss = criterion(outputs, us_batch)
 
         # Backward pass and optimization
@@ -87,11 +131,11 @@ def forward(self, y):
         print(f'Epoch [{epoch+1}/{num_epochs}], RMSE: {np.sqrt(running_loss/len(train_loader)):.4f}')
 
 # Evaluation on test dataset
-model.eval()
+inverse_kinematic_model.eval()
 test_loss = 0.0
 with torch.no_grad():
     for ys_batch, us_batch in test_loader:
-        outputs = model(ys_batch)
+        outputs = inverse_kinematic_model(ys_batch)
         loss = criterion(outputs, us_batch)
         test_loss += loss.item()
 
@@ -101,7 +145,7 @@ def forward(self, y):
 u_min, u_max = -0.35, 0.35
 test_loss_random = 0.0
 mse_criterion = nn.MSELoss()
-model.eval()
+inverse_kinematic_model.eval()
 with torch.no_grad():
     for ys_batch, us_batch in test_loader:
         random_preds = torch.FloatTensor(us_batch.size()).uniform_(u_min, u_max)