diff --git a/logs/example_experiment/2022/11_01/16_01_50_0/log.pkl b/logs/example_experiment/2022/11_01/16_01_50_0/log.pkl new file mode 100644 index 0000000..2b87f82 Binary files /dev/null and b/logs/example_experiment/2022/11_01/16_01_50_0/log.pkl differ diff --git a/scripts/actuator_net/__init__.py b/scripts/actuator_net/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/scripts/actuator_net/eval.py b/scripts/actuator_net/eval.py new file mode 100644 index 0000000..01e3154 --- /dev/null +++ b/scripts/actuator_net/eval.py @@ -0,0 +1,22 @@ +from utils import train_actuator_network_and_plot_predictions +from glob import glob + + +log_dir_root = "../../logs/" +log_dir = "example_experiment/2022/11_01/16_01_50_0" + +# Evaluates the existing actuator network by default +load_pretrained_model = True +actuator_network_path = "../../resources/actuator_nets/unitree_go1.pt" + +log_dirs = glob(f"{log_dir_root}{log_dir}/", recursive=True) + +if len(log_dirs) == 0: raise FileNotFoundError(f"No log files found in {log_dir_root}{log_dir}/") + +for log_dir in log_dirs: + try: + train_actuator_network_and_plot_predictions(log_dir[:11], log_dir[11:], actuator_network_path=actuator_network_path, load_pretrained_model=load_pretrained_model) + except FileNotFoundError: + print(f"Couldn't find log.pkl in {log_dir}") + except EOFError: + print(f"Incomplete log.pkl in {log_dir}") \ No newline at end of file diff --git a/scripts/actuator_net/train.py b/scripts/actuator_net/train.py new file mode 100644 index 0000000..ff6a9a0 --- /dev/null +++ b/scripts/actuator_net/train.py @@ -0,0 +1,26 @@ +from utils import train_actuator_network_and_plot_predictions +from glob import glob + +log_dir_root = "../../logs/" +log_dir = "example_experiment/2022/11_01/16_01_50_0" + +# Evaluates the existing actuator network by default +# load_pretrained_model = True +# actuator_network_path = "../../resources/actuator_nets/unitree_go1.pt" + +# Uncomment these lines to train a new actuator network +load_pretrained_model = False +actuator_network_path = "../../resources/actuator_nets/unitree_go1_new.pt" + + +log_dirs = glob(f"{log_dir_root}{log_dir}/", recursive=True) + +if len(log_dirs) == 0: raise FileNotFoundError(f"No log files found in {log_dir_root}{log_dir}/") + +for log_dir in log_dirs: + try: + train_actuator_network_and_plot_predictions(log_dir[:11], log_dir[11:], actuator_network_path=actuator_network_path, load_pretrained_model=load_pretrained_model) + except FileNotFoundError: + print(f"Couldn't find log.pkl in {log_dir}") + except EOFError: + print(f"Incomplete log.pkl in {log_dir}") \ No newline at end of file diff --git a/scripts/actuator_net/utils.py b/scripts/actuator_net/utils.py new file mode 100644 index 0000000..bae15f4 --- /dev/null +++ b/scripts/actuator_net/utils.py @@ -0,0 +1,232 @@ +import os +import pickle as pkl +from matplotlib import pyplot as plt +import time +import imageio +import numpy as np +from tqdm import tqdm +from glob import glob + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.utils.data import Dataset +from torch.utils.data import DataLoader +from torch.optim import Adam + +class ActuatorDataset(Dataset): + def __init__(self, data): + self.data = data + + def __len__(self): + return len(self.data['joint_states']) + + def __getitem__(self, idx): + return {k: v[idx] for k,v in self.data.items()} + +class Act(nn.Module): + def __init__(self, act, slope=0.05): + super(Act, self).__init__() + self.act = act + self.slope = slope + self.shift = torch.log(torch.tensor(2.0)).item() + + def forward(self, input): + if self.act == "relu": + return F.relu(input) + elif self.act == "leaky_relu": + return F.leaky_relu(input) + elif self.act == "sp": + return F.softplus(input, beta=1.) + elif self.act == "leaky_sp": + return F.softplus(input, beta=1.) - self.slope * F.relu(-input) + elif self.act == "elu": + return F.elu(input, alpha=1.) + elif self.act == "leaky_elu": + return F.elu(input, alpha=1.) - self.slope * F.relu(-input) + elif self.act == "ssp": + return F.softplus(input, beta=1.) - self.shift + elif self.act == "leaky_ssp": + return ( + F.softplus(input, beta=1.) - + self.slope * F.relu(-input) - + self.shift + ) + elif self.act == "tanh": + return torch.tanh(input) + elif self.act == "leaky_tanh": + return torch.tanh(input) + self.slope * input + elif self.act == "swish": + return torch.sigmoid(input) * input + elif self.act == "softsign": + return F.softsign(input) + else: + raise RuntimeError(f"Undefined activation called {self.act}") + +def build_mlp(in_dim, units, layers, out_dim, + act='relu', layer_norm=False, act_final=False): + mods = [nn.Linear(in_dim, units), Act(act)] + for i in range(layers-1): + mods += [nn.Linear(units, units), Act(act)] + mods += [nn.Linear(units, out_dim)] + if act_final: + mods += [Act(act)] + if layer_norm: + mods += [nn.LayerNorm(out_dim)] + return nn.Sequential(*mods) + +def train_actuator_network(xs, ys, actuator_network_path): + + print(xs.shape, ys.shape) + + num_data = xs.shape[0] + num_train = num_data // 5 * 4 + num_test = num_data - num_train + + dataset = ActuatorDataset({"joint_states": xs, "tau_ests": ys}) + train_set, val_set = torch.utils.data.random_split(dataset, [num_train, num_test]) + train_loader = DataLoader(train_set, batch_size=128, shuffle=True) + test_loader = DataLoader(val_set, batch_size=128, shuffle=True) + + model = build_mlp(in_dim=6, units=32, layers=2, out_dim=1, act='softsign') + + lr = 8e-4 + opt = Adam(model.parameters(), lr=lr, eps=1e-8, weight_decay=0.0) + + epochs = 100 + device = 'cuda:0' + + model = model.to(device) + for epoch in range(epochs): + epoch_loss = 0 + ct = 0 + for batch in train_loader: + data = batch['joint_states'].to(device) + y_pred = model(data) + + opt.zero_grad() + + y_label = batch['tau_ests'].to(device) + + tau_est_loss = ((y_pred - y_label) ** 2).mean() + loss = tau_est_loss + + loss.backward() + opt.step() + epoch_loss += loss.detach().cpu().numpy() + ct += 1 + epoch_loss /= ct + + test_loss = 0 + mae = 0 + ct = 0 + if epoch % 1 == 0: + with torch.no_grad(): + for batch in test_loader: + data = batch['joint_states'].to(device) + y_pred = model(data) + + y_label = batch['tau_ests'].to(device) + + tau_est_loss = ((y_pred - y_label) ** 2).mean() + loss = tau_est_loss + test_mae = (y_pred - y_label).abs().mean() + + test_loss += loss + mae += test_mae + ct += 1 + test_loss /= ct + mae /= ct + + print( + f'epoch: {epoch} | loss: {epoch_loss:.4f} | test loss: {test_loss:.4f} | mae: {mae:.4f}') + + model_scripted = torch.jit.script(model) # Export to TorchScript + model_scripted.save(actuator_network_path) # Save + return model + +def train_actuator_network_and_plot_predictions(log_dir_root, log_dir, actuator_network_path, load_pretrained_model=False): + + log_path = log_dir_root + log_dir + "log.pkl" + print(log_path) + with open(log_path, 'rb') as file: + data = pkl.load(file) + + datas = data['hardware_closed_loop'][1] + + if len(datas) < 1: + return + + tau_ests = np.zeros((len(datas), 12)) + torques = np.zeros((len(datas), 12)) + joint_positions = np.zeros((len(datas), 12)) + joint_position_targets = np.zeros((len(datas), 12)) + joint_velocities = np.zeros((len(datas), 12)) + + if "tau_est" not in datas[0].keys(): + return + + for i in range(len(datas)): + tau_ests[i, :] = datas[i]["tau_est"] + torques[i, :] = datas[i]["torques"] + joint_positions[i, :] = datas[i]["joint_pos"] + joint_position_targets[i, :] = datas[i]["joint_pos_target"] + joint_velocities[i, :] = datas[i]["joint_vel"] + + timesteps = np.array(range(len(datas))) / 50.0 + + import matplotlib.pyplot as plt + + joint_position_errors = joint_positions - joint_position_targets + joint_velocities = joint_velocities + + joint_position_errors = torch.tensor(joint_position_errors, dtype=torch.float) + joint_velocities = torch.tensor(joint_velocities, dtype=torch.float) + tau_ests = torch.tensor(tau_ests, dtype=torch.float) + + xs = [] + ys = [] + step = 2 + # all joints are equal + for i in range(12): + xs_joint = [joint_position_errors[2:-step+1, i:i+1], + joint_position_errors[1:-step, i:i+1], + joint_position_errors[:-step-1, i:i+1], + joint_velocities[2:-step+1, i:i+1], + joint_velocities[1:-step, i:i+1], + joint_velocities[:-step-1, i:i+1]] + + tau_ests_joint = [tau_ests[step:-1, i:i+1]] + + xs_joint = torch.cat(xs_joint, dim=1) + xs += [xs_joint] + ys += tau_ests_joint + + xs = torch.cat(xs, dim=0) + ys = torch.cat(ys, dim=0) + + if load_pretrained_model: + model = torch.jit.load(actuator_network_path).to('cpu') + else: + model = train_actuator_network(xs, ys, actuator_network_path).to("cpu") + + tau_preds = model(xs).detach().reshape(12, -1).T + + plot_length = 300 + + timesteps = timesteps[:plot_length] + torques = torques[step:plot_length+step] + tau_ests = tau_ests[step:plot_length+step] + tau_preds = tau_preds[:plot_length] + + fig, axs = plt.subplots(6, 2, figsize=(14, 6)) + axs = np.array(axs).flatten() + for i in range(12): + axs[i].plot(timesteps, torques[:, i], label="idealized torque") + axs[i].plot(timesteps, tau_ests[:, i], label="true torque") + axs[i].plot(timesteps, tau_preds[:, i], linestyle='--', label="actuator model predicted torque") + plt.legend() + + plt.show() + +