NonScalarized Lagrange SAC_V2.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "importing Jupyter notebook from Helper_class2.ipynb\n"
     ]
    }
   ],
   "source": [
    "#Path related stuff\n",
    "import os\n",
    "\n",
    "#Math stuff\n",
    "import numpy as np\n",
    "import random\n",
    "\n",
    "#Environment stuff\n",
    "import gym\n",
    "import mujoco_py\n",
    "\n",
    "#3D Object stuff\n",
    "import pyvista as pv\n",
    "\n",
    "#Pytorch Stuff \n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.distributions import Normal\n",
    "import torch.nn.functional as F\n",
    "from torch.optim import Adam\n",
    "from torch.optim import NAdam\n",
    "\n",
    "#To import other notebooks\n",
    "import import_ipynb\n",
    "\n",
    "#Matplotlib stuff\n",
    "#import matplotlib.pyplot as plt\n",
    "#Helper class imports\n",
    "from Helper_class2 import ThreeD_Obstacle\n",
    "from Helper_class2 import Features\n",
    "from Helper_class2 import AgentUtils\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ReplayBuffer:\n",
    "  def __init__(self, buffer_size, batch_size, seed = 4):\n",
    "    self.buffer = []\n",
    "    self.max_size = buffer_size\n",
    "    self.batch_size = batch_size\n",
    "    self.random_generator = np.random.RandomState(seed)\n",
    "\n",
    "  def append(self, state, action, reward, terminal, next_state, cost):\n",
    "    if(len(self.buffer) == self.max_size):\n",
    "      del self.buffer[0]\n",
    "\n",
    "    self.buffer.append([state, action, reward, terminal, next_state, cost])\n",
    "\n",
    "  def sample(self):\n",
    "    batch = random.sample(self.buffer, self.batch_size)\n",
    "\n",
    "    state, action, reward, terminal, next_state, cost = map(np.stack, zip(*batch))\n",
    "    \n",
    "    return state, action, reward, terminal, next_state, cost\n",
    "\n",
    "  def get_buffer_size(self):\n",
    "    return len(self.buffer)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "LOG_SIG_MAX = 2\n",
    "LOG_SIG_MIN = -20\n",
    "epsilon = 1e-6\n",
    "\n",
    "#Initialize Policy weights\n",
    "def weights_init_(m):\n",
    "    if isinstance(m, nn.Linear):\n",
    "        torch.nn.init.xavier_uniform_(m.weight, gain=1)\n",
    "        torch.nn.init.constant_(m.bias, 0)\n",
    "\n",
    "class QNetwork(nn.Module):\n",
    "    def __init__(self, num_inputs, num_actions, hidden_dim):\n",
    "        super(QNetwork, self).__init__()\n",
    "        \n",
    "        # Q1 architecture\n",
    "        self.linear1 = nn.Linear(num_inputs + num_actions, hidden_dim)\n",
    "        self.linear2 = nn.Linear(hidden_dim, hidden_dim)\n",
    "        self.linear3 = nn.Linear(hidden_dim, hidden_dim)\n",
    "        self.linear4 = nn.Linear(hidden_dim, 1)\n",
    "\n",
    "        # Q2 architecture\n",
    "        self.linear5 = nn.Linear(num_inputs + num_actions, hidden_dim)\n",
    "        self.linear6 = nn.Linear(hidden_dim, hidden_dim)\n",
    "        self.linear7 = nn.Linear(hidden_dim, hidden_dim)\n",
    "        self.linear8 = nn.Linear(hidden_dim, 1)\n",
    "\n",
    "        self.apply(weights_init_)\n",
    "\n",
    "    def forward(self, state, action):\n",
    "        xu = torch.cat([state, action], 1)\n",
    "        \n",
    "        x1 = F.relu(self.linear1(xu))\n",
    "        x1 = F.relu(self.linear2(x1))\n",
    "        x1 = F.relu(self.linear3(x1))\n",
    "        \n",
    "        x1 = self.linear4(x1)\n",
    "\n",
    "        x2 = F.relu(self.linear5(xu))\n",
    "        x2 = F.relu(self.linear6(x2))\n",
    "        x2 = F.relu(self.linear7(x2))\n",
    "        x2 = self.linear8(x2)\n",
    "\n",
    "        return x1, x2\n",
    "\n",
    "\n",
    "class GaussianPolicy(nn.Module):\n",
    "    def __init__(self, num_inputs, num_actions, hidden_dim, action_range=None):\n",
    "        super(GaussianPolicy, self).__init__()\n",
    "        \n",
    "        self.linear1 = nn.Linear(num_inputs, hidden_dim)\n",
    "        self.linear2 = nn.Linear(hidden_dim, hidden_dim)\n",
    "        self.linear3 = nn.Linear(hidden_dim, hidden_dim)\n",
    "\n",
    "        self.mean_linear = nn.Linear(hidden_dim, num_actions)\n",
    "        self.log_std_linear = nn.Linear(hidden_dim, num_actions)\n",
    "\n",
    "        self.apply(weights_init_)\n",
    "\n",
    "        # action rescaling\n",
    "        if action_range is None:\n",
    "            self.action_scale = torch.tensor(1.)\n",
    "            self.action_bias = torch.tensor(0.)\n",
    "        else:\n",
    "            self.action_scale = torch.FloatTensor([\n",
    "                (np.max(action_range) - np.min(action_range)) / 2.\n",
    "                ])\n",
    "            self.action_bias = torch.FloatTensor([\n",
    "                (np.max(action_range) - np.max(action_range)) / 2.\n",
    "                ])\n",
    "\n",
    "    def forward(self, state):\n",
    "        x = F.relu(self.linear1(state))\n",
    "        x = F.relu(self.linear2(x))\n",
    "        x = F.relu(self.linear3(x))\n",
    "\n",
    "        mean = self.mean_linear(x)\n",
    "        log_std = self.log_std_linear(x)\n",
    "        log_std = torch.clamp(log_std, min=LOG_SIG_MIN, max=LOG_SIG_MAX)\n",
    "        return mean, log_std\n",
    "\n",
    "    def sample(self, state):\n",
    "        mean, log_std = self.forward(state)\n",
    "        std = log_std.exp()\n",
    "        normal = Normal(mean, std)\n",
    "        x_t = normal.rsample()  # for reparameterization trick (mean + std * N(0,1))\n",
    "        y_t = torch.tanh(x_t)\n",
    "        action = y_t * self.action_scale + self.action_bias\n",
    "        log_prob = normal.log_prob(x_t)\n",
    "        # Enforcing Action Bound\n",
    "        log_prob -= torch.log(self.action_scale * (1 - y_t.pow(2)) + epsilon)\n",
    "        log_prob = log_prob.sum(1, keepdim=True)\n",
    "        \n",
    "        mean = torch.tanh(mean) * self.action_scale + self.action_bias\n",
    "        return action, log_prob, mean\n",
    "\n",
    "    def to(self, device):\n",
    "        self.action_scale = self.action_scale.to(device)\n",
    "        self.action_bias = self.action_bias.to(device)\n",
    "        return super(GaussianPolicy, self).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SAC_Agent(object):\n",
    "    def __init__(self,observation_space, action_space, gamma, lr,  tau, alpha, hidden_network_size, \n",
    "                 action_range, batch_size, agent_util):\n",
    "        self.agent_util = agent_util\n",
    "\n",
    "        self.REPLAY_BATCH_SIZE = batch_size\n",
    "        self.replay_buffer = ReplayBuffer(300000, self.REPLAY_BATCH_SIZE)\n",
    "\n",
    "        self.gamma = gamma\n",
    "        self.lr = lr\n",
    "        self.tau = tau\n",
    "        self.alpha = alpha\n",
    "\n",
    "        self.device = \"cuda:0\" if torch.cuda.is_available() else \"cpu\"\n",
    "        \n",
    "        self.LEARNING_RATE_LAM = 0.001\n",
    "        self.lamb = torch.tensor(-10.0, requires_grad=True, device=self.device)\n",
    "        self.lamb_optimiser = torch.optim.Adam([self.lamb], lr=self.LEARNING_RATE_LAM, weight_decay=1e-5)\n",
    "\n",
    "        self.cost_lim = 0.001\n",
    "\n",
    "        self.iterator = 1\n",
    "\n",
    "        self.critic = QNetwork(observation_space[0], action_space[0], hidden_network_size).to(device=self.device)\n",
    "        self.critic_optim = Adam(self.critic.parameters(), lr=self.lr)\n",
    "\n",
    "        self.critic_target = QNetwork(observation_space[0], action_space[0], hidden_network_size).to(self.device)\n",
    "        \n",
    "\n",
    "        self.critic_cost = QNetwork(observation_space[0], action_space[0], hidden_network_size).to(device=self.device)\n",
    "        self.critic_cost_optim = Adam(self.critic_cost.parameters(), lr=self.lr)\n",
    "        \n",
    "        self.critic_target_cost = QNetwork(observation_space[0], action_space[0], hidden_network_size).to(self.device)\n",
    "        \n",
    "        #perfoming a soft copy of the parameters\n",
    "        for target_param, param in zip( self.critic_target.parameters(), self.critic.parameters()):\n",
    "            target_param.data.copy_( param.data * self.tau + (1-self.tau) * target_param.data)\n",
    "\n",
    "        for target_param, param in zip( self.critic_target_cost.parameters(), self.critic_cost.parameters()):\n",
    "            target_param.data.copy_( param.data * self.tau + (1-self.tau) * target_param.data)\n",
    "\n",
    "\n",
    "        self.policy = GaussianPolicy(observation_space[0], action_space[0], hidden_network_size, action_range).to(self.device)\n",
    "        self.policy_optim = Adam(self.policy.parameters(), lr=self.lr)\n",
    "\n",
    "\n",
    "        # self.target_entropy = -torch.prod(torch.Tensor(action_space).to(self.device)).item()\n",
    "        # self.log_alpha = torch.zeros(1, requires_grad=True, device=self.device)\n",
    "        # self.alpha_optim = NAdam([self.log_alpha], lr=self.lr)\n",
    "\n",
    "\n",
    "    def select_action(self, state, evaluate=False):\n",
    "        state = torch.FloatTensor(state).to(self.device).unsqueeze(0)\n",
    "        if evaluate is False:\n",
    "            action, _, _ = self.policy.sample(state)\n",
    "\n",
    "        return action.detach().cpu().numpy()[0]\n",
    "    \n",
    "    \n",
    "    def gradient_clipping(self, model_parameters, clip_val):\n",
    "        torch.nn.utils.clip_grad_norm_(model_parameters, max_norm=clip_val, norm_type=2)\n",
    "\n",
    "    \n",
    "    def update_parameters(self):\n",
    "        # Sample a batch from memory\n",
    "\n",
    "        state_batch, action_batch, reward_batch, mask_batch, next_state_batch, cost_batch = self.replay_buffer.sample()\n",
    "\n",
    "        state_batch = torch.FloatTensor(state_batch).to(self.device)\n",
    "        next_state_batch = torch.FloatTensor(next_state_batch).to(self.device)\n",
    "        action_batch = torch.FloatTensor(action_batch).to(self.device)\n",
    "        reward_batch = torch.FloatTensor(reward_batch).to(self.device).unsqueeze(1)\n",
    "        mask_batch = torch.FloatTensor(mask_batch).to(self.device).unsqueeze(1)\n",
    "        cost_batch = torch.FloatTensor(cost_batch).to(self.device).unsqueeze(1)\n",
    "\n",
    "        with torch.no_grad():\n",
    "            next_state_action, next_state_log_pi, _ = self.policy.sample(next_state_batch)\n",
    "            \n",
    "            qf1_next_target, qf2_next_target = self.critic_target(next_state_batch, next_state_action)\n",
    "            min_qf_next_target = torch.min(qf1_next_target, qf2_next_target) - self.alpha * next_state_log_pi\n",
    "            next_q_value = (reward_batch) + (1-mask_batch) * self.gamma * (min_qf_next_target)\n",
    "\n",
    "            qf1_next_target_cost, _ = self.critic_target_cost(next_state_batch, next_state_action)\n",
    "            #min_qf_next_target_cost = torch.min(qf1_next_target_cost, qf2_next_target_cost)\n",
    "            next_q_value_cost = cost_batch + (1-mask_batch) * self.gamma * (qf1_next_target_cost)\n",
    "       \n",
    "\n",
    "        qf1, qf2 = self.critic(state_batch, action_batch) # Two Q-functions to mitigate positive bias in the policy improvement step\n",
    "        qf1_c, _ = self.critic_cost(state_batch, action_batch) # Two Q-functions to mitigate positive bias in the policy improvement step\n",
    "\n",
    "        qf1_loss = F.mse_loss(qf1, next_q_value)\n",
    "        qf2_loss = F.mse_loss(qf2, next_q_value)\n",
    "\n",
    "        qf1_cost_loss = F.mse_loss(qf1_c, next_q_value_cost)  # JQ = 𝔼(st,at)~D[0.5(Q1(st,at) - r(st,at) - γ(𝔼st+1~p[V(st+1)]))^2]\n",
    "\n",
    "\n",
    "        qf_loss = (qf1_loss + qf2_loss).mean()\n",
    "        qf_cost_loss = qf1_cost_loss.mean()\n",
    "\n",
    "\n",
    "        self.critic_optim.zero_grad()\n",
    "        qf_loss.backward()\n",
    "        self.critic_optim.step()\n",
    "\n",
    "        self.critic_cost_optim.zero_grad()\n",
    "        qf_cost_loss.backward()\n",
    "        self.critic_cost_optim.step()\n",
    "\n",
    "\n",
    "        for target_param, param in zip( self.critic_target.parameters(), self.critic.parameters()):\n",
    "            target_param.data.copy_( param.data * self.tau + (1-self.tau) * target_param.data)\n",
    "        \n",
    "        for target_param, param in zip( self.critic_target_cost.parameters(), self.critic_cost.parameters()):\n",
    "            target_param.data.copy_( param.data * self.tau + (1-self.tau) * target_param.data)\n",
    "\n",
    "\n",
    "        pi, log_pi, _ = self.policy.sample(state_batch)\n",
    "        \n",
    "\n",
    "        qf1_pi, qf2_pi = self.critic(state_batch, pi)\n",
    "        min_qf_pi = torch.min(qf1_pi, qf2_pi)\n",
    "        first_term = (self.alpha * log_pi) - min_qf_pi\n",
    "\n",
    "        qf1_pi_c, _ = self.critic_cost(state_batch, pi)\n",
    "        \n",
    "        second_term = 10 * self.lamb * qf1_pi_c\n",
    "\n",
    "        policy_loss = (first_term + second_term).mean()# Jπ = 𝔼st∼D,εt∼N[α * logπ(f(εt;st)|st) − Q(st,f(εt;st))]\n",
    "\n",
    "\n",
    "        self.policy_optim.zero_grad()\n",
    "        policy_loss.backward()\n",
    "        self.policy_optim.step()\n",
    "\n",
    "\n",
    "        #Training lambda\n",
    "        #Training lambda\n",
    "        self.iterator += 1\n",
    "        \n",
    "        if self.iterator % 3 == 0:\n",
    "\n",
    "            violation = cost_batch - self.cost_lim\n",
    "            \n",
    "            self.log_lam = torch.nn.functional.softplus(self.lamb, beta=0.5)\n",
    "            lambda_loss = -self.log_lam*(violation).detach()\n",
    "            lambda_loss = lambda_loss.sum()\n",
    "\n",
    "            self.lamb_optimiser.zero_grad()\n",
    "            lambda_loss.backward()\n",
    "            self.lamb_optimiser.step()\n",
    "\n",
    "            self.iterator = 1\n",
    "\n",
    "        \n",
    "        return qf1_loss.item(), qf2_loss.item(), policy_loss.item()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found 2 GPUs for rendering. Using device 0.\n"
     ]
    }
   ],
   "source": [
    "mj_path = mujoco_py.utils.discover_mujoco()\n",
    "xml_path = os.path.join(mj_path, 'model', 'ur5.xml')\n",
    "model = mujoco_py.load_model_from_path(xml_path)\n",
    "sim = mujoco_py.MjSim(model)\n",
    "viewer = mujoco_py.MjRenderContextOffscreen(sim, -1)\n",
    "\n",
    "##agent.critic.state_dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def models_performance_tracker(seed, network_size, q1_avg_loss, q2_avg_loss, policy_avg_loss, avg_reward, steps, crashes, total_violations, PATH):\n",
    "    torch.save({\n",
    "            'seed': seed,\n",
    "            'network_size': network_size,\n",
    "            'q1_avg_loss': q1_avg_loss,\n",
    "            'q2_avg_loss': q2_avg_loss,\n",
    "            'policy_avg_loss': policy_avg_loss,\n",
    "            'avg_reward': avg_reward,\n",
    "            'steps': steps,\n",
    "            'crashes': crashes,\n",
    "            'violations': total_violations,\n",
    "            #'model_state_dict': agent.critic.state_dict(),\n",
    "            #'optimizer_state_dict': optimizer.state_dict(),\n",
    "            }, PATH)\n",
    "    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Simulation is unstable. Taking corrective action...\n",
      "crashed\n",
      "epoch:  0 ,  step :  1211 ,  avg_reward :  -0.9138151455504\n",
      "intersected\n",
      "epoch:  1 ,  step :  861 ,  avg_reward :  -0.873299730538498\n",
      "crashed\n",
      "epoch:  2 ,  step :  348 ,  avg_reward :  -0.7333814807492082\n"
     ]
    }
   ],
   "source": [
    "EPOCHS = 10000\n",
    "TARGET_POS = [-0.11034687, -0.59504954, 0.2978785 ] #based on [0, -0.22, 1.67, 0.66, 1.57, 0.22]\n",
    "\n",
    "SEEDs = [0]\n",
    "NETWORKs_SEARCH = [128] #This was 128 with seed 4,5 when it was working\n",
    "\n",
    "#Object stuff\n",
    "object_position1 = [-0.11034687, -0.59504954, 0.2978785 + 0.2] #the plus is the axis radius, to be on the same plne\n",
    "\n",
    "# layer1_shere_semi_axis = [0.17,0.17,0.17]\n",
    "# layer2_shere_semi_axis = [0.35,0.35,0.35]\n",
    "l1_radius = 0.35\n",
    "l2_radius = 0.34\n",
    "\n",
    "threeD_object1 = ThreeD_Obstacle(simulation = sim, sphere_or_ellipsoid= \"Sphere\", \n",
    "                                 radius_or_axis = 0.2 , my_position = object_position1,\n",
    "                                 safety_l1 = l1_radius,\n",
    "                                 safety_l2 = l2_radius,\n",
    "                                 target_pos = TARGET_POS,\n",
    "                                   translation_rate= [0.008, 0, 0])\n",
    "\n",
    "initial_action = [-2.42, -1.57, -0.597, -1.48, 3.14, -1.42]\n",
    "\n",
    "#Arm stuff\n",
    "arm_utils = AgentUtils(sim, threeD_object1, initial_action, TARGET_POS) \n",
    "arm_feature_object = Features(sim, threeD_object1, arm_utils)\n",
    "\n",
    "for seed in SEEDs:\n",
    "    alpha = 0.05\n",
    "    random.seed(seed)\n",
    "    np.random.seed(seed)\n",
    "    torch.manual_seed(seed)\n",
    "\n",
    "    for network in NETWORKs_SEARCH:\n",
    "\n",
    "        agent = SAC_Agent((48,), (6,), 0.99, 0.0001, 0.005, alpha, network, \n",
    "                          [-0.0698132, 0.0698132], 256, arm_utils)\n",
    "        \n",
    "        q1_loss_avgs = np.array([])\n",
    "        q2_loss_avgs = np.array([])\n",
    "        policy_loss_avgs = np.array([])\n",
    "        average_rewards = np.array([])\n",
    "        steps = np.array([])\n",
    "        crashes = np.array([])\n",
    "        total_violations = np.array([])\n",
    "        lambda_episodes = []\n",
    "        for i in range(EPOCHS):\n",
    "            done = 0\n",
    "            distance = 0\n",
    "            step = 1\n",
    "            minus_on_reward = 0\n",
    "            \n",
    "            #storing data\n",
    "            q1_loss = np.array([])\n",
    "            q2_loss = np.array([])\n",
    "            policy_loss = np.array([])\n",
    "            rewards = np.array([])\n",
    "            violations = np.array([])\n",
    "            crashed = 0\n",
    "\n",
    "            #taking the initial step\n",
    "            sim.data.qpos[:] = initial_action\n",
    "            sim.step()\n",
    "\n",
    "            #Getting features of the arm robot.\n",
    "            calculated_features = arm_feature_object.get_all_features().copy()\n",
    "            observation = calculated_features\n",
    "            \n",
    "            lambda_step = []\n",
    "            \n",
    "            time_step_T = 5000\n",
    "            for _ in range(time_step_T):\n",
    "                \n",
    "                action = agent.select_action(observation).copy()\n",
    "                \n",
    "                # current_qpos = sim.data.qpos.copy()\n",
    "                # new_qpos = current_qpos + action\n",
    "                sim.data.qpos[:] = arm_utils.get_joint_angles(action)\n",
    "\n",
    "\n",
    "                try:\n",
    "                    sim.step()\n",
    "                except mujoco_py.MujocoException as e:\n",
    "                    if \"Got MuJoCo Warning: Nan, Inf or huge value in QACC\" in str(e):\n",
    "                        print(\"Simulation is unstable. Taking corrective action...\")\n",
    "                        # Do something to correct the instability, e.g. reset the simulation or adjust the control parameters\n",
    "                        sim.data.qpos[:] = [-2.42, -1.57, -0.597, -1.48, 3.14, -1.41]\n",
    "                        sim.step() #required to update the joints and EE\n",
    "                    else:\n",
    "                        # Handle other MujocoExceptions as needed\n",
    "                        print(f\"Unexpected MujocoException: {e}\")\n",
    "\n",
    "              \n",
    "                cost = threeD_object1.compute_link_distance_to_object_when_in_l1_by_closest_points_for_cost()\n",
    "                violations = np.append(violations,cost)\n",
    "                \n",
    "                if( threeD_object1.object_intruded().sum() >= 1):\n",
    "                    done = 1\n",
    "                    crashed = 1\n",
    "                    #max_distance_cost = np.abs(rewards).sum()\n",
    "                    #threeD_object1.reset()\n",
    "\n",
    "                    print(\"crashed\")\n",
    "\n",
    "                reward = arm_utils.get_reward()\n",
    "\n",
    "                #translate the object\n",
    "                threeD_object1.linearly_translate_object_back_and_forward()\n",
    "\n",
    "                #Getting features of the arm robot.\n",
    "                calculated_features = arm_feature_object.get_all_features().copy()\n",
    "                new_observation = calculated_features\n",
    "\n",
    "                if(arm_utils.check_EE_is_in_the_target_area() == 1):#function returns a 1 if the EE coordinates are in the sphere, else 0\n",
    "                    done = 1\n",
    "\n",
    "                rewards = np.append(rewards,reward)\n",
    "                \n",
    "\n",
    "                agent.replay_buffer.append(observation, action, reward, done, new_observation, cost)\n",
    "                \n",
    "\n",
    "                observation = new_observation\n",
    "\n",
    "                if( agent.replay_buffer.get_buffer_size() > agent.REPLAY_BATCH_SIZE ):\n",
    "                    q1_l, q2_l, policy_l = agent.update_parameters()\n",
    "\n",
    "                    q1_loss = np.append(q1_loss, q1_l)\n",
    "                    q2_loss = np.append(q2_loss, q2_l)\n",
    "                    policy_loss = np.append(policy_loss, policy_l)\n",
    "                   \n",
    "\n",
    "                if(done == 1):\n",
    "                    arm_utils = AgentUtils(sim, threeD_object1, initial_action, TARGET_POS) #resetting the object for the initial position\n",
    "                    break\n",
    "                \n",
    "                step+=1\n",
    "\n",
    "            print(\"epoch: \", i, \",  step : \",step, \",  avg_reward : \", rewards.mean())\n",
    "\n",
    "            q1_loss_avgs = np.append(q1_loss_avgs, q1_loss.mean())\n",
    "            q2_loss_avgs = np.append(q2_loss_avgs, q2_loss.mean())\n",
    "            policy_loss_avgs = np.append(policy_loss_avgs, policy_loss.mean())\n",
    "            average_rewards = np.append(average_rewards, rewards.mean())\n",
    "            steps = np.append(steps, step)\n",
    "            crashes = np.append(crashes,crashed)\n",
    "            total_violations = np.append(total_violations,violations.sum())\n",
    "            \n",
    "            \n",
    "\n",
    "        name = \"BackupFinalNonScalarizedLagrange_alpha=\"+str(alpha) +\"_cost=0.001_seed=:\"+str(seed)+\".pt\"\n",
    "        Path = \"/home/padjei/Desktop/Models/Reviewers/\" + name  \n",
    "        \n",
    "        models_performance_tracker(seed=seed, network_size=network, q1_avg_loss=q1_loss_avgs, q2_avg_loss=q2_loss_avgs, \n",
    "                                   policy_avg_loss=policy_loss_avgs, avg_reward = average_rewards, steps=steps, crashes=crashes, \n",
    "                                    total_violations=total_violations, PATH=Path)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "rlenv1",
   "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.8.16"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}