LE_notebook_vectorized.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "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"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Dynamics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics import mean_squared_error\n",
    "from autoencoder import Autoencoder\n",
    "import math\n",
    "import mujoco"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Vision"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "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 = Autoencoder(latent_dim=50, input_shape=(64, 64, 3))\n",
    "        self.autoencoder.encoder.load_weights(\"models/encoder50.h5\")\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",
    "        image = cv2.resize(image, (64, 64))\n",
    "        result = self.autoencoder.encoder.predict(np.array([image]), verbose=0)[0]\n",
    "        # cv2.imwrite(f\"/Users/cowolff/Documents/GitHub/s.language_experiments/images/{agent}_{self.index}.png\", image)\n",
    "        return 0, result, False, {}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Communication"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Communication:\n",
    "    def __init__(self, environment):\n",
    "        self.environment = environment\n",
    "        self.observation_space = {\"low\": [0, 0, 0, 0], \"high\": [1, 1, 1, 1]}\n",
    "        self.action_space = {\"low\": [0, 0, 0, 0], \"high\": [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]\n",
    "            if \"target_color\" in self.environment.data_store.keys():\n",
    "                utterance[np.argmax(self.environment.data_store[\"target_color\"])] = 1\n",
    "            observation = utterance\n",
    "        elif agent == \"sender\":\n",
    "            utterance = [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]\n",
    "        else:\n",
    "            print(\"Dafaq is going on here?\")\n",
    "        return 0, observation, False, {}"
   ]
  },
  {
   "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 of the current target object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "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",
    "            for choice in self.choices:\n",
    "                if (color == self.environment.get_data(choice + \"_geom\")[\"color\"]).all():\n",
    "                    self.environment.data_store[\"target\"] = choice\n",
    "                    self.environment.data_store[\"target_color\"] = self.environment.get_data(choice + \"_geom\")[\"color\"]\n",
    "                    self.environment.data_store[\"last_distance\"] = copy.deepcopy(self.environment.distance(\"receiver_geom\", choice + \"_geom\"))\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]\n",
    "            color = self.environment.data_store[\"target_color\"]\n",
    "            reference[np.argmax(color)] = 1\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": 6,
   "metadata": {},
   "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]\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",
    "                reference = [0, 0, 0, 0]\n",
    "                color = self.environment.data_store[\"target_color\"]\n",
    "                reference[np.argmax(color)] = 1\n",
    "                self.currentSend = np.add(self.currentSend, self.environment.data_store[agent][\"utterance_max\"])\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": 7,
   "metadata": {},
   "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",
    "                for choice in choices:\n",
    "                    if (color == mujoco_gym.get_data(choice + \"_geom\")[\"color\"]).all():\n",
    "                        mujoco_gym.data_store[\"target\"] = choice\n",
    "                        mujoco_gym.data_store[\"target_color\"] = mujoco_gym.get_data(choice + \"_geom\")[\"color\"]\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": 8,
   "metadata": {},
   "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": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_env(config_dict):\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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/cowolff/miniconda3/envs/Ray/lib/python3.9/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",
      "/Users/cowolff/miniconda3/envs/Ray/lib/python3.9/site-packages/keras/src/engine/training_v1.py:2359: UserWarning: `Model.state_updates` will be removed in a future version. This property should not be used in TensorFlow 2.0, as `updates` are applied automatically.\n",
      "  updates=self.state_updates,\n"
     ]
    }
   ],
   "source": [
    "xml_files = [\"levels/\" + file for file in os.listdir(\"levels/\")]\n",
    "agents = [\"sender\", \"receiver\"]\n",
    "num_envs = 1\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",
    "envs = [make_env(config_dict) for _ in range(num_envs)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[<wrappers.normalizeRewards.NormalizeReward object at 0x10b7f6c70>]\n"
     ]
    }
   ],
   "source": [
    "print(envs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# RL Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "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": {},
   "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": {},
   "outputs": [],
   "source": [
    "def update_agent(agent, buffer, optimizer, next_obs, next_done, envs, 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",
    "    env = envs[0]\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": {},
   "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": {},
   "outputs": [],
   "source": [
    "exp_name = \"Sender box\"\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 = False\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_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": {},
   "outputs": [],
   "source": [
    "def flatten_list(data):\n",
    "    flat_list = []\n",
    "    for item in data:\n",
    "        flat_list.extend(item.values())\n",
    "    return flat_list\n",
    "\n",
    "def reverse_flatten_list_with_agent_list(flat_list, agent_names):\n",
    "    reversed_data = []\n",
    "    for i in range(0, len(flat_list), len(agent_names)):\n",
    "        # Creating a dictionary for each set of lists corresponding to the agent names\n",
    "        entry = {agent_names[j]: flat_list[i + j] for j in range(len(agent_names))}\n",
    "        reversed_data.append(entry)\n",
    "    return reversed_data\n",
    "\n",
    "def get_observations_reset(envs, device):\n",
    "    obs = {}\n",
    "    infos = {}\n",
    "    agents = envs[0].agents\n",
    "    for agent in agents:\n",
    "        obs[agent] = []\n",
    "        infos[agent] = []\n",
    "    for env in envs:\n",
    "        next_obs, info = env.reset()\n",
    "        for agent in agents:\n",
    "            obs[agent].append(next_obs[agent])\n",
    "            infos[agent].append(info[agent])\n",
    "    for agent in agents:\n",
    "        obs[agent] = torch.Tensor(np.array(obs[agent])).to(device)\n",
    "    return obs, infos\n",
    "\n",
    "def get_data_step(envs, device, actions):\n",
    "    obs = {}\n",
    "    rewards = {}\n",
    "    terminations = {}\n",
    "    truncations = {}\n",
    "    infos = {}\n",
    "    new_actions = {}\n",
    "    for i, env in enumerate(envs):\n",
    "        current_actions = {agent: actions[agent][i] for agent in env.agents}\n",
    "        next_obs, next_rewards, next_terminations, next_truncations, next_infos = env.step(current_actions)\n",
    "        for agent in env.agents:\n",
    "            if agent not in obs.keys():\n",
    "                obs[agent] = []\n",
    "                rewards[agent] = []\n",
    "                terminations[agent] = []\n",
    "                truncations[agent] = []\n",
    "                infos[agent] = []\n",
    "                new_actions[agent] = []\n",
    "            obs[agent].append(next_obs[agent])\n",
    "            rewards[agent].append(next_rewards[agent])\n",
    "            terminations[agent].append(next_terminations[agent])\n",
    "            truncations[agent].append(next_truncations[agent])\n",
    "            infos[agent].append(next_infos[agent])\n",
    "            new_actions[agent].append(next_obs[agent])\n",
    "    return obs, rewards, terminations, truncations, infos, new_actions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0% (0 of 488) |                        | Elapsed Time: 0:00:00 ETA:  --:--:--\n"
     ]
    },
    {
     "ename": "TypeError",
     "evalue": "unsupported operand type(s) for -: 'float' and 'list'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[39], line 126\u001b[0m\n\u001b[1;32m    122\u001b[0m             \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mSPS:\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;28mint\u001b[39m(global_step \u001b[38;5;241m/\u001b[39m (time\u001b[38;5;241m.\u001b[39mtime() \u001b[38;5;241m-\u001b[39m start_time)), \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mAverage Reward:\u001b[39m\u001b[38;5;124m\"\u001b[39m, epoch_rewards[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msender\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m/\u001b[39m epoch_runs, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mAverage Lengtht:\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;28msum\u001b[39m(epoch_lengths[\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m100\u001b[39m:]) \u001b[38;5;241m/\u001b[39m \u001b[38;5;28mmin\u001b[39m(\u001b[38;5;28mlen\u001b[39m(epoch_runs), \u001b[38;5;241m100\u001b[39m))\n\u001b[1;32m    124\u001b[0m     next_done \u001b[38;5;241m=\u001b[39m current_dones\n\u001b[0;32m--> 126\u001b[0m \u001b[43mupdate_agent\u001b[49m\u001b[43m(\u001b[49m\u001b[43msender\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mbuffer_sender\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43msender_optimizer\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mnext_obs\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[43m \u001b[49m\u001b[43mnext_done\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[43m \u001b[49m\u001b[43menvs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mbatch_size\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mupdate_epochs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mminibatch_size\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mclip_coef\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mvf_coef\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43ment_coef\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmax_grad_norm\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtarget_kl\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mclip_vloss\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mnorm_adv\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mgae_lambda\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mgae\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mgamma\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdevice\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mstep\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    127\u001b[0m update_agent(receiver, buffer_receiver, receiver_optimizer, next_obs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mreceiver\u001b[39m\u001b[38;5;124m\"\u001b[39m], next_done[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mreceiver\u001b[39m\u001b[38;5;124m\"\u001b[39m], envs, 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\u001b[1;32m    129\u001b[0m writer\u001b[38;5;241m.\u001b[39madd_scalar(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mcharts/learning_rate\u001b[39m\u001b[38;5;124m\"\u001b[39m, sender_optimizer\u001b[38;5;241m.\u001b[39mparam_groups[\u001b[38;5;241m0\u001b[39m][\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mlr\u001b[39m\u001b[38;5;124m\"\u001b[39m], global_step)\n",
      "Cell \u001b[0;32mIn[14], line 11\u001b[0m, in \u001b[0;36mupdate_agent\u001b[0;34m(agent, buffer, optimizer, next_obs, next_done, envs, 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)\u001b[0m\n\u001b[1;32m      9\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m t \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mreversed\u001b[39m(\u001b[38;5;28mrange\u001b[39m(num_steps)):\n\u001b[1;32m     10\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m t \u001b[38;5;241m==\u001b[39m num_steps \u001b[38;5;241m-\u001b[39m \u001b[38;5;241m1\u001b[39m:\n\u001b[0;32m---> 11\u001b[0m         nextnonterminal \u001b[38;5;241m=\u001b[39m \u001b[38;5;241;43m1.0\u001b[39;49m\u001b[43m \u001b[49m\u001b[38;5;241;43m-\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mnext_done\u001b[49m\n\u001b[1;32m     12\u001b[0m         nextvalues \u001b[38;5;241m=\u001b[39m next_value\n\u001b[1;32m     13\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n",
      "\u001b[0;31mTypeError\u001b[0m: unsupported operand type(s) for -: 'float' and 'list'"
     ]
    },
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mDer Kernel ist beim Ausführen von Code in der aktuellen Zelle oder einer vorherigen Zelle abgestürzt. \n",
      "\u001b[1;31mBitte überprüfen Sie den Code in der/den Zelle(n), um eine mögliche Fehlerursache zu identifizieren. \n",
      "\u001b[1;31mKlicken Sie <a href='https://aka.ms/vscodeJupyterKernelCrash'>hier</a>, um weitere Informationen zu erhalten. \n",
      "\u001b[1;31mWeitere Informationen finden Sie unter Jupyter <a href='command:jupyter.viewOutput'>Protokoll</a>."
     ]
    }
   ],
   "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",
    "next_obs, infos = get_observations_reset(envs, device)\n",
    "\n",
    "sender, sender_optimizer = initialize_agent(envs[0], device, learning_rate)\n",
    "receiver, receiver_optimizer = initialize_agent(envs[0], device, learning_rate)\n",
    "\n",
    "buffer_sender = Buffer(num_steps, envs[0], num_envs, device)\n",
    "buffer_receiver = Buffer(num_steps, envs[0], num_envs, device)\n",
    "\n",
    "global_step = 0\n",
    "\n",
    "next_done = {\"sender\": np.zeros(num_envs), \"receiver\": np.zeros(num_envs)}\n",
    "\n",
    "num_updates = total_timesteps // batch_size\n",
    "train_start = time.time()\n",
    "\n",
    "epoch_lengths = []\n",
    "current_length = 0\n",
    "start_time = time.time()\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",
    "        buffer_sender.dones[step] = torch.tensor(next_done[\"sender\"])\n",
    "        buffer_receiver.dones[step] = torch.tensor(next_done[\"receiver\"])\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",
    "        actions = {\"sender\": sender_action.cpu().numpy(), \"receiver\": receiver_action.cpu().numpy()}\n",
    "        next_obs, reward, terminations, truncations, infos, new_actions = get_data_step(envs, device, actions)\n",
    "\n",
    "        next_obs = {key: torch.Tensor(new_actions[key]).to(device) for key in new_actions}\n",
    "\n",
    "        current_dones = {key: terminations.get(key, False) or truncations.get(key, False) for key in set(terminations) | set(truncations)}\n",
    "\n",
    "        current_rewards[\"sender\"].append(reward[\"sender\"])\n",
    "        current_rewards[\"receiver\"].append(reward[\"receiver\"])\n",
    "\n",
    "        buffer_sender.rewards[step] = torch.tensor(np.array(reward[\"sender\"]).flatten()).to(device)\n",
    "        buffer_receiver.rewards[step] = torch.tensor(np.array(reward[\"receiver\"]).flatten()).to(device)\n",
    "        \n",
    "        if any(current_dones[\"sender\"]) or any(current_dones[\"receiver\"]):\n",
    "            indizes = np.array(flatten_list(current_dones))\n",
    "            true_indices = np.nonzero(indizes)[0]\n",
    "\n",
    "            for index in true_indices:\n",
    "                env = envs[int(index / 2)]\n",
    "                epoch_lengths.append(env.env.env.timestep)\n",
    "                obs, info = env.reset()\n",
    "                next_obs[\"sender\"][index] = torch.Tensor(obs[\"sender\"]).to(device)\n",
    "                next_obs[\"receiver\"][index] = torch.Tensor(obs[\"receiver\"]).to(device)\n",
    "                \n",
    "                epoch_rewards[\"sender\"] += sum(current_rewards[\"sender\"][index])\n",
    "                epoch_rewards[\"receiver\"] += sum(current_rewards[\"receiver\"][index])\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",
    "                print(\"SPS:\", int(global_step / (time.time() - start_time)), \"Average Reward:\", epoch_rewards[\"sender\"] / epoch_runs, \"Average Lengtht:\", sum(epoch_lengths[-100:]) / min(len(epoch_runs), 100))\n",
    "        \n",
    "        next_done = current_dones\n",
    "\n",
    "    update_agent(sender, buffer_sender, sender_optimizer, next_obs[\"sender\"], next_done[\"sender\"], envs, 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\"], envs, 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",
    "    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": null,
   "metadata": {},
   "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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Test_Communication:\n",
    "    def __init__(self, environment):\n",
    "        self.environment = environment\n",
    "        self.observation_space = {\"low\": [0, 0, 0, 0], \"high\": [1, 1, 1, 1]}\n",
    "        self.action_space = {\"low\": [0, 0, 0, 0], \"high\": [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]\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]\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]\n",
    "        else:\n",
    "            print(\"Dafaq is going on here?\")\n",
    "        return 0, observation, False, {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/cowolff/miniconda3/envs/Ray/lib/python3.9/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",
      "/Users/cowolff/miniconda3/envs/Ray/lib/python3.9/site-packages/keras/src/engine/training_v1.py:2359: UserWarning: `Model.state_updates` will be removed in a future version. This property should not be used in TensorFlow 2.0, as `updates` are applied automatically.\n",
      "  updates=self.state_updates,\n"
     ]
    }
   ],
   "source": [
    "xml_files = [\"levels/\" + file for file in os.listdir(\"levels/\")]\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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "epochs = 30\n",
    "num_steps = 1024\n",
    "lengths = []\n",
    "\n",
    "for epoch in 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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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": {
  "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}