LE_notebook_shapes.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "One-hot encoded 6-bit utterance vector for 3 colors and 2 shapes\n",
    "New AE trained on 3 colors and 2 shapes"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/imtezcan/Repositories/CogSci/EBIMAS/venv-ebimas/lib/python3.10/site-packages/pygame/pkgdata.py:25: DeprecationWarning: pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html\n",
      "  from pkg_resources import resource_stream, resource_exists\n",
      "/Users/imtezcan/Repositories/CogSci/EBIMAS/venv-ebimas/lib/python3.10/site-packages/pkg_resources/__init__.py:2871: DeprecationWarning: Deprecated call to `pkg_resources.declare_namespace('google')`.\n",
      "Implementing implicit namespace packages (as specified in PEP 420) is preferred to `pkg_resources.declare_namespace`. See https://setuptools.pypa.io/en/latest/references/keywords.html#keyword-namespace-packages\n",
      "  declare_namespace(pkg)\n",
      "/Users/imtezcan/Repositories/CogSci/EBIMAS/venv-ebimas/lib/python3.10/site-packages/pkg_resources/__init__.py:2871: DeprecationWarning: Deprecated call to `pkg_resources.declare_namespace('mpl_toolkits')`.\n",
      "Implementing implicit namespace packages (as specified in PEP 420) is preferred to `pkg_resources.declare_namespace`. See https://setuptools.pypa.io/en/latest/references/keywords.html#keyword-namespace-packages\n",
      "  declare_namespace(pkg)\n",
      "/Users/imtezcan/Repositories/CogSci/EBIMAS/venv-ebimas/lib/python3.10/site-packages/pkg_resources/__init__.py:2871: DeprecationWarning: Deprecated call to `pkg_resources.declare_namespace('sphinxcontrib')`.\n",
      "Implementing implicit namespace packages (as specified in PEP 420) is preferred to `pkg_resources.declare_namespace`. See https://setuptools.pypa.io/en/latest/references/keywords.html#keyword-namespace-packages\n",
      "  declare_namespace(pkg)\n",
      "/Users/imtezcan/Repositories/CogSci/EBIMAS/venv-ebimas/lib/python3.10/site-packages/ray/tune/logger/tensorboardx.py:35: DeprecationWarning: `np.bool8` is a deprecated alias for `np.bool_`.  (Deprecated NumPy 1.24)\n",
      "  VALID_NP_HPARAMS = (np.bool8, np.float32, np.float64, np.int32, np.int64)\n",
      "/Users/imtezcan/Repositories/CogSci/EBIMAS/venv-ebimas/lib/python3.10/site-packages/ray/tune/logger/tensorboardx.py:161: DeprecationWarning: `np.bool8` is a deprecated alias for `np.bool_`.  (Deprecated NumPy 1.24)\n",
      "  VALID_NP_HPARAMS = (np.bool8, np.float32, np.float64, np.int32, np.int64)\n",
      "/Users/imtezcan/Repositories/CogSci/EBIMAS/venv-ebimas/lib/python3.10/site-packages/tensorflow_probability/python/__init__.py:57: DeprecationWarning: distutils Version classes are deprecated. Use packaging.version instead.\n",
      "  if (distutils.version.LooseVersion(tf.__version__) <\n",
      "/Users/imtezcan/Repositories/CogSci/EBIMAS/venv-ebimas/lib/python3.10/site-packages/tensorflow_probability/python/__init__.py:58: DeprecationWarning: distutils Version classes are deprecated. Use packaging.version instead.\n",
      "  distutils.version.LooseVersion(required_tensorflow_version)):\n",
      "/Users/imtezcan/Repositories/CogSci/EBIMAS/venv-ebimas/lib/python3.10/site-packages/torch/utils/tensorboard/__init__.py:4: DeprecationWarning: distutils Version classes are deprecated. Use packaging.version instead.\n",
      "  if not hasattr(tensorboard, \"__version__\") or LooseVersion(\n",
      "/Users/imtezcan/Repositories/CogSci/EBIMAS/venv-ebimas/lib/python3.10/site-packages/torch/utils/tensorboard/__init__.py:6: DeprecationWarning: distutils Version classes are deprecated. Use packaging.version instead.\n",
      "  ) < LooseVersion(\"1.15\"):\n"
     ]
    }
   ],
   "source": [
    "from MuJoCo_Gym.mujoco_rl import MuJoCoRL\n",
    "from MuJoCo_Gym.wrappers import GymnasiumWrapper, GymWrapper\n",
    "from gymnasium.wrappers.frame_stack import FrameStack\n",
    "from gymnasium.experimental.wrappers import NormalizeObservationV0\n",
    "from dynamics import *\n",
    "import argparse\n",
    "import os\n",
    "import random\n",
    "import time\n",
    "from distutils.util import strtobool\n",
    "import gym\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.distributions.normal import Normal\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "\n",
    "from wrappers.record_episode_statistics import RecordEpisodeStatistics\n",
    "from wrappers.frame_stack import FrameStack\n",
    "from wrappers.normalizeObservation import NormalizeObservation\n",
    "from wrappers.normalizeRewards import NormalizeReward\n",
    "\n",
    "from progressbar import progressbar\n",
    "import tensorflow as tf\n",
    "\n",
    "# Check if eager execution is enabled\n",
    "if not tf.executing_eagerly():\n",
    "    tf.compat.v1.enable_eager_execution()"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-15T10:16:47.226484Z",
     "start_time": "2024-03-15T10:16:39.815225Z"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Dynamics"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Vision"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "outputs": [],
   "source": [
    "from autoencoder import load_autoencoder_model, get_a_single_image_embedding\n",
    "\n",
    "class Image:\n",
    "    def __init__(self, environment):\n",
    "        self.environment = environment\n",
    "        self.observation_space = {\"low\": [0 for _ in range(50)], \"high\": [1 for _ in range(50)]}\n",
    "        self.action_space = {\"low\": [], \"high\": []}\n",
    "        self.autoencoder = load_autoencoder_model(\"models/3colors_2shapes_50.tf\")\n",
    "        self.index = 0\n",
    "\n",
    "    def dynamic(self, agent, actions):\n",
    "        self.index = self.index + 1\n",
    "        image = self.environment.get_camera_data(agent + \"_camera\")\n",
    "        result = get_a_single_image_embedding(self.autoencoder, image)\n",
    "        # cv2.imwrite(f\"/Users/cowolff/Documents/GitHub/s.language_experiments/images/{agent}_{self.index}.png\", image)\n",
    "        return 0, result, False, {}"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-15T10:16:48.821671Z",
     "start_time": "2024-03-15T10:16:48.777981Z"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Communication"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "outputs": [],
   "source": [
    "def get_shape(shape_id):\n",
    "    if shape_id == 6:  # box\n",
    "        return np.array([1, 0])\n",
    "    elif shape_id == 2:  # sphere\n",
    "        return np.array([0, 1])\n",
    "    else:\n",
    "        raise ValueError(\"Unknown shape\")\n",
    "    \n",
    "def encode_to_one_hot_vector(color, shape):\n",
    "    combined_vector = [0, 0, 0, 0, 0, 0]\n",
    "    color_index = np.argmax(color)\n",
    "    shape_index = np.argmax(shape)\n",
    "    # [Red Box, Red Sphere, Green Box, Green Sphere, Blue Box, Blue Sphere]\n",
    "    combined_index = color_index * 2 + shape_index\n",
    "    combined_vector[combined_index] = 1\n",
    "    return combined_vector"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-15T10:16:49.835756Z",
     "start_time": "2024-03-15T10:16:49.829559Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "outputs": [],
   "source": [
    "class Communication:\n",
    "    def __init__(self, environment):\n",
    "        self.environment = environment\n",
    "        self.observation_space = {\"low\": [0, 0, 0, 0, 0, 0], \"high\": [1, 1, 1, 1, 1, 1]}\n",
    "        self.action_space = {\"low\": [0, 0, 0, 0, 0, 0], \"high\": [1, 1, 1, 1, 1, 1]}\n",
    "\n",
    "    def dynamic(self, agent, actions):\n",
    "        if \"utterance\" not in self.environment.data_store[agent].keys():\n",
    "            self.environment.data_store[agent][\"utterance\"] = None\n",
    "        if agent == \"receiver\":\n",
    "            utterance = [0, 0, 0, 0, 0, 0]\n",
    "            if \"target_color\" in self.environment.data_store.keys() and \"target_shape\" in self.environment.data_store.keys():\n",
    "                utterance = encode_to_one_hot_vector(self.environment.data_store[\"target_color\"], self.environment.data_store[\"target_shape\"])\n",
    "            observation = utterance\n",
    "            # print(f'Comm Receiver Utterance: {utterance}')  # TODO remove\n",
    "        elif agent == \"sender\":\n",
    "            utterance = [0, 0, 0, 0, 0, 0]\n",
    "            utterance[np.argmax(actions)] = 1\n",
    "            self.environment.data_store[agent][\"utterance\"] = actions\n",
    "            self.environment.data_store[agent][\"utterance_max\"] = utterance\n",
    "            # print(f'Comm Sender Utterance: {actions}, Utterance_max: {utterance}')  # TODO remove\n",
    "            observation = [0, 0, 0, 0, 0, 0]\n",
    "        else:\n",
    "            print(\"Unknown agent\")\n",
    "        return 0, observation, False, {}"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-15T10:16:50.829506Z",
     "start_time": "2024-03-15T10:16:50.825607Z"
    }
   }
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reward Dynamic Function\n",
    "The reward function for the receiver gives a positive reward for the agent to move closer to the current target object.<br/>\n",
    "The sender in turn gets a positive reward if it chooses the correct word encoding for the color and shape of the current target object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "outputs": [],
   "source": [
    "from sklearn.metrics import mean_squared_error\n",
    "from autoencoder import Autoencoder\n",
    "import math\n",
    "import mujoco"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-15T10:16:53.154500Z",
     "start_time": "2024-03-15T10:16:53.143574Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-15T10:16:53.647691Z",
     "start_time": "2024-03-15T10:16:53.638444Z"
    }
   },
   "outputs": [],
   "source": [
    "class Reward:\n",
    "    def __init__(self, environment):\n",
    "        self.environment = environment\n",
    "        self.observation_space = {\"low\": [], \"high\": []}\n",
    "        self.action_space = {\"low\": [], \"high\": []}\n",
    "        self.choices = [\"choice_1\", \"choice_2\"]\n",
    "\n",
    "    def dynamic(self, agent, actions):\n",
    "        if not \"target\" in self.environment.data_store.keys():\n",
    "            color = self.environment.get_data(\"reference_geom\")[\"color\"]\n",
    "            shape = get_shape(self.environment.get_data(\"reference_geom\")[\"shape\"])\n",
    "            # print(f'Color: {color}, Shape: {shape}')  # TODO remove\n",
    "            for choice in self.choices:\n",
    "                choice_color = self.environment.get_data(choice + \"_geom\")[\"color\"]\n",
    "                choice_shape = get_shape(self.environment.get_data(choice + \"_geom\")[\"shape\"])\n",
    "                # Check if both color and shape match\n",
    "                if (color == choice_color).all() and (shape == choice_shape).all():\n",
    "                    self.environment.data_store[\"target\"] = choice\n",
    "                    self.environment.data_store[\"target_color\"] = choice_color\n",
    "                    self.environment.data_store[\"target_shape\"] = choice_shape\n",
    "                    self.environment.data_store[\"last_distance\"] = copy.deepcopy(self.environment.distance(\"receiver_geom\", choice + \"_geom\"))\n",
    "                    break \n",
    "        if agent == \"receiver\":\n",
    "            target = self.environment.data_store[\"target\"]\n",
    "            new_distance = self.environment.distance(\"receiver_geom\", target + \"_geom\")\n",
    "            reward = (self.environment.data_store[\"last_distance\"] - new_distance) * 10\n",
    "            self.environment.data_store[\"last_distance\"] = copy.deepcopy(new_distance)\n",
    "        elif agent == \"sender\":\n",
    "            reference = [0, 0, 0, 0, 0, 0]\n",
    "            color = self.environment.data_store[\"target_color\"]\n",
    "            shape = self.environment.data_store[\"target_shape\"]\n",
    "            \n",
    "            reference = encode_to_one_hot_vector(color, shape)\n",
    "            reward = 0\n",
    "            if \"utterance\" in self.environment.data_store[agent].keys():\n",
    "                reward = -1 * mean_squared_error(reference, self.environment.data_store[agent][\"utterance\"])\n",
    "        return reward, [], False, {}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Accuracy Measurement Dynamic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-15T10:16:55.127407Z",
     "start_time": "2024-03-15T10:16:55.110181Z"
    }
   },
   "outputs": [],
   "source": [
    "class Accuracy:\n",
    "    def __init__(self, environment):\n",
    "        self.environment = environment\n",
    "        self.observation_space = {\"low\": [], \"high\": []}\n",
    "        self.action_space = {\"low\": [], \"high\": []}\n",
    "        self.accuracies = []\n",
    "        self.variances = []\n",
    "        self.sendAccuracies = []\n",
    "        self.sendVariances = []\n",
    "        self.currentSend = []\n",
    "        self.report_accuracy = {\"sender\": 0, \"receiver\": 0}\n",
    "\n",
    "    def dynamic(self, agent, actions):\n",
    "        choices = [\"choice_1\", \"choice_2\"]\n",
    "        variance = {\"choice_1\":1, \"choice_2\":-1}\n",
    "        if \"target\" in self.environment.data_store.keys():\n",
    "            if \"sendVariances\" not in self.environment.data_store.keys():\n",
    "                self.environment.data_store[\"sendVariances\"] = True\n",
    "                self.currentSend = [0, 0, 0, 0, 0, 0]\n",
    "            target = self.environment.data_store[\"target\"]\n",
    "            # if any(self.environment.collision(ankle, target + \"_geom\") for ankle in [\"left_leg_geom_2\", \"left_ankle_geom_2\", \"right_leg_geom_2\", \"right_ankle_geom_2\", \"back_leg_geom_2\", \"third_ankle_geom_2\", \"rightback_leg_geom_2\", \"fourth_ankle_geom_2\"]):\n",
    "            if self.environment.collision(\"receiver_geom\", target + \"_geom\"):\n",
    "                self.accuracies.append(1)\n",
    "                self.variances.append(variance[target])\n",
    "\n",
    "                if len(self.variances) > 50:\n",
    "                    report_variance = 1 - abs(sum(self.variances[-50:]) / 50)\n",
    "                    self.report_accuracy[\"receiver\"] = sum(self.accuracies[-50:]) / 50\n",
    "            # elif any(self.environment.collision(ankle, [choice for choice in choices if choice != target][0] + \"_geom\") for ankle in [\"left_leg_geom_2\", \"left_ankle_geom_2\", \"right_leg_geom_2\", \"right_ankle_geom_2\", \"back_leg_geom_2\", \"third_ankle_geom_2\", \"rightback_leg_geom_2\", \"fourth_ankle_geom_2\"]):\n",
    "            elif self.environment.collision(\"receiver_geom\", [choice for choice in choices if choice != target][0] + \"_geom\"):\n",
    "                self.accuracies.append(0)\n",
    "                self.variances.append(variance[[choice for choice in choices if choice != target][0]])\n",
    "\n",
    "                if len(self.variances) > 50:\n",
    "                    report_variance = 1 - abs(sum(self.variances[-50:]) / 50)\n",
    "                    self.report_accuracy[0] = sum(self.accuracies[-50:]) / 50\n",
    "            if \"utterance_max\" in self.environment.data_store[agent].keys():\n",
    "                color = self.environment.data_store[\"target_color\"]\n",
    "                shape = self.environment.data_store[\"target_shape\"]\n",
    "                reference = encode_to_one_hot_vector(color, shape)\n",
    "                self.currentSend = np.add(self.currentSend, self.environment.data_store[agent][\"utterance_max\"])\n",
    "                # print(f'Utterance Max: {self.environment.data_store[agent][\"utterance_max\"]}')  # TODO remove\n",
    "                # print(f'Reference: {reference}')  # TODO remove\n",
    "\n",
    "                if self.environment.data_store[agent][\"utterance_max\"]  == reference:\n",
    "                    self.sendAccuracies.append(1)\n",
    "                else:\n",
    "                    self.sendAccuracies.append(0)\n",
    "        return 0, [], False, {}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Reward Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-15T10:16:56.088444Z",
     "start_time": "2024-03-15T10:16:56.079499Z"
    }
   },
   "outputs": [],
   "source": [
    "def target_reward(mujoco_gym, agent):\n",
    "    if agent == \"receiver\":\n",
    "        choices = [\"choice_1\", \"choice_2\"]\n",
    "        if not \"target\" in mujoco_gym.data_store.keys():\n",
    "                color = mujoco_gym.get_data(\"reference_geom\")[\"color\"]\n",
    "                shape = get_shape(mujoco_gym.get_data(\"reference_geom\")[\"shape\"])\n",
    "                # print(f'Target Reward Color: {color}, Shape: {shape}')  # TODO remove\n",
    "                for choice in choices:\n",
    "                    choice_shape = get_shape(mujoco_gym.get_data(choice + \"_geom\")[\"shape\"])\n",
    "                    if (color == mujoco_gym.get_data(choice + \"_geom\")[\"color\"]).all() and (shape == choice_shape).all():\n",
    "                        mujoco_gym.data_store[\"target\"] = choice\n",
    "                        mujoco_gym.data_store[\"target_color\"] = mujoco_gym.get_data(choice + \"_geom\")[\"color\"]\n",
    "                        mujoco_gym.data_store[\"target_shape\"] = choice_shape\n",
    "                        # print(f'Target Reward target shape: {choice_shape}')  # TODO remove\n",
    "                        break\n",
    "        \n",
    "        target = mujoco_gym.data_store[\"target\"]\n",
    "        # for ankle in [\"left_leg_geom_2\", \"left_ankle_geom_2\", \"right_leg_geom_2\", \"right_ankle_geom_2\", \"back_leg_geom_2\", \"third_ankle_geom_2\", \"rightback_leg_geom_2\", \"fourth_ankle_geom_2\"]:\n",
    "        for ankle in [\"receiver_geom\"]:\n",
    "            if mujoco_gym.collision(ankle, target + \"_geom\"):\n",
    "                return 1\n",
    "            elif mujoco_gym.collision(ankle, [choice for choice in choices if choice != target][0] + \"_geom\"):\n",
    "                return -1\n",
    "    return 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-15T10:16:57.120745Z",
     "start_time": "2024-03-15T10:16:57.116981Z"
    }
   },
   "outputs": [],
   "source": [
    "def collision_reward(mujoco_gym, agent):\n",
    "    for border in [\"border1_geom\", \"border2_geom\", \"border3_geom\", \"border4_geom\", \"border5_geom\"]:\n",
    "        # for ankle in [\"left_leg_geom_2\", \"left_ankle_geom_2\", \"right_leg_geom_2\", \"right_ankle_geom_2\", \"back_leg_geom_2\", \"third_ankle_geom_2\", \"rightback_leg_geom_2\", \"fourth_ankle_geom_2\"]:\n",
    "        for ankle in [agent + \"_geom\"]:\n",
    "            if mujoco_gym.collision(border, ankle):\n",
    "                return -0.1\n",
    "    return 0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Environment Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-15T10:16:58.250565Z",
     "start_time": "2024-03-15T10:16:58.247496Z"
    }
   },
   "outputs": [],
   "source": [
    "def make_env(config_dict):\n",
    "    def thunk():\n",
    "        window = 5\n",
    "        env = MuJoCoRL(config_dict=config_dict)\n",
    "        # env = GymWrapper(env, \"receiver\")\n",
    "        # env = FrameStack(env, 4)\n",
    "        env = NormalizeObservation(env)\n",
    "        env = NormalizeReward(env)\n",
    "        # env = RecordEpisodeStatistics(env)\n",
    "        return env\n",
    "\n",
    "    return thunk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-15T10:17:11.937600Z",
     "start_time": "2024-03-15T10:17:10.213342Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/imtezcan/Repositories/CogSci/EBIMAS/venv-ebimas/lib/python3.10/site-packages/gymnasium/spaces/box.py:130: UserWarning: \u001B[33mWARN: Box bound precision lowered by casting to float32\u001B[0m\n",
      "  gym.logger.warn(f\"Box bound precision lowered by casting to {self.dtype}\")\n"
     ]
    }
   ],
   "source": [
    "# xml_files = [\"levels/\" + file for file in os.listdir(\"levels/\")]\n",
    "xml_files = [\"levels_shape/\" + file for file in os.listdir(\"levels_shape/\")]\n",
    "agents = [\"sender\", \"receiver\"]\n",
    "\n",
    "config_dict = {\"xmlPath\":xml_files, \n",
    "                   \"agents\":agents, \n",
    "                   \"rewardFunctions\":[collision_reward, target_reward], \n",
    "                   \"doneFunctions\":[target_done, border_done], \n",
    "                   \"skipFrames\":5,\n",
    "                   \"environmentDynamics\":[Image, Reward, Communication, Accuracy],\n",
    "                   \"freeJoint\":True,\n",
    "                   \"renderMode\":False,\n",
    "                   \"maxSteps\":1024,\n",
    "                   \"agentCameras\":True}\n",
    "\n",
    "env = make_env(config_dict)()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# RL Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-15T10:17:13.240410Z",
     "start_time": "2024-03-15T10:17:13.232846Z"
    }
   },
   "outputs": [],
   "source": [
    "def layer_init(layer, std=np.sqrt(2), bias_const=0.0):\n",
    "    torch.nn.init.orthogonal_(layer.weight, std)\n",
    "    torch.nn.init.constant_(layer.bias, bias_const)\n",
    "    return layer\n",
    "\n",
    "class Agent(nn.Module):\n",
    "    def __init__(self, envs):\n",
    "        super(Agent, self).__init__()\n",
    "        self.critic = nn.Sequential(\n",
    "            nn.Flatten(),\n",
    "            layer_init(nn.Linear(np.array(envs.observation_space.shape).prod(), 128)),\n",
    "            nn.Tanh(),\n",
    "            layer_init(nn.Linear(128, 128)),\n",
    "            nn.Tanh(),\n",
    "            layer_init(nn.Linear(128, 1), std=1.0),\n",
    "        )\n",
    "        self.actor_mean = nn.Sequential(\n",
    "            nn.Flatten(),\n",
    "            layer_init(nn.Linear(np.array(envs.observation_space.shape).prod(), 128)),\n",
    "            nn.Tanh(),\n",
    "            layer_init(nn.Linear(128, 128)),\n",
    "            nn.Tanh(),\n",
    "            layer_init(nn.Linear(128, np.prod(envs.action_space.shape)), std=0.01),\n",
    "        )\n",
    "        self.actor_logstd = nn.Parameter(torch.zeros(1, np.prod(envs.action_space.shape)))\n",
    "\n",
    "    def get_value(self, x):\n",
    "        return self.critic(x)\n",
    "\n",
    "    def get_action_and_value(self, x, action=None):\n",
    "        action_mean = self.actor_mean(x)\n",
    "        action_logstd = self.actor_logstd.expand_as(action_mean)\n",
    "        action_std = torch.exp(action_logstd)\n",
    "        probs = Normal(action_mean, action_std)\n",
    "        if action is None:\n",
    "            action = probs.sample()\n",
    "        return action, probs.log_prob(action).sum(1), probs.entropy().sum(1), self.critic(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Buffer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-15T10:17:13.975963Z",
     "start_time": "2024-03-15T10:17:13.972694Z"
    }
   },
   "outputs": [],
   "source": [
    "class Buffer():\n",
    "    def __init__(self, num_steps, envs, num_envs, device):\n",
    "        self.obs = torch.zeros((num_steps, num_envs) + envs.observation_space.shape).to(device)\n",
    "        self.actions = torch.zeros((num_steps, num_envs) + envs.action_space.shape).to(device)\n",
    "        self.logprobs = torch.zeros((num_steps, num_envs)).to(device)\n",
    "        self.rewards = torch.zeros((num_steps, num_envs)).to(device)\n",
    "        self.dones = torch.zeros((num_steps, num_envs)).to(device)\n",
    "        self.values = torch.zeros((num_steps, num_envs)).to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Update Function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-15T10:17:14.485160Z",
     "start_time": "2024-03-15T10:17:14.476395Z"
    }
   },
   "outputs": [],
   "source": [
    "def update_agent(agent, buffer, optimizer, next_obs, next_done, env, batch_size, update_epochs, minibatch_size, clip_coef, vf_coef, ent_coef, max_grad_norm, target_kl, clip_vloss, norm_adv, gae_lambda, gae, gamma, device, num_steps):\n",
    "\n",
    "    with torch.no_grad():\n",
    "        next_value = agent.get_value(next_obs).reshape(1, -1)\n",
    "        if gae:\n",
    "            advantages = torch.zeros_like(buffer.rewards).to(device)\n",
    "            lastgaelam = 0\n",
    "            for t in reversed(range(num_steps)):\n",
    "                if t == num_steps - 1:\n",
    "                    nextnonterminal = 1.0 - next_done\n",
    "                    nextvalues = next_value\n",
    "                else:\n",
    "                    nextnonterminal = 1.0 - buffer.dones[t + 1]\n",
    "                    nextvalues = buffer.values[t + 1]\n",
    "                delta = buffer.rewards[t] + gamma * nextvalues * nextnonterminal - buffer.values[t]\n",
    "                advantages[t] = lastgaelam = delta + gamma * gae_lambda * nextnonterminal * lastgaelam\n",
    "            returns = advantages + buffer.values\n",
    "        else:\n",
    "            returns = torch.zeros_like(buffer.rewards).to(device)\n",
    "            for t in reversed(range(num_steps)):\n",
    "                if t == num_steps - 1:\n",
    "                    nextnonterminal = 1.0 - next_done\n",
    "                    next_return = next_value\n",
    "                else:\n",
    "                    nextnonterminal = 1.0 - buffer.dones[t + 1]\n",
    "                    next_return = returns[t + 1]\n",
    "                returns[t] = buffer.rewards[t] + gamma * nextnonterminal * next_return\n",
    "            advantages = returns - buffer.values\n",
    "\n",
    "    # flatten the batch\n",
    "    b_obs = buffer.obs.reshape((-1,) + env.observation_space.shape)\n",
    "    b_logprobs = buffer.logprobs.reshape(-1)\n",
    "    b_actions = buffer.actions.reshape((-1,) + env.action_space.shape)\n",
    "    b_advantages = advantages.reshape(-1)\n",
    "    b_returns = returns.reshape(-1)\n",
    "    b_values = buffer.values.reshape(-1)\n",
    "\n",
    "    # Optimizing the policy and value network\n",
    "    b_inds = np.arange(batch_size)\n",
    "    clipfracs = []\n",
    "    for epoch in range(update_epochs):\n",
    "        np.random.shuffle(b_inds)\n",
    "        for start in range(0, batch_size, minibatch_size):\n",
    "            end = start + minibatch_size\n",
    "            mb_inds = b_inds[start:end]\n",
    "\n",
    "            _, newlogprob, entropy, newvalue = agent.get_action_and_value(b_obs[mb_inds], b_actions[mb_inds])\n",
    "            logratio = newlogprob - b_logprobs[mb_inds]\n",
    "            ratio = logratio.exp()\n",
    "\n",
    "            with torch.no_grad():\n",
    "                # calculate approx_kl http://joschu.net/blog/kl-approx.html\n",
    "                old_approx_kl = (-logratio).mean()\n",
    "                approx_kl = ((ratio - 1) - logratio).mean()\n",
    "                clipfracs += [((ratio - 1.0).abs() > clip_coef).float().mean().item()]\n",
    "\n",
    "            mb_advantages = b_advantages[mb_inds]\n",
    "            if norm_adv:\n",
    "                mb_advantages = (mb_advantages - mb_advantages.mean()) / (mb_advantages.std() + 1e-8)\n",
    "\n",
    "            # Policy loss\n",
    "            pg_loss1 = -mb_advantages * ratio\n",
    "            pg_loss2 = -mb_advantages * torch.clamp(ratio, 1 - clip_coef, 1 + clip_coef)\n",
    "            pg_loss = torch.max(pg_loss1, pg_loss2).mean()\n",
    "\n",
    "            # Value loss\n",
    "            newvalue = newvalue.view(-1)\n",
    "            if clip_vloss:\n",
    "                v_loss_unclipped = (newvalue - b_returns[mb_inds]) ** 2\n",
    "                v_clipped = b_values[mb_inds] + torch.clamp(\n",
    "                    newvalue - b_values[mb_inds],\n",
    "                    -clip_coef,\n",
    "                    clip_coef,\n",
    "                )\n",
    "                v_loss_clipped = (v_clipped - b_returns[mb_inds]) ** 2\n",
    "                v_loss_max = torch.max(v_loss_unclipped, v_loss_clipped)\n",
    "                v_loss = 0.5 * v_loss_max.mean()\n",
    "            else:\n",
    "                v_loss = 0.5 * ((newvalue - b_returns[mb_inds]) ** 2).mean()\n",
    "\n",
    "            entropy_loss = entropy.mean()\n",
    "            loss = pg_loss - ent_coef * entropy_loss + v_loss * vf_coef\n",
    "\n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            nn.utils.clip_grad_norm_(agent.parameters(), max_grad_norm)\n",
    "            optimizer.step()\n",
    "\n",
    "        if target_kl is not None:\n",
    "            if approx_kl > target_kl:\n",
    "                break\n",
    "\n",
    "    y_pred, y_true = b_values.cpu().numpy(), b_returns.cpu().numpy()\n",
    "    var_y = np.var(y_true)\n",
    "    explained_var = np.nan if var_y == 0 else 1 - np.var(y_true - y_pred) / var_y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Helper Function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-15T10:17:15.544952Z",
     "start_time": "2024-03-15T10:17:15.541964Z"
    }
   },
   "outputs": [],
   "source": [
    "def initialize_agent(env, device, learning_rate):\n",
    "    agent = Agent(env).to(device)\n",
    "    optimizer = optim.Adam(agent.parameters(), lr=learning_rate, eps=1e-5)\n",
    "    return agent, optimizer\n",
    "\n",
    "def get_action_and_update_buffer(agent, obs, buffer, step):\n",
    "    with torch.no_grad():\n",
    "        action, logprob, _, value = agent.get_action_and_value(obs)\n",
    "        buffer.values[step] = value.flatten()\n",
    "    buffer.actions[step] = action\n",
    "    buffer.logprobs[step] = logprob\n",
    "    return action\n",
    "\n",
    "def reset_environment(env, device):\n",
    "    next_obs, infos = env.reset()\n",
    "    next_obs = {k: torch.Tensor(v).unsqueeze(0).to(device) for k, v in next_obs.items()}\n",
    "    return next_obs, infos"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-15T10:17:16.080972Z",
     "start_time": "2024-03-15T10:17:16.065127Z"
    }
   },
   "outputs": [],
   "source": [
    "exp_name = \"3 Colors 2 Shapes New AE 6-bit Vector\"\n",
    "\n",
    "learning_rate = 1e-5\n",
    "seed = 1\n",
    "# total_timesteps = 20000000\n",
    "total_timesteps = 1000000\n",
    "torch_deterministic = True\n",
    "cuda = False\n",
    "mps = True\n",
    "track = False\n",
    "wandb_project_name = \"ppo-implementation-details\"\n",
    "wandb_entity = None\n",
    "capture_video = False\n",
    "\n",
    "# Algorithm-specific arguments\n",
    "num_envs = 1\n",
    "num_steps = 2048\n",
    "anneal_lr = True\n",
    "gae = True\n",
    "gamma = 0.99\n",
    "gae_lambda = 0.95\n",
    "num_minibatches = 128\n",
    "update_epochs = 10\n",
    "norm_adv = True\n",
    "clip_coef = 0.2\n",
    "clip_vloss = True\n",
    "ent_coef = 0.0\n",
    "vf_coef = 0.5\n",
    "max_grad_norm = 0.5\n",
    "target_kl = None\n",
    "store_freq = 20\n",
    "\n",
    "# Calculate derived variables\n",
    "batch_size = int(num_envs * num_steps)\n",
    "minibatch_size = int(batch_size // num_minibatches)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-15T10:17:46.063932Z",
     "start_time": "2024-03-15T10:17:17.551234Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0% (0 of 488) |                        | Elapsed Time: 0:00:00 ETA:  --:--:--"
     ]
    },
    {
     "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[17], line 65\u001B[0m\n\u001B[1;32m     62\u001B[0m sender_action \u001B[38;5;241m=\u001B[39m get_action_and_update_buffer(sender, next_obs[\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124msender\u001B[39m\u001B[38;5;124m\"\u001B[39m], buffer_sender, step)\n\u001B[1;32m     63\u001B[0m receiver_action \u001B[38;5;241m=\u001B[39m get_action_and_update_buffer(receiver, next_obs[\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mreceiver\u001B[39m\u001B[38;5;124m\"\u001B[39m], buffer_receiver, step)\n\u001B[0;32m---> 65\u001B[0m next_obs, reward, terminations, truncations, info \u001B[38;5;241m=\u001B[39m \u001B[43menv\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mstep\u001B[49m\u001B[43m(\u001B[49m\u001B[43m{\u001B[49m\u001B[38;5;124;43m\"\u001B[39;49m\u001B[38;5;124;43msender\u001B[39;49m\u001B[38;5;124;43m\"\u001B[39;49m\u001B[43m:\u001B[49m\u001B[43m \u001B[49m\u001B[43msender_action\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mcpu\u001B[49m\u001B[43m(\u001B[49m\u001B[43m)\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mnumpy\u001B[49m\u001B[43m(\u001B[49m\u001B[43m)\u001B[49m\u001B[43m[\u001B[49m\u001B[38;5;241;43m0\u001B[39;49m\u001B[43m]\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;124;43m\"\u001B[39;49m\u001B[38;5;124;43mreceiver\u001B[39;49m\u001B[38;5;124;43m\"\u001B[39;49m\u001B[43m:\u001B[49m\u001B[43m \u001B[49m\u001B[43mreceiver_action\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mcpu\u001B[49m\u001B[43m(\u001B[49m\u001B[43m)\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mnumpy\u001B[49m\u001B[43m(\u001B[49m\u001B[43m)\u001B[49m\u001B[43m[\u001B[49m\u001B[38;5;241;43m0\u001B[39;49m\u001B[43m]\u001B[49m\u001B[43m}\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m     66\u001B[0m current_rewards[\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124msender\u001B[39m\u001B[38;5;124m\"\u001B[39m]\u001B[38;5;241m.\u001B[39mappend(reward[\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124msender\u001B[39m\u001B[38;5;124m\"\u001B[39m])\n\u001B[1;32m     67\u001B[0m current_rewards[\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mreceiver\u001B[39m\u001B[38;5;124m\"\u001B[39m]\u001B[38;5;241m.\u001B[39mappend(reward[\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mreceiver\u001B[39m\u001B[38;5;124m\"\u001B[39m])\n",
      "File \u001B[0;32m~/Repositories/CogSci/EBIMAS/s.language_experiments/wrappers/normalizeRewards.py:80\u001B[0m, in \u001B[0;36mNormalizeReward.step\u001B[0;34m(self, action)\u001B[0m\n\u001B[1;32m     78\u001B[0m \u001B[38;5;28;01mdef\u001B[39;00m \u001B[38;5;21mstep\u001B[39m(\u001B[38;5;28mself\u001B[39m, action):\n\u001B[1;32m     79\u001B[0m \u001B[38;5;250m    \u001B[39m\u001B[38;5;124;03m\"\"\"Steps through the environment, normalizing the rewards returned.\"\"\"\u001B[39;00m\n\u001B[0;32m---> 80\u001B[0m     obs, rews, terminateds, truncateds, infos \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43menv\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mstep\u001B[49m\u001B[43m(\u001B[49m\u001B[43maction\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m     81\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[39mis_vector_env:\n\u001B[1;32m     82\u001B[0m         rews \u001B[38;5;241m=\u001B[39m {agent: np\u001B[38;5;241m.\u001B[39marray(rews[agent]) \u001B[38;5;28;01mfor\u001B[39;00m agent \u001B[38;5;129;01min\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39magents}\n",
      "File \u001B[0;32m~/Repositories/CogSci/EBIMAS/s.language_experiments/wrappers/normalizeObservation.py:74\u001B[0m, in \u001B[0;36mNormalizeObservation.step\u001B[0;34m(self, actions)\u001B[0m\n\u001B[1;32m     72\u001B[0m \u001B[38;5;28;01mdef\u001B[39;00m \u001B[38;5;21mstep\u001B[39m(\u001B[38;5;28mself\u001B[39m, actions):\n\u001B[1;32m     73\u001B[0m \u001B[38;5;250m    \u001B[39m\u001B[38;5;124;03m\"\"\"Steps through the environment and normalizes the observation.\"\"\"\u001B[39;00m\n\u001B[0;32m---> 74\u001B[0m     obs, rews, terminateds, truncateds, infos \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43menv\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mstep\u001B[49m\u001B[43m(\u001B[49m\u001B[43mactions\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m     75\u001B[0m     \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39mis_vector_env:\n\u001B[1;32m     76\u001B[0m         obs \u001B[38;5;241m=\u001B[39m {agent: \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39mnormalize(obs[agent]) \u001B[38;5;28;01mfor\u001B[39;00m agent \u001B[38;5;129;01min\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39magents}\n",
      "File \u001B[0;32m~/Repositories/CogSci/EBIMAS/venv-ebimas/lib/python3.10/site-packages/MuJoCo_Gym/mujoco_rl.py:262\u001B[0m, in \u001B[0;36mMuJoCoRL.step\u001B[0;34m(self, action)\u001B[0m\n\u001B[1;32m    259\u001B[0m data_store_copies \u001B[38;5;241m=\u001B[39m [copy\u001B[38;5;241m.\u001B[39mdeepcopy(\u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39mdata_store) \u001B[38;5;28;01mfor\u001B[39;00m _ \u001B[38;5;129;01min\u001B[39;00m \u001B[38;5;28mrange\u001B[39m(\u001B[38;5;28mlen\u001B[39m(\u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39menvironment_dynamics))]\n\u001B[1;32m    260\u001B[0m original_data_store \u001B[38;5;241m=\u001B[39m copy\u001B[38;5;241m.\u001B[39mdeepcopy(\u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39mdata_store)\n\u001B[0;32m--> 262\u001B[0m observations, rewards, terminations, infos, data_store_copies \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m__apply_dynamics\u001B[49m\u001B[43m(\u001B[49m\u001B[43maction\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mobservations\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mrewards\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mterminations\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43minfos\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mdata_store_copies\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m    264\u001B[0m \u001B[38;5;66;03m# self.data_store = original_data_store\u001B[39;00m\n\u001B[1;32m    265\u001B[0m \u001B[38;5;66;03m# for data in data_store_copies:\u001B[39;00m\n\u001B[1;32m    266\u001B[0m \u001B[38;5;66;03m#     self.data_store = update_deep(self.data_store, data)\u001B[39;00m\n\u001B[1;32m    268\u001B[0m \u001B[38;5;28;01mfor\u001B[39;00m reward \u001B[38;5;129;01min\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39mreward_functions:\n",
      "File \u001B[0;32m~/Repositories/CogSci/EBIMAS/venv-ebimas/lib/python3.10/site-packages/MuJoCo_Gym/mujoco_rl.py:230\u001B[0m, in \u001B[0;36mMuJoCoRL.__apply_dynamics\u001B[0;34m(self, action, observations, rewards, terminations, infos, data_store_copies)\u001B[0m\n\u001B[1;32m    228\u001B[0m dynamic_indizes \u001B[38;5;241m=\u001B[39m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39maction_routing[\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mdynamic\u001B[39m\u001B[38;5;124m\"\u001B[39m][dynamic\u001B[38;5;241m.\u001B[39m\u001B[38;5;18m__class__\u001B[39m\u001B[38;5;241m.\u001B[39m\u001B[38;5;18m__name__\u001B[39m]\n\u001B[1;32m    229\u001B[0m dynamic_actions \u001B[38;5;241m=\u001B[39m action[agent][dynamic_indizes[\u001B[38;5;241m0\u001B[39m]:dynamic_indizes[\u001B[38;5;241m1\u001B[39m]]\n\u001B[0;32m--> 230\u001B[0m reward, obs, done, info \u001B[38;5;241m=\u001B[39m \u001B[43mdynamic\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mdynamic\u001B[49m\u001B[43m(\u001B[49m\u001B[43magent\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mdynamic_actions\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m    231\u001B[0m observations[agent] \u001B[38;5;241m=\u001B[39m np\u001B[38;5;241m.\u001B[39mconcatenate((observations[agent], obs))\n\u001B[1;32m    232\u001B[0m rewards[agent] \u001B[38;5;241m+\u001B[39m\u001B[38;5;241m=\u001B[39m reward\n",
      "Cell \u001B[0;32mIn[2], line 14\u001B[0m, in \u001B[0;36mImage.dynamic\u001B[0;34m(self, agent, actions)\u001B[0m\n\u001B[1;32m     12\u001B[0m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39mindex \u001B[38;5;241m=\u001B[39m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39mindex \u001B[38;5;241m+\u001B[39m \u001B[38;5;241m1\u001B[39m\n\u001B[1;32m     13\u001B[0m image \u001B[38;5;241m=\u001B[39m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39menvironment\u001B[38;5;241m.\u001B[39mget_camera_data(agent \u001B[38;5;241m+\u001B[39m \u001B[38;5;124m\"\u001B[39m\u001B[38;5;124m_camera\u001B[39m\u001B[38;5;124m\"\u001B[39m)\n\u001B[0;32m---> 14\u001B[0m result \u001B[38;5;241m=\u001B[39m \u001B[43mget_a_single_image_embedding\u001B[49m\u001B[43m(\u001B[49m\u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mautoencoder\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mimage\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m     15\u001B[0m \u001B[38;5;66;03m# cv2.imwrite(f\"/Users/cowolff/Documents/GitHub/s.language_experiments/images/{agent}_{self.index}.png\", image)\u001B[39;00m\n\u001B[1;32m     16\u001B[0m \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[38;5;241m0\u001B[39m, result, \u001B[38;5;28;01mFalse\u001B[39;00m, {}\n",
      "File \u001B[0;32m~/Repositories/CogSci/EBIMAS/s.language_experiments/autoencoder.py:102\u001B[0m, in \u001B[0;36mget_a_single_image_embedding\u001B[0;34m(autoencoder, image)\u001B[0m\n\u001B[1;32m    100\u001B[0m \u001B[38;5;250m\u001B[39m\u001B[38;5;124;03m\"\"\"Load and encode a single image.\"\"\"\u001B[39;00m\n\u001B[1;32m    101\u001B[0m image \u001B[38;5;241m=\u001B[39m preprocess_image(image)\n\u001B[0;32m--> 102\u001B[0m encoded_image \u001B[38;5;241m=\u001B[39m \u001B[43mautoencoder\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mencode\u001B[49m\u001B[43m(\u001B[49m\u001B[43mimage\u001B[49m\u001B[43m[\u001B[49m\u001B[38;5;28;43;01mNone\u001B[39;49;00m\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[38;5;241;43m.\u001B[39;49m\u001B[43m]\u001B[49m\u001B[43m)\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mnumpy\u001B[49m\u001B[43m(\u001B[49m\u001B[43m)\u001B[49m[\u001B[38;5;241m0\u001B[39m]\n\u001B[1;32m    103\u001B[0m \u001B[38;5;28;01mreturn\u001B[39;00m encoded_image\n",
      "File \u001B[0;32m~/Repositories/CogSci/EBIMAS/venv-ebimas/lib/python3.10/site-packages/tensorflow/python/framework/ops.py:1141\u001B[0m, in \u001B[0;36m_EagerTensorBase.numpy\u001B[0;34m(self)\u001B[0m\n\u001B[1;32m   1118\u001B[0m \u001B[38;5;250m\u001B[39m\u001B[38;5;124;03m\"\"\"Copy of the contents of this Tensor into a NumPy array or scalar.\u001B[39;00m\n\u001B[1;32m   1119\u001B[0m \n\u001B[1;32m   1120\u001B[0m \u001B[38;5;124;03mUnlike NumPy arrays, Tensors are immutable, so this method has to copy\u001B[39;00m\n\u001B[0;32m   (...)\u001B[0m\n\u001B[1;32m   1138\u001B[0m \u001B[38;5;124;03m    NumPy dtype.\u001B[39;00m\n\u001B[1;32m   1139\u001B[0m \u001B[38;5;124;03m\"\"\"\u001B[39;00m\n\u001B[1;32m   1140\u001B[0m \u001B[38;5;66;03m# TODO(slebedev): Consider avoiding a copy for non-CPU or remote tensors.\u001B[39;00m\n\u001B[0;32m-> 1141\u001B[0m maybe_arr \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m_numpy\u001B[49m\u001B[43m(\u001B[49m\u001B[43m)\u001B[49m  \u001B[38;5;66;03m# pylint: disable=protected-access\u001B[39;00m\n\u001B[1;32m   1142\u001B[0m \u001B[38;5;28;01mreturn\u001B[39;00m maybe_arr\u001B[38;5;241m.\u001B[39mcopy() \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;28misinstance\u001B[39m(maybe_arr, np\u001B[38;5;241m.\u001B[39mndarray) \u001B[38;5;28;01melse\u001B[39;00m maybe_arr\n",
      "File \u001B[0;32m~/Repositories/CogSci/EBIMAS/venv-ebimas/lib/python3.10/site-packages/tensorflow/python/framework/ops.py:1107\u001B[0m, in \u001B[0;36m_EagerTensorBase._numpy\u001B[0;34m(self)\u001B[0m\n\u001B[1;32m   1105\u001B[0m \u001B[38;5;28;01mdef\u001B[39;00m \u001B[38;5;21m_numpy\u001B[39m(\u001B[38;5;28mself\u001B[39m):\n\u001B[1;32m   1106\u001B[0m   \u001B[38;5;28;01mtry\u001B[39;00m:\n\u001B[0;32m-> 1107\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_numpy_internal\u001B[49m\u001B[43m(\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m   1108\u001B[0m   \u001B[38;5;28;01mexcept\u001B[39;00m core\u001B[38;5;241m.\u001B[39m_NotOkStatusException \u001B[38;5;28;01mas\u001B[39;00m e:  \u001B[38;5;66;03m# pylint: disable=protected-access\u001B[39;00m\n\u001B[1;32m   1109\u001B[0m     \u001B[38;5;28;01mraise\u001B[39;00m core\u001B[38;5;241m.\u001B[39m_status_to_exception(e) \u001B[38;5;28;01mfrom\u001B[39;00m \u001B[38;5;28mNone\u001B[39m\n",
      "\u001B[0;31mKeyboardInterrupt\u001B[0m: "
     ]
    }
   ],
   "source": [
    "torch.set_default_dtype(torch.float32)\n",
    "\n",
    "run_name = f\"{exp_name}__{seed}__{int(time.time())}\"\n",
    "\n",
    "writer = SummaryWriter(f\"runs/{run_name}\")\n",
    "\n",
    "writer.add_text(\"environment/level_number\", str(len(xml_files)), 0)\n",
    "writer.add_text(\"environment/agents\", ', '.join(agents), 0)\n",
    "writer.add_text(\"hyperparameters/learning_rate\", str(learning_rate), 0)\n",
    "writer.add_text(\"hyperparameters/network_size\", ', '.join(str(e) for e in [512, 256]), 0)\n",
    "writer.add_text(\"hyperparameters/batch\", str(minibatch_size), 0)\n",
    "\n",
    "# TRY NOT TO MODIFY: seeding\n",
    "random.seed(seed)\n",
    "np.random.seed(seed)\n",
    "torch.manual_seed(seed)\n",
    "torch.backends.cudnn.deterministic = torch_deterministic\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() and cuda else \"cpu\")\n",
    "\n",
    "obs, infos = env.reset()\n",
    "\n",
    "sender, sender_optimizer = initialize_agent(env, device, learning_rate)\n",
    "receiver, receiver_optimizer = initialize_agent(env, device, learning_rate)\n",
    "\n",
    "buffer_sender = Buffer(num_steps, env, num_envs, device)\n",
    "buffer_receiver = Buffer(num_steps, env, num_envs, device)\n",
    "\n",
    "global_step = 0\n",
    "start_time = time.time()\n",
    "next_obs, infos = reset_environment(env, device)\n",
    "\n",
    "next_done = {\"sender\": torch.zeros(num_envs).to(device), \"receiver\": torch.zeros(num_envs).to(device)}\n",
    "\n",
    "num_updates = total_timesteps // batch_size\n",
    "train_start = time.time()\n",
    "\n",
    "epoch_lengths = []\n",
    "current_length = 0\n",
    "\n",
    "for update in progressbar(range(1, num_updates + 1), redirect_stdout=True):\n",
    "    # Annealing the rate if instructed to do so.\n",
    "    if anneal_lr:\n",
    "        frac = 1.0 - (update - 1.0) / num_updates\n",
    "        lrnow = frac * learning_rate\n",
    "        sender_optimizer.param_groups[0][\"lr\"] = lrnow\n",
    "        receiver_optimizer.param_groups[0][\"lr\"] = lrnow\n",
    "    \n",
    "    epoch_rewards = {\"sender\":0, \"receiver\":0}\n",
    "    current_rewards = {\"sender\":[], \"receiver\":[]}\n",
    "    variances = {\"sender\":[], \"receiver\":[]}\n",
    "    epoch_runs = 0\n",
    "    episode_accuracies = 0\n",
    "    episode_sendAccuracies = 0\n",
    "    for step in range(0, num_steps):\n",
    "        global_step += 1 * num_envs\n",
    "        current_length += 1\n",
    "        buffer_sender.obs[step] = next_obs[\"sender\"]\n",
    "        buffer_receiver.obs[step] = next_obs[\"receiver\"]\n",
    "\n",
    "\n",
    "        sender_action = get_action_and_update_buffer(sender, next_obs[\"sender\"], buffer_sender, step)\n",
    "        receiver_action = get_action_and_update_buffer(receiver, next_obs[\"receiver\"], buffer_receiver, step)\n",
    "\n",
    "        next_obs, reward, terminations, truncations, info = env.step({\"sender\": sender_action.cpu().numpy()[0], \"receiver\": receiver_action.cpu().numpy()[0]})\n",
    "        current_rewards[\"sender\"].append(reward[\"sender\"])\n",
    "        current_rewards[\"receiver\"].append(reward[\"receiver\"])\n",
    "        next_obs = {\"sender\": torch.Tensor(next_obs[\"sender\"]).unsqueeze(0).to(device), \"receiver\": torch.Tensor(next_obs[\"receiver\"]).unsqueeze(0).to(device)}\n",
    "\n",
    "        if terminations[\"sender\"] or terminations[\"receiver\"] or truncations[\"sender\"] or truncations[\"receiver\"]:\n",
    "            next_obs, infos = reset_environment(env, device)\n",
    "            epoch_rewards[\"sender\"] += sum(current_rewards[\"sender\"])\n",
    "            epoch_rewards[\"receiver\"] += sum(current_rewards[\"receiver\"])\n",
    "\n",
    "            epoch_lengths.append(current_length)\n",
    "            current_length = 0\n",
    "\n",
    "            dynamic = env.env.env.environment_dynamics[3]\n",
    "\n",
    "            if len(dynamic.sendAccuracies) > 512:\n",
    "                episode_sendAccuracies = sum(dynamic.sendAccuracies[-512:]) / 512\n",
    "                del dynamic.sendAccuracies[:-513]\n",
    "                writer.add_scalar(\"charts/sender/accuracies\", episode_sendAccuracies, global_step)\n",
    "\n",
    "            if len(dynamic.accuracies) > 4:\n",
    "                window = min(15, len(dynamic.accuracies))\n",
    "                episode_accuracies = sum(dynamic.accuracies[-1 * window:]) / window\n",
    "                writer.add_scalar(\"charts/receiver/accuracies\", episode_accuracies, global_step)\n",
    "                if window == 15:\n",
    "                    del dynamic.accuracies[:-16]\n",
    "\n",
    "            if len(dynamic.variances) > 4:\n",
    "                window = min(15, len(dynamic.variances))\n",
    "                current_variance = sum(dynamic.variances[-1 * window:]) / window\n",
    "                writer.add_scalar(\"charts/receiver_variance\", current_variance, global_step)\n",
    "                if window == 15:\n",
    "                    del dynamic.variances[:-16]\n",
    "\n",
    "            if len(epoch_lengths) > 3:\n",
    "                window = min(10, len(epoch_lengths))\n",
    "                epoch_length = sum(epoch_lengths[-1 * window:]) / window\n",
    "                writer.add_scalar(\"charts/episodic_length\", epoch_length, global_step)\n",
    "                if window == 10:\n",
    "                    del epoch_lengths[:-11]\n",
    "            epoch_runs += 1\n",
    "        \n",
    "        buffer_sender.rewards[step] = torch.tensor(reward[\"sender\"]).to(device).view(-1)\n",
    "        buffer_receiver.rewards[step] = torch.tensor(reward[\"receiver\"]).to(device).view(-1)\n",
    "        next_done = {\"sender\": torch.Tensor([terminations[\"sender\"]]).to(device), \"receiver\": torch.Tensor([terminations[\"receiver\"]]).to(device)}\n",
    "    if update % store_freq == 0:\n",
    "        torch.save(sender, \"models/model\" + str(start_time) + \".pth\")\n",
    "        torch.save(receiver, \"models/model\" + str(start_time) + \".pth\")\n",
    "\n",
    "    update_agent(sender, buffer_sender, sender_optimizer, next_obs[\"sender\"], next_done[\"sender\"], env, batch_size, update_epochs, minibatch_size, clip_coef, vf_coef, ent_coef, max_grad_norm, target_kl, clip_vloss, norm_adv, gae_lambda, gae, gamma, device, step)\n",
    "    update_agent(receiver, buffer_receiver, receiver_optimizer, next_obs[\"receiver\"], next_done[\"receiver\"], env, batch_size, update_epochs, minibatch_size, clip_coef, vf_coef, ent_coef, max_grad_norm, target_kl, clip_vloss, norm_adv, gae_lambda, gae, gamma, device, step)\n",
    "\n",
    "    writer.add_scalar(\"charts/learning_rate\", sender_optimizer.param_groups[0][\"lr\"], global_step)\n",
    "    writer.add_scalar(\"charts/sender/episodic_return\", epoch_rewards[\"sender\"] / epoch_runs, global_step)\n",
    "    writer.add_scalar(\"charts/receiver/episodic_return\", epoch_rewards[\"receiver\"] / epoch_runs, global_step)\n",
    "    print(\"SPS:\", int(global_step / (time.time() - start_time)), \"Average Reward:\", epoch_rewards[\"sender\"] / epoch_runs)\n",
    "    writer.add_scalar(\"charts/SPS\", int(global_step / (time.time() - start_time)), global_step)\n",
    "\n",
    "torch.save(sender, \"models/model\" + str(start_time) + \".pth\")\n",
    "torch.save(receiver, \"models/model\" + str(start_time) + \".pth\")\n",
    "writer.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-07T00:28:26.183844Z",
     "start_time": "2024-03-07T00:28:26.102909Z"
    }
   },
   "outputs": [],
   "source": [
    "torch.save(sender, \"models/sender\" + str(start_time) + \".pth\")\n",
    "torch.save(receiver, \"models/receiver\" + str(start_time) + \".pth\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Direct communication test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-12T09:53:23.686542Z",
     "start_time": "2024-03-12T09:53:23.627754Z"
    }
   },
   "outputs": [],
   "source": [
    "class Test_Communication:\n",
    "    def __init__(self, environment):\n",
    "        self.environment = environment\n",
    "        self.observation_space = {\"low\": [0, 0, 0, 0, 0, 0], \"high\": [1, 1, 1, 1, 1, 1]}\n",
    "        self.action_space = {\"low\": [0, 0, 0, 0, 0, 0], \"high\": [1, 1, 1, 1, 1, 1]}\n",
    "\n",
    "    def dynamic(self, agent, actions):\n",
    "        if \"utterance\" not in self.environment.data_store[agent].keys():\n",
    "            self.environment.data_store[agent][\"utterance\"] = None\n",
    "        if agent == \"receiver\":\n",
    "            utterance = [0, 0, 0, 0, 0, 0]\n",
    "            if \"utterance_max\" in self.environment.data_store[\"sender\"].keys():\n",
    "                observation = self.environment.data_store[\"sender\"][\"utterance_max\"]\n",
    "            else:\n",
    "                observation = utterance\n",
    "        elif agent == \"sender\":\n",
    "            utterance = [0, 0, 0, 0, 0, 0]\n",
    "            utterance[np.argmax(actions)] = 1\n",
    "            self.environment.data_store[agent][\"utterance\"] = actions\n",
    "            self.environment.data_store[agent][\"utterance_max\"] = utterance\n",
    "            observation = [0, 0, 0, 0, 0, 0]\n",
    "        else:\n",
    "            print(\"Dafaq is going on here?\")\n",
    "        return 0, observation, False, {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-12T09:53:24.006505Z",
     "start_time": "2024-03-12T09:53:23.632358Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/imtezcan/Repositories/CogSci/EBIMAS/venv-ebimas/lib/python3.10/site-packages/gymnasium/spaces/box.py:130: UserWarning: \u001B[33mWARN: Box bound precision lowered by casting to float32\u001B[0m\n",
      "  gym.logger.warn(f\"Box bound precision lowered by casting to {self.dtype}\")\n"
     ]
    }
   ],
   "source": [
    "xml_files = [\"levels_shape_backup/modified/\" + file for file in os.listdir(\"levels_shape_backup/modified/\")]\n",
    "agents = [\"sender\", \"receiver\"]\n",
    "\n",
    "config_dict = {\"xmlPath\":xml_files, \n",
    "                   \"agents\":agents, \n",
    "                   \"rewardFunctions\":[collision_reward, target_reward], \n",
    "                   \"doneFunctions\":[target_done, border_done], \n",
    "                   \"skipFrames\":5,\n",
    "                   \"environmentDynamics\":[Image, Reward, Test_Communication, Accuracy],\n",
    "                   \"freeJoint\":True,\n",
    "                   \"renderMode\":False,\n",
    "                   \"maxSteps\":1024,\n",
    "                   \"agentCameras\":True}\n",
    "\n",
    "env = make_env(config_dict)()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-12T09:53:52.000788Z",
     "start_time": "2024-03-12T09:53:24.004042Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 30/30 [00:27<00:00,  1.10it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy: 0.5\n",
      "Variance: 1.0\n",
      "Send Accuracy: 0.6164179104477612\n",
      "Length: 63.96666666666667\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "from tqdm import tqdm\n",
    "\n",
    "epochs = 30\n",
    "num_steps = 1024\n",
    "lengths = []\n",
    "\n",
    "random.seed(seed)\n",
    "np.random.seed(seed)\n",
    "torch.manual_seed(seed)\n",
    "torch.backends.cudnn.deterministic = torch_deterministic\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() and cuda else \"cpu\")\n",
    "\n",
    "for epoch in tqdm(range(epochs)):\n",
    "    next_obs, infos = reset_environment(env, device)\n",
    "\n",
    "    next_obs = {k: torch.Tensor(v).unsqueeze(0).to(device) for k, v in next_obs.items()}\n",
    "    next_done = {\"sender\": torch.zeros(num_envs).to(device), \"receiver\": torch.zeros(num_envs).to(device)}\n",
    "\n",
    "    for step in range(0, num_steps):\n",
    "        sender_action = sender.get_action_and_value(next_obs[\"sender\"])[0]\n",
    "        receiver_action = receiver.get_action_and_value(next_obs[\"receiver\"])[0]\n",
    "\n",
    "        next_obs, reward, terminations, truncations, info = env.step({\"sender\": sender_action.cpu().numpy()[0], \"receiver\": receiver_action.cpu().numpy()[0]})\n",
    "        next_obs = {\"sender\": torch.Tensor(next_obs[\"sender\"]).unsqueeze(0).to(device), \"receiver\": torch.Tensor(next_obs[\"receiver\"]).unsqueeze(0).to(device)}\n",
    "\n",
    "        if terminations[\"sender\"] or terminations[\"receiver\"] or truncations[\"sender\"] or truncations[\"receiver\"]:\n",
    "            next_obs, infos = reset_environment(env, device)\n",
    "            lengths.append(step)\n",
    "            break\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy: 0.5\n",
      "Variance: 1.0\n",
      "Send Accuracy: 0.6164179104477612\n",
      "Length: 63.96666666666667\n"
     ]
    }
   ],
   "source": [
    "dynamic = env.env.env.environment_dynamics[3]\n",
    "print(\"Accuracy:\", sum(dynamic.accuracies) / len(dynamic.accuracies))\n",
    "print(\"Variance:\", sum(dynamic.variances) / len(dynamic.variances))\n",
    "print(\"Send Accuracy:\", sum(dynamic.sendAccuracies) / len(dynamic.sendAccuracies))\n",
    "print(\"Length:\", sum(lengths) / len(lengths))"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-12T09:53:52.005140Z",
     "start_time": "2024-03-12T09:53:51.194632Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [],
   "metadata": {
    "collapsed": false
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Ray",
   "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.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}