content: 
  import numpy as np
  import torch
  from torch import nn
  from torch.nn import functional as F
  from torch.utils.data import DataLoader
  import os
  import h5py
  import math
  import wandb
  import pickle
  import matplotlib.pyplot as plt
  from copy import deepcopy
  from tqdm import tqdm
  from utils import find_all_hdf5
  from imitate_episodes import repeater, compute_dict_mean
  
  import IPython
  e = IPython.embed
  
  def main():
      ### Idea
      # input : o o o o o o # observed speed 
      # target: a a a a a a # commanded speed
      # at test time, input desired speed profile and convert that to command
  
      #########################################################
      history_len = 50
      future_len = 50
      prediction_len = 50
      batch_size_train = 16
      batch_size_val  = 16
      lr = 1e-4
      weight_decay = 1e-4
  
      num_steps = 10000
      validate_every = 2000
      save_every = 2000
  
      expr_name = f'actuator_network_test_{history_len}_{future_len}_{prediction_len}'
      ckpt_dir = f'/scr/tonyzhao/train_logs/{expr_name}' if os.getlogin() == 'tonyzhao' else f'./ckpts/{expr_name}'
      dataset_dir = '/scr/tonyzhao/compressed_datasets/aloha_mobile_fork/' if os.getlogin() == 'tonyzhao' else '/home/zfu/data/aloha_mobile_fork/'
      #########################################################
      assert(history_len + future_len >= prediction_len)
      assert(future_len % prediction_len == 0)
  
      wandb.init(project="mobile-aloha2", reinit=True, entity="mobile-aloha2", name=expr_name) # mode='disabled', 
  
      if not os.path.isdir(ckpt_dir):
          os.makedirs(ckpt_dir)
  
      dataset_path_list = find_all_hdf5(dataset_dir, skip_mirrored_data=True)
      dataset_path_list = [n for n in dataset_path_list if 'replayed' in n]
      num_episodes = len(dataset_path_list)
  
      # obtain train test split
      train_ratio = 0.9
      shuffled_episode_ids = np.random.permutation(num_episodes)
      train_episode_ids = shuffled_episode_ids[:int(train_ratio * num_episodes)]
      val_episode_ids = shuffled_episode_ids[int(train_ratio * num_episodes):]
      print(f'\n\nData from: {dataset_dir}\n- Train on {len(train_episode_ids)} episodes\n- Test on {len(val_episode_ids)} episodes\n\n')
  
      # obtain normalization stats for qpos and action
      # if load_pretrain:
      #     with open(os.path.join('/home/zfu/interbotix_ws/src/act/ckpts/pretrain_all', 'dataset_stats.pkl'), 'rb') as f:
      #         norm_stats = pickle.load(f)
      #     print('Loaded pretrain dataset stats')
      norm_stats, all_episode_len = get_norm_stats(dataset_path_list)
      train_episode_len = [all_episode_len[i] for i in train_episode_ids]
      val_episode_len = [all_episode_len[i] for i in val_episode_ids]
      assert(all_episode_len[0] % prediction_len == 0)
  
      # save dataset stats
      stats_path = os.path.join(ckpt_dir, f'actuator_net_stats.pkl')
      with open(stats_path, 'wb') as f:
          pickle.dump(norm_stats, f)
  
      # construct dataset and dataloader
      train_dataset = EpisodicDataset(dataset_path_list, norm_stats, train_episode_ids, train_episode_len, history_len, future_len, prediction_len)
      val_dataset = EpisodicDataset(dataset_path_list, norm_stats, val_episode_ids, val_episode_len, history_len, future_len, prediction_len)
      train_dataloader = DataLoader(train_dataset, batch_size=batch_size_train, shuffle=True, pin_memory=True, num_workers=1, prefetch_factor=1)
      val_dataloader = DataLoader(val_dataset, batch_size=batch_size_val, shuffle=True, pin_memory=True, num_workers=1, prefetch_factor=1)
      
      policy = ActuatorNetwork(prediction_len).cuda()
      optimizer = torch.optim.AdamW(policy.parameters(), lr=lr, weight_decay=weight_decay)
  
      n_parameters = sum(p.numel() for p in policy.parameters() if p.requires_grad)
      print("number of parameters: %.2fM" % (n_parameters/1e6,))
  
      min_val_loss = np.inf
      best_ckpt_info = None
      train_dataloader = repeater(train_dataloader)
      for step in tqdm(range(num_steps+1)):
          # validation
          if step % validate_every == 0:
              print('validating')
  
              with torch.inference_mode():
                  policy.eval()
                  validation_dicts = []
                  for batch_idx, data in enumerate(val_dataloader):
                      observed_speed, commanded_speed = data
                      out, forward_dict = policy(observed_speed.cuda(), commanded_speed.cuda())
                      validation_dicts.append(forward_dict)
  
                  validation_summary = compute_dict_mean(validation_dicts)
  
                  epoch_val_loss = validation_summary['loss']
                  if epoch_val_loss < min_val_loss:
                      min_val_loss = epoch_val_loss
                      best_ckpt_info = (step, min_val_loss, deepcopy(policy.state_dict()))
              for k in list(validation_summary.keys()):
                  validation_summary[f'val_{k}'] = validation_summary.pop(k)            
              wandb.log(validation_summary, step=step)
              print(f'Val loss:   {epoch_val_loss:.5f}')
              summary_string = ''
              for k, v in validation_summary.items():
                  summary_string += f'{k}: {v.item():.3f} '
              print(summary_string)
  
              visualize_prediction(dataset_path_list, val_episode_ids, policy, norm_stats, history_len, future_len, prediction_len, ckpt_dir, step, 'val')
              visualize_prediction(dataset_path_list, train_episode_ids, policy, norm_stats, history_len, future_len, prediction_len, ckpt_dir, step, 'train')
  
  
          # training
          policy.train()
          optimizer.zero_grad()
          data = next(train_dataloader)
          observed_speed, commanded_speed = data
          out, forward_dict = policy(observed_speed.cuda(), commanded_speed.cuda())
          # backward
          loss = forward_dict['loss']
          loss.backward()
          optimizer.step()
          wandb.log(forward_dict, step=step) # not great, make training 1-2% slower
  
          if step % save_every == 0:
              ckpt_path = os.path.join(ckpt_dir, f'actuator_net_step_{step}.ckpt')
              torch.save(policy.state_dict(), ckpt_path)
  
      ckpt_path = os.path.join(ckpt_dir, f'actuator_net_last.ckpt')
      torch.save(policy.state_dict(), ckpt_path)
  
      best_step, min_val_loss, best_state_dict = best_ckpt_info
      ckpt_path = os.path.join(ckpt_dir, f'actuator_net_step_{best_step}.ckpt')
      torch.save(best_state_dict, ckpt_path)
      print(f'Training finished:\nval loss {min_val_loss:.6f} at step {best_step}')
  
  
  def visualize_prediction(dataset_path_list, episode_ids, policy, norm_stats, history_len, future_len, prediction_len, ckpt_dir, step, name):
      num_vis = 2
      episode_ids = episode_ids[:num_vis]
      vis_path = [dataset_path_list[i] for i in episode_ids]
  
      for i, dataset_path in enumerate(vis_path):
          try:
              with h5py.File(dataset_path, 'r') as root:
                  commanded_speed = root['/base_action'][()]
                  observed_speed = root['/obs_tracer'][()]
          except Exception as ee:
              print(f'Error loading {dataset_path} in get_norm_stats')
              print(ee)
              quit()
          
          # commanded_speed = (commanded_speed - norm_stats["commanded_speed_mean"]) / norm_stats["commanded_speed_std"]
          norm_observed_speed = (observed_speed - norm_stats["observed_speed_mean"]) / norm_stats["observed_speed_std"]
          out_unnorm_fn = lambda x: (x * norm_stats["commanded_speed_std"]) + norm_stats["commanded_speed_mean"]
  
          history_pad = np.zeros((history_len, 2))
          future_pad = np.zeros((future_len, 2))
          norm_observed_speed = np.concatenate([history_pad, norm_observed_speed, future_pad], axis=0)
  
          episode_len = commanded_speed.shape[0]
  
          all_pred = []
          for t in range(0, episode_len, prediction_len):
              offset_start_ts = t + history_len
              policy_input = norm_observed_speed[offset_start_ts-history_len: offset_start_ts+future_len]
              policy_input = torch.from_numpy(policy_input).float().unsqueeze(dim=0).cuda()
              pred = policy(policy_input)
              pred = pred.detach().cpu().numpy()[0]
              all_pred += out_unnorm_fn(pred).tolist()
          all_pred = np.array(all_pred)
  
          plot_path = os.path.join(ckpt_dir, f'{name}{i}_step{step}_linear')
          plt.figure()
          plt.plot(commanded_speed[:, 0], label='commanded_speed_linear')
          plt.plot(observed_speed[:, 0], label='observed_speed_linear')
          plt.plot(all_pred[:, 0],  label='pred_commanded_speed_linear')
          # plot vertical grey dotted lines every prediction_len
          for t in range(0, episode_len, prediction_len):
              plt.axvline(t, linestyle='--', color='grey')
          plt.legend()
          plt.savefig(plot_path)
          plt.close()
  
          plot_path = os.path.join(ckpt_dir, f'{name}{i}_step{step}_angular')
          plt.figure()
          plt.plot(commanded_speed[:, 1], label='commanded_speed_angular')
          plt.plot(observed_speed[:, 1], label='observed_speed_angular')
          plt.plot(all_pred[:, 1], label='pred_commanded_speed_angular')
          # plot vertical dotted lines every prediction_len
          for t in range(0, episode_len, prediction_len):
              plt.axvline(t, linestyle='--', color='grey')
          plt.legend()
          plt.savefig(plot_path)
          plt.close()
  
  
  
  class ActuatorNetwork(nn.Module):
  
      def __init__(self, prediction_len):
          super().__init__()
          d_model = 256
          encoder_layer = nn.TransformerEncoderLayer(d_model=d_model, nhead=8)
          self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=3)
          self.pe = PositionalEncoding(d_model)
          self.in_proj = nn.Linear(2, d_model)
          self.out_proj = nn.Linear(d_model, 2)
          self.prediction_len = prediction_len
  
      def forward(self, src, tgt=None):
          if tgt is not None: # training time
              # (batch, seq, feature) -> (seq, batch, feature)
              src = self.in_proj(src)
              src = torch.einsum('b s d -> s b d', src)
              src = self.pe(src)
              out = self.transformer(src)
              
              tgt = torch.einsum('b s d -> s b d', tgt)
              assert(self.prediction_len == tgt.shape[0])
              out = out[0: self.prediction_len] # take first few tokens only for prediction
              out = self.out_proj(out)
  
              l2_loss = loss = F.mse_loss(out, tgt)
              loss_dict = {'loss': l2_loss}
              out = torch.einsum('s b d -> b s d', out)
              return out, loss_dict
          else:
              src = self.in_proj(src)
              src = torch.einsum('b s d -> s b d', src)
              src = self.pe(src)
              out = self.transformer(src)
              out = out[0: self.prediction_len] # take first few tokens only for prediction
              out = self.out_proj(out)
              out = torch.einsum('s b d -> b s d', out)
              return out
  
  
  
  class PositionalEncoding(nn.Module):
      def __init__(self, d_model: int, dropout: float = 0.1, max_len: int = 5000):
          super().__init__()
          self.dropout = nn.Dropout(p=dropout)
          position = torch.arange(max_len).unsqueeze(1)
          div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))
          pe = torch.zeros(max_len, 1, d_model)
          pe[:, 0, 0::2] = torch.sin(position * div_term)
          pe[:, 0, 1::2] = torch.cos(position * div_term)
          self.register_buffer('pe', pe)
  
      def forward(self, x):
          """
          Arguments:
              x: Tensor, shape ``[seq_len, batch_size, embedding_dim]``
          """
          x = x + self.pe[:x.size(0)]
          return self.dropout(x)
  
  def get_norm_stats(dataset_path_list):
      all_commanded_speed = []
      all_observed_speed = []
      all_episode_len = []
      for dataset_path in dataset_path_list:
          try:
              with h5py.File(dataset_path, 'r') as root:
                  commanded_speed = root['/base_action'][()]
                  observed_speed = root['/obs_tracer'][()]
          except Exception as e:
              print(f'Error loading {dataset_path} in get_norm_stats')
              print(e)
              quit()
          all_commanded_speed.append(torch.from_numpy(commanded_speed))
          all_observed_speed.append(torch.from_numpy(observed_speed))
          all_episode_len.append(len(commanded_speed))
      all_commanded_speed = torch.cat(all_commanded_speed, dim=0)
      all_observed_speed = torch.cat(all_observed_speed, dim=0)
  
      # normalize all_commanded_speed
      commanded_speed_mean = all_commanded_speed.mean(dim=[0]).float()
      commanded_speed_std = all_commanded_speed.std(dim=[0]).float()
      commanded_speed_std = torch.clip(commanded_speed_std, 1e-2, np.inf) # clipping
  
      # normalize all_observed_speed
      observed_speed_mean = all_observed_speed.mean(dim=[0]).float()
      observed_speed_std = all_observed_speed.std(dim=[0]).float()
      observed_speed_std = torch.clip(observed_speed_std, 1e-2, np.inf) # clipping
  
      stats = {"commanded_speed_mean": commanded_speed_mean.numpy(), "commanded_speed_std": commanded_speed_std.numpy(),
               "observed_speed_mean": observed_speed_mean.numpy(), "observed_speed_std": observed_speed_std.numpy()}
  
      return stats, all_episode_len
  
  
  class EpisodicDataset(torch.utils.data.Dataset):
      def __init__(self, dataset_path_list, norm_stats, episode_ids, episode_len, history_len, future_len, prediction_len):
          super(EpisodicDataset).__init__()
          self.episode_ids = episode_ids
          self.dataset_path_list = dataset_path_list
          self.norm_stats = norm_stats
          self.episode_len = episode_len
          self.cumulative_len = np.cumsum(self.episode_len)
          self.max_episode_len = max(episode_len)
          self.history_len = history_len
          self.future_len = future_len
          self.prediction_len = prediction_len
          self.is_sim = False
          self.history_pad = np.zeros((self.history_len, 2))
          self.future_pad = np.zeros((self.future_len, 2))
          self.prediction_pad = np.zeros((self.prediction_len, 2))
          self.__getitem__(0) # initialize self.is_sim
  
      def __len__(self):
          return sum(self.episode_len)
  
      def _locate_transition(self, index):
          assert index < self.cumulative_len[-1]
          episode_index = np.argmax(self.cumulative_len > index) # argmax returns first True index
          start_ts = index - (self.cumulative_len[episode_index] - self.episode_len[episode_index])
          episode_id = self.episode_ids[episode_index]
          return episode_id, start_ts
  
      def __getitem__(self, index):
          episode_id, start_ts = self._locate_transition(index)
          dataset_path = self.dataset_path_list[episode_id]
          try:
              # print(dataset_path)
              with h5py.File(dataset_path, 'r') as root:
                  commanded_speed = root['/base_action'][()]
                  observed_speed = root['/obs_tracer'][()]
                  observed_speed = np.concatenate([self.history_pad, observed_speed, self.future_pad], axis=0)
                  commanded_speed = np.concatenate([commanded_speed, self.prediction_pad], axis=0)
  
                  offset_start_ts = start_ts + self.history_len
                  commanded_speed = commanded_speed[start_ts: start_ts+self.prediction_len]
                  observed_speed = observed_speed[offset_start_ts-self.history_len: offset_start_ts+self.future_len]
  
              commanded_speed = torch.from_numpy(commanded_speed).float()
              observed_speed = torch.from_numpy(observed_speed).float()
  
              # normalize to mean 0 std 1
              commanded_speed = (commanded_speed - self.norm_stats["commanded_speed_mean"]) / self.norm_stats["commanded_speed_std"]
              observed_speed = (observed_speed - self.norm_stats["observed_speed_mean"]) / self.norm_stats["observed_speed_std"]
  
          except:
              print(f'Error loading {dataset_path} in __getitem__')
              quit()
  
          # print(image_data.dtype, qpos_data.dtype, action_data.dtype, is_pad.dtype)
          return observed_speed, commanded_speed
  
  
  
  
  if __name__ == '__main__':
      main()
  