inference_low_trans.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "pybullet build time: Nov 28 2023 23:45:17\n",
      "/usr/lib/python3/dist-packages/scipy/__init__.py:146: UserWarning: A NumPy version >=1.17.3 and <1.25.0 is required for this version of SciPy (detected version 1.26.0\n",
      "  warnings.warn(f\"A NumPy version >={np_minversion} and <{np_maxversion}\"\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[33mROBOMIMIC WARNING(\n",
      "    No private macro file found!\n",
      "    It is recommended to use a private macro file\n",
      "    To setup, run: python /home/franka_deoxys/robomimic/robomimic/scripts/setup_macros.py\n",
      ")\u001b[0m\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/franka_deoxys/.local/lib/python3.10/site-packages/matplotlib/projections/__init__.py:63: UserWarning: Unable to import Axes3D. This may be due to multiple versions of Matplotlib being installed (e.g. as a system package and as a pip package). As a result, the 3D projection is not available.\n",
      "  warnings.warn(\"Unable to import Axes3D. This may be due to multiple versions of \"\n"
     ]
    }
   ],
   "source": [
    "import argparse\n",
    "import json\n",
    "import os\n",
    "import pprint\n",
    "import sys\n",
    "from pathlib import Path\n",
    "\n",
    "import cv2\n",
    "import h5py\n",
    "# import hydra\n",
    "import numpy as np\n",
    "import torch\n",
    "import torchvision\n",
    "import yaml\n",
    "from easydict import EasyDict\n",
    "# from hydra.experimental import compose, initialize\n",
    "# from omegaconf import DictConfig, OmegaConf\n",
    "from PIL import Image\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "from torchvision import transforms\n",
    "\n",
    "from deoxys import config_root\n",
    "# Control import\n",
    "from deoxys.franka_interface import FrankaInterface\n",
    "# from deoxys.utils import RobotStateRawObsDictGenerator, YamlConfig\n",
    "from deoxys.utils import  YamlConfig\n",
    "from deoxys.utils.config_utils import robot_config_parse_args\n",
    "from deoxys.utils.input_utils import input2action\n",
    "from deoxys.utils.io_devices import SpaceMouse\n",
    "from deoxys.utils.log_utils import get_deoxys_example_logger\n",
    "from deoxys.experimental.motion_utils import follow_joint_traj, reset_joints_to\n",
    "\n",
    "\n",
    "import argparse\n",
    "import json\n",
    "import numpy as np\n",
    "import time\n",
    "import os\n",
    "import shutil\n",
    "import psutil\n",
    "import sys\n",
    "import socket\n",
    "import traceback\n",
    "\n",
    "from collections import OrderedDict\n",
    "\n",
    "import torch\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "import robomimic\n",
    "import robomimic.utils.train_utils as TrainUtils\n",
    "import robomimic.utils.torch_utils as TorchUtils\n",
    "import robomimic.utils.obs_utils as ObsUtils\n",
    "import robomimic.utils.env_utils as EnvUtils\n",
    "import robomimic.utils.file_utils as FileUtils\n",
    "from robomimic.config import config_factory\n",
    "from robomimic.algo import algo_factory, RolloutPolicy\n",
    "from robomimic.utils.log_utils import PrintLogger, DataLogger, flush_warnings\n",
    "\n",
    "\n",
    "import robomimic.utils.train_utils as TrainUtils\n",
    "import robomimic.utils.torch_utils as TorchUtils\n",
    "import robomimic.utils.tensor_utils as TensorUtils\n",
    "import robomimic.utils.obs_utils as ObsUtils\n",
    "import robomimic.utils.env_utils as EnvUtils\n",
    "import robomimic.utils.file_utils as FileUtils\n",
    "\n",
    "import imageio \n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from robomimic.envs.wrappers import EnvWrapper\n",
    "from copy import deepcopy\n",
    "import textwrap\n",
    "import numpy as np\n",
    "from collections import deque\n",
    "import time \n",
    "# from camera_utils import MyRealSense, CVCamera\n",
    "\n",
    "np.set_printoptions(precision=3, suppress=True)\n",
    "\n",
    "logger = get_deoxys_example_logger()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "device(type='cuda', index=0)"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "device = TorchUtils.get_torch_device(try_to_use_cuda=True)\n",
    "device "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class BaseRawObsDictGenerator:\n",
    "    def __init__(self, *args, **kwargs):\n",
    "        self.last_obs_dict = None\n",
    "\n",
    "    def get_raw_obs_dict(self, state_info):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "           state_info (dict): A dictionary of robot state + images\n",
    "        \"\"\"\n",
    "        obs_dict = {}\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def load(self):\n",
    "        raise NotImplementedError\n",
    "\n",
    "\n",
    "class RobotStateRawObsDictGenerator(BaseRawObsDictGenerator):\n",
    "    def __init__(self, *args, **kwargs):\n",
    "        super().__init__(*args, **kwargs)\n",
    "\n",
    "    def load(self, obs_dict, key, value, check_valid=True):\n",
    "        \"\"\"\n",
    "        This is to check if the data is correct or not. Sometimes the data will be all zero depending on the networking conditions.\n",
    "        \"\"\"\n",
    "        if (\n",
    "            (\n",
    "                np.sum(np.abs(value)) == 0.0\n",
    "                and key in [\"ee_states\", \"joint_states\", \"gripper_states\"]\n",
    "            )\n",
    "            and check_valid\n",
    "            and self.last_obs_dict is not None\n",
    "        ):\n",
    "            value = self.last_obs_dict[key]\n",
    "        obs_dict[key] = value\n",
    "\n",
    "    def get_raw_obs_dict(self, state_info):\n",
    "        last_state = state_info[\"last_state\"]\n",
    "        last_gripper_state = state_info[\"last_gripper_state\"]\n",
    "        obs_dict = {}\n",
    "\n",
    "        ee_states = np.array(last_state.O_T_EE)\n",
    "        joint_states = np.array(last_state.q)\n",
    "        gripper_states = np.array([last_gripper_state.width])\n",
    "\n",
    "        self.load(obs_dict, \"ee_states\", ee_states)\n",
    "        self.load(obs_dict, \"joint_states\", joint_states)\n",
    "        # Gripper widh will probably become zero\n",
    "        self.load(obs_dict, \"gripper_states\", gripper_states, check_valid=False)\n",
    "\n",
    "        for state in [\"ee_states\", \"joint_states\", \"gripper_states\"]:\n",
    "            if (\n",
    "                np.sum(np.abs(obs_dict[state])) <= 1e-6\n",
    "                and self.last_obs_dict is not None\n",
    "            ):\n",
    "                print(f\"{state} missing!!!!\")\n",
    "                obs_dict[state] = self.last_obs_dict[state]\n",
    "        self.last_obs_dict = obs_dict\n",
    "        return obs_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def stacked_get_init(init_obs, num_frames):\n",
    "    obs_history = {}\n",
    "    for k in init_obs:\n",
    "        obs_history[k] = deque([init_obs[k][None] for _ in range(num_frames)], maxlen=num_frames,)\n",
    "    obs = { k : np.concatenate(obs_history[k], axis=0) for k in obs_history }\n",
    "    return obs_history, obs \n",
    "\n",
    "def stacked_add_new(obs_history, new_obs):\n",
    "    for k in new_obs:\n",
    "        if 'timesteps' in k or 'actions' in k: continue\n",
    "        obs_history[k].append(new_obs[k][None])\n",
    "\n",
    "    obs= { k : np.concatenate(obs_history[k], axis=0) for k in obs_history }\n",
    "    return obs_history, obs "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# dataset_path=\"/home/franka_deoxys/data_franka/block_green_on_blue/franka_block_gonb_50.hdf5\"\n",
    "# ckpt_path = \"/home/franka_deoxys/data_franka/block_green_on_blue/policy_low/model_epoch_120.pth\"\n",
    "# assert os.path.exists(ckpt_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_path=\"/home/franka_deoxys/data_franka/block_green_on_blue/franka_gonb50.hdf5\"\n",
    "ckpt_path = \"/home/franka_deoxys/data_franka/block_green_on_blue/policy_low/trans_epoch_60.pth\"\n",
    "assert os.path.exists(ckpt_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/franka_deoxys/robomimic/robomimic/utils/file_utils.py:186: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
      "  ckpt_dict = torch.load(ckpt_path)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "============= Initialized Observation Utils with Obs Spec =============\n",
      "\n",
      "using obs modality: low_dim with keys: ['gripper_states', 'ee_states', 'joint_states']\n",
      "using obs modality: rgb with keys: []\n",
      "using obs modality: depth with keys: []\n",
      "using obs modality: scan with keys: []\n",
      "Created GPT_Backbone model with number of parameters: 18906112\n",
      "\n",
      "============= Initialized Observation Utils with Obs Spec =============\n",
      "\n",
      "using obs modality: low_dim with keys: ['gripper_states', 'ee_states', 'joint_states']\n",
      "using obs modality: rgb with keys: []\n",
      "using obs modality: depth with keys: []\n",
      "using obs modality: scan with keys: []\n",
      "SequenceDataset: loading dataset into memory...\n",
      "100%|██████████| 50/50 [00:00<00:00, 1353.21it/s]\n",
      "SequenceDataset: caching get_item calls...\n",
      "100%|██████████| 3280/3280 [00:00<00:00, 22225.86it/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "50"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# restore policy\n",
    "policy, ckpt_dict = FileUtils.policy_from_checkpoint(ckpt_path=ckpt_path, device=device, verbose=False)\n",
    "ckpt_dict = FileUtils.maybe_dict_from_checkpoint(ckpt_path=ckpt_path, ckpt_dict=None)\n",
    "# algo name and config from model dict\n",
    "algo_name, _ = FileUtils.algo_name_from_checkpoint(ckpt_dict=ckpt_dict)\n",
    "config, _ = FileUtils.config_from_checkpoint(algo_name=algo_name, ckpt_dict=ckpt_dict, verbose=False)\n",
    "\n",
    "config.unlock()\n",
    "config.train.data=dataset_path\n",
    "config.lock() \n",
    "\n",
    "\n",
    "ObsUtils.initialize_obs_utils_with_config(config)\n",
    "\n",
    "shape_meta = FileUtils.get_shape_metadata_from_dataset(\n",
    "    dataset_path=config.train.data,\n",
    "    all_obs_keys=config.all_obs_keys,\n",
    "    verbose=False\n",
    ")\n",
    "trainset, validset = TrainUtils.load_data_for_training(\n",
    "    config, obs_keys=shape_meta[\"all_obs_keys\"])\n",
    "len(trainset.demos)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = policy.policy\n",
    "model.set_eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[33m[Deoxys Examples WARNING] \u001b[0mThis is a very dummy control policy!!!!! (Ln 1)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_105536/2597759035.py:1: DeprecationWarning: The 'warn' method is deprecated, use 'warning' instead\n",
      "  logger.warn(\"This is a very dummy control policy!!!!!\")\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Opening SpaceMouse device\n",
      "Manufacturer: 3Dconnexion\n",
      "Product: SpaceMouse Wireless\n",
      "\n",
      "Control                       \tCommand\n",
      "Right button                  \treset simulation\n",
      "Left button (hold)            \tclose gripper\n",
      "Move mouse laterally          \tmove arm horizontally in x-y plane\n",
      "Move mouse vertically         \tmove arm vertically\n",
      "Twist mouse about an axis     \trotate arm about a corresponding axis\n",
      "ESC                           \tquit\n",
      "\n"
     ]
    }
   ],
   "source": [
    "logger.warn(\"This is a very dummy control policy!!!!!\")\n",
    "\n",
    "args = robot_config_parse_args()\n",
    "robot_interface = FrankaInterface(os.path.join(config_root, args.interface_cfg))\n",
    "controller_cfg = YamlConfig(\n",
    "    os.path.join(config_root, args.controller_cfg)\n",
    ").as_easydict()\n",
    "controller_type = args.controller_type\n",
    "\n",
    "spacemouse = SpaceMouse(vendor_id=9583, product_id=50734)\n",
    "spacemouse.start_control()\n",
    "\n",
    "raw_obs_dict_generator = RobotStateRawObsDictGenerator()\n",
    "\n",
    "import time\n",
    "\n",
    "time.sleep(0.3)\n",
    "dummy_torch_model = torch.nn.Linear(7, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Namespace(interface_cfg='charmander.yml', folder=PosixPath('/home/franka_deoxys/.local/share/jupyter/runtime/kernel-v2-12320ZTwlhMLLGbMN.json'), controller_type='OSC_POSE', controller_cfg='osc-position-controller.yml')"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "args "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(58, 7)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "file=trainset.hdf5_file\n",
    "demo = file['data']['demo_5']\n",
    "actions=demo['actions']\n",
    "actions.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "camera_ids [0 1]\n",
      "num_demos 50\n",
      "total 3280\n"
     ]
    }
   ],
   "source": [
    "# file[\"data\"].attrs[\"config\"]\n",
    "for key in file[\"data\"].attrs.keys():\n",
    "    print(key, file[\"data\"].attrs[key])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(66, 7)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "joint_sequence = file[\"data/demo_1/obs/joint_states\"]\n",
    "joint_sequence.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def set_gripper(open=True):\n",
    "    d=-1. if open else 1.0\n",
    "    action_close=np.array([ 0.,  0., -0.,  0.,  0., -0., d])\n",
    "    robot_interface.control(\n",
    "        controller_type=controller_type,\n",
    "        action=action_close,\n",
    "        controller_cfg=controller_cfg,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[30m[Deoxys Examples INFO] \u001b[0mResetting to the initial configuration (Ln 1)\n",
      "JOINT_POSITION\n"
     ]
    }
   ],
   "source": [
    "logger.info(\"Resetting to the initial configuration\")\n",
    "reset_joints_to(robot_interface, joint_sequence[0])\n",
    "set_gripper(open=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set_gripper(open=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "# for action in actions:\n",
    "#     robot_interface.control(\n",
    "#         controller_type=controller_type,\n",
    "#         action=action,\n",
    "#         controller_cfg=controller_cfg,\n",
    "#     )\n",
    "# robot_interface.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_current_obs_lowdim(robot_interface):\n",
    "    last_state = robot_interface._state_buffer[-1]\n",
    "    last_gripper_state = robot_interface._gripper_state_buffer[-1]\n",
    "\n",
    "    obs_dict = raw_obs_dict_generator.get_raw_obs_dict(\n",
    "    {\"last_state\": last_state, \"last_gripper_state\": last_gripper_state}) \n",
    "    return obs_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ee_states (1, 16)\n",
      "joint_states (1, 7)\n",
      "gripper_states (1, 1)\n"
     ]
    }
   ],
   "source": [
    "obs_dict=get_current_obs_lowdim(robot_interface)\n",
    "for key in obs_dict.keys():\n",
    "    obs_dict[key]=obs_dict[key][None]\n",
    "    print(key, obs_dict[key].shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_frames=10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 0.126, -0.02 , -0.306, -0.003, -0.005, -0.005, -1.   ],\n",
       "      dtype=float32)"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "init_obs = get_current_obs_lowdim(robot_interface)\n",
    "obs_history, obs = stacked_get_init(init_obs, num_frames) \n",
    "obs=TensorUtils.to_torch(obs, device=device)\n",
    "\n",
    "batch_obs={key:obs[key][None] for key in obs} \n",
    "pred = model.get_action(obs_dict=batch_obs , goal_dict=None)\n",
    "ac=pred.detach().cpu().numpy().squeeze()\n",
    "ac"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[22], line 16\u001b[0m\n\u001b[1;32m     14\u001b[0m obs\u001b[38;5;241m=\u001b[39mTensorUtils\u001b[38;5;241m.\u001b[39mto_torch(obs, device\u001b[38;5;241m=\u001b[39mdevice)\n\u001b[1;32m     15\u001b[0m batch_obs\u001b[38;5;241m=\u001b[39m{key:obs[key][\u001b[38;5;28;01mNone\u001b[39;00m] \u001b[38;5;28;01mfor\u001b[39;00m key \u001b[38;5;129;01min\u001b[39;00m obs} \n\u001b[0;32m---> 16\u001b[0m pred \u001b[38;5;241m=\u001b[39m \u001b[43mmodel\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mget_action\u001b[49m\u001b[43m(\u001b[49m\u001b[43mobs_dict\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mbatch_obs\u001b[49m\u001b[43m \u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mgoal_dict\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m)\u001b[49m\n\u001b[1;32m     17\u001b[0m action\u001b[38;5;241m=\u001b[39mpred\u001b[38;5;241m.\u001b[39mdetach()\u001b[38;5;241m.\u001b[39mcpu()\u001b[38;5;241m.\u001b[39mnumpy()\u001b[38;5;241m.\u001b[39msqueeze() \n\u001b[1;32m     18\u001b[0m \u001b[38;5;66;03m# action=action[:6] \u001b[39;00m\n",
      "File \u001b[0;32m~/robomimic/robomimic/algo/bc.py:769\u001b[0m, in \u001b[0;36mBC_Transformer.get_action\u001b[0;34m(self, obs_dict, goal_dict)\u001b[0m\n\u001b[1;32m    759\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    760\u001b[0m \u001b[38;5;124;03mGet policy action outputs.\u001b[39;00m\n\u001b[1;32m    761\u001b[0m \u001b[38;5;124;03mArgs:\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    765\u001b[0m \u001b[38;5;124;03m    action (torch.Tensor): action tensor\u001b[39;00m\n\u001b[1;32m    766\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    767\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mnets\u001b[38;5;241m.\u001b[39mtraining\n\u001b[0;32m--> 769\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mnets\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mpolicy\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m(\u001b[49m\u001b[43mobs_dict\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mactions\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mgoal_dict\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mgoal_dict\u001b[49m\u001b[43m)\u001b[49m[:, \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m, :]\n",
      "File \u001b[0;32m~/.local/lib/python3.10/site-packages/torch/nn/modules/module.py:1553\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1551\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m   1552\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1553\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/.local/lib/python3.10/site-packages/torch/nn/modules/module.py:1562\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1557\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1558\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1559\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1560\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1561\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1562\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1564\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m   1565\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
      "File \u001b[0;32m~/robomimic/robomimic/models/policy_nets.py:1326\u001b[0m, in \u001b[0;36mTransformerGMMActorNetwork.forward\u001b[0;34m(self, obs_dict, actions, goal_dict)\u001b[0m\n\u001b[1;32m   1316\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mforward\u001b[39m(\u001b[38;5;28mself\u001b[39m, obs_dict, actions\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m, goal_dict\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m):\n\u001b[1;32m   1317\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m   1318\u001b[0m \u001b[38;5;124;03m    Samples actions from the policy distribution.\u001b[39;00m\n\u001b[1;32m   1319\u001b[0m \u001b[38;5;124;03m    Args:\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m   1324\u001b[0m \u001b[38;5;124;03m        action (torch.Tensor): batch of actions from policy distribution\u001b[39;00m\n\u001b[1;32m   1325\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[0;32m-> 1326\u001b[0m     out \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mforward_train\u001b[49m\u001b[43m(\u001b[49m\u001b[43mobs_dict\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mobs_dict\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mactions\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mactions\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mgoal_dict\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mgoal_dict\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1327\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m out\u001b[38;5;241m.\u001b[39msample()\n",
      "File \u001b[0;32m~/robomimic/robomimic/models/policy_nets.py:1278\u001b[0m, in \u001b[0;36mTransformerGMMActorNetwork.forward_train\u001b[0;34m(self, obs_dict, actions, goal_dict, low_noise_eval)\u001b[0m\n\u001b[1;32m   1274\u001b[0m     goal_dict \u001b[38;5;241m=\u001b[39m TensorUtils\u001b[38;5;241m.\u001b[39munsqueeze_expand_at(goal_dict, size\u001b[38;5;241m=\u001b[39mobs_dict[mod]\u001b[38;5;241m.\u001b[39mshape[\u001b[38;5;241m1\u001b[39m], dim\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1\u001b[39m)\n\u001b[1;32m   1276\u001b[0m forward_kwargs \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mdict\u001b[39m(obs\u001b[38;5;241m=\u001b[39mobs_dict, goal\u001b[38;5;241m=\u001b[39mgoal_dict)\n\u001b[0;32m-> 1278\u001b[0m outputs \u001b[38;5;241m=\u001b[39m \u001b[43mMIMO_Transformer\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mforward\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mforward_kwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1280\u001b[0m means \u001b[38;5;241m=\u001b[39m outputs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmean\u001b[39m\u001b[38;5;124m\"\u001b[39m]\n\u001b[1;32m   1281\u001b[0m scales \u001b[38;5;241m=\u001b[39m outputs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mscale\u001b[39m\u001b[38;5;124m\"\u001b[39m]\n",
      "File \u001b[0;32m~/robomimic/robomimic/models/obs_nets.py:1072\u001b[0m, in \u001b[0;36mMIMO_Transformer.forward\u001b[0;34m(self, **inputs)\u001b[0m\n\u001b[1;32m   1070\u001b[0m     transformer_embeddings \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39minput_embedding(transformer_inputs)\n\u001b[1;32m   1071\u001b[0m     \u001b[38;5;66;03m# pass encoded sequences through transformer\u001b[39;00m\n\u001b[0;32m-> 1072\u001b[0m     transformer_encoder_outputs \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mnets\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mtransformer\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m]\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mforward\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtransformer_embeddings\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1074\u001b[0m transformer_outputs \u001b[38;5;241m=\u001b[39m transformer_encoder_outputs\n\u001b[1;32m   1075\u001b[0m \u001b[38;5;66;03m# apply decoder to each timestep of sequence to get a dictionary of outputs\u001b[39;00m\n",
      "File \u001b[0;32m~/robomimic/robomimic/models/transformers.py:424\u001b[0m, in \u001b[0;36mGPT_Backbone.forward\u001b[0;34m(self, inputs)\u001b[0m\n\u001b[1;32m    422\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mforward\u001b[39m(\u001b[38;5;28mself\u001b[39m, inputs):\n\u001b[1;32m    423\u001b[0m     \u001b[38;5;28;01massert\u001b[39;00m inputs\u001b[38;5;241m.\u001b[39mshape[\u001b[38;5;241m1\u001b[39m:] \u001b[38;5;241m==\u001b[39m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mcontext_length, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39membed_dim), inputs\u001b[38;5;241m.\u001b[39mshape\n\u001b[0;32m--> 424\u001b[0m     x \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mnets\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mtransformer\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m(\u001b[49m\u001b[43minputs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    425\u001b[0m     transformer_output \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mnets[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124moutput_ln\u001b[39m\u001b[38;5;124m\"\u001b[39m](x)\n\u001b[1;32m    426\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m transformer_output\n",
      "File \u001b[0;32m~/.local/lib/python3.10/site-packages/torch/nn/modules/module.py:1553\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1551\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m   1552\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1553\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/.local/lib/python3.10/site-packages/torch/nn/modules/module.py:1562\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1557\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1558\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1559\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1560\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1561\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1562\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1564\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m   1565\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
      "File \u001b[0;32m~/.local/lib/python3.10/site-packages/torch/nn/modules/container.py:219\u001b[0m, in \u001b[0;36mSequential.forward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    217\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mforward\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;28minput\u001b[39m):\n\u001b[1;32m    218\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m module \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m:\n\u001b[0;32m--> 219\u001b[0m         \u001b[38;5;28minput\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[43mmodule\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m    220\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28minput\u001b[39m\n",
      "File \u001b[0;32m~/.local/lib/python3.10/site-packages/torch/nn/modules/module.py:1553\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1551\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m   1552\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1553\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/.local/lib/python3.10/site-packages/torch/nn/modules/module.py:1562\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1557\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1558\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1559\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1560\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1561\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1562\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1564\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m   1565\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
      "File \u001b[0;32m~/robomimic/robomimic/models/transformers.py:292\u001b[0m, in \u001b[0;36mSelfAttentionBlock.forward\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m    288\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    289\u001b[0m \u001b[38;5;124;03mForward pass - chain self-attention + MLP blocks, with residual connections and layer norms.\u001b[39;00m\n\u001b[1;32m    290\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    291\u001b[0m x \u001b[38;5;241m=\u001b[39m x \u001b[38;5;241m+\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mnets[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mattention\u001b[39m\u001b[38;5;124m\"\u001b[39m](\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mnets[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mln1\u001b[39m\u001b[38;5;124m\"\u001b[39m](x))\n\u001b[0;32m--> 292\u001b[0m x \u001b[38;5;241m=\u001b[39m x \u001b[38;5;241m+\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mnets[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmlp\u001b[39m\u001b[38;5;124m\"\u001b[39m](\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mnets\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mln2\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m)\u001b[49m)\n\u001b[1;32m    293\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m x\n",
      "File \u001b[0;32m~/.local/lib/python3.10/site-packages/torch/nn/modules/module.py:1553\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1551\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m   1552\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1553\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/.local/lib/python3.10/site-packages/torch/nn/modules/module.py:1562\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1557\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1558\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1559\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1560\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1561\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1562\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1564\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m   1565\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
      "File \u001b[0;32m~/.local/lib/python3.10/site-packages/torch/nn/modules/normalization.py:202\u001b[0m, in \u001b[0;36mLayerNorm.forward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    201\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mforward\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;28minput\u001b[39m: Tensor) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Tensor:\n\u001b[0;32m--> 202\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mF\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlayer_norm\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    203\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mnormalized_shape\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mweight\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbias\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43meps\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/.local/lib/python3.10/site-packages/torch/nn/functional.py:2576\u001b[0m, in \u001b[0;36mlayer_norm\u001b[0;34m(input, normalized_shape, weight, bias, eps)\u001b[0m\n\u001b[1;32m   2572\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m has_torch_function_variadic(\u001b[38;5;28minput\u001b[39m, weight, bias):\n\u001b[1;32m   2573\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m handle_torch_function(\n\u001b[1;32m   2574\u001b[0m         layer_norm, (\u001b[38;5;28minput\u001b[39m, weight, bias), \u001b[38;5;28minput\u001b[39m, normalized_shape, weight\u001b[38;5;241m=\u001b[39mweight, bias\u001b[38;5;241m=\u001b[39mbias, eps\u001b[38;5;241m=\u001b[39meps\n\u001b[1;32m   2575\u001b[0m     )\n\u001b[0;32m-> 2576\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlayer_norm\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mnormalized_shape\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mweight\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mbias\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43meps\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbackends\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcudnn\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43menabled\u001b[49m\u001b[43m)\u001b[49m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "actions = []\n",
    "for _ in range(200):\n",
    "    action, grasp = input2action(\n",
    "        device=spacemouse,\n",
    "        controller_type=controller_type,\n",
    "    )\n",
    "    if action is None:\n",
    "        break\n",
    "    if len(robot_interface._state_buffer) == 0:\n",
    "        continue\n",
    "\n",
    "    obs_t = get_current_obs_lowdim(robot_interface)\n",
    "    obs_history, obs = stacked_add_new(obs_history, obs_t) \n",
    "    obs=TensorUtils.to_torch(obs, device=device)\n",
    "    batch_obs={key:obs[key][None] for key in obs} \n",
    "    pred = model.get_action(obs_dict=batch_obs , goal_dict=None)\n",
    "    action=pred.detach().cpu().numpy().squeeze() \n",
    "    # action=action[:6] \n",
    "\n",
    "    actions.append(action)\n",
    "    robot_interface.control(\n",
    "        controller_type=controller_type,\n",
    "        action=action,\n",
    "        controller_cfg=controller_cfg,\n",
    "    )\n",
    "\n",
    "    dt=0.01\n",
    "    time.sleep(dt)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Worked great."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "robot_interface.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "56"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(actions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 0.126 -0.02  -0.306 -0.003 -0.005 -0.005 -1.   ]\n",
      "[ 0.126 -0.021 -0.306 -0.003 -0.006 -0.005 -1.   ]\n",
      "[ 0.127 -0.02  -0.306 -0.003 -0.005 -0.005 -1.   ]\n",
      "[ 0.136 -0.02  -0.31  -0.003 -0.006 -0.005 -1.   ]\n",
      "[ 0.149 -0.019 -0.314 -0.003 -0.006 -0.005 -1.   ]\n",
      "[ 0.163 -0.018 -0.319 -0.003 -0.006 -0.005 -1.   ]\n",
      "[ 0.181 -0.018 -0.326 -0.003 -0.006 -0.005 -1.   ]\n",
      "[ 0.2   -0.017 -0.333 -0.003 -0.006 -0.005 -1.   ]\n",
      "[ 0.222 -0.016 -0.342 -0.003 -0.006 -0.005 -1.   ]\n",
      "[ 0.241 -0.015 -0.349 -0.003 -0.006 -0.005 -1.   ]\n"
     ]
    }
   ],
   "source": [
    "actions[:10]\n",
    "for ac in actions[:10]:\n",
    "    print(ac)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}