Tianshou_Training_GoDot_PPO_AEC.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import datetime\n",
    "from typing import Optional, Tuple\n",
    "import json\n",
    "import numpy as np\n",
    "import torch\n",
    "\n",
    "from tianshou.data import Collector, VectorReplayBuffer, PrioritizedVectorReplayBuffer\n",
    "from tianshou.env import DummyVectorEnv\n",
    "from tianshou.env.pettingzoo_env_parallel import PettingZooParallelEnv\n",
    "from tianshou.env.pettingzoo_env import PettingZooEnv\n",
    "from pettingzoo.utils import parallel_to_aec\n",
    "\n",
    "#from PettingZooParallelEnv import PettingZooParallelEnv\n",
    "\n",
    "\n",
    "from tianshou.policy import PPOPolicy\n",
    "from tianshou.trainer import OnpolicyTrainer\n",
    "\n",
    "from tianshou.utils.net.common import ActorCritic, DataParallelNet, Net\n",
    "from tianshou.utils.net.discrete import Actor, Critic\n",
    "\n",
    "from tianshou.policy import BasePolicy, DQNPolicy, MultiAgentPolicyManager, RandomPolicy, RainbowPolicy\n",
    "from tianshou.trainer import OffpolicyTrainer\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "from DNN_B_ACE import DNN_B_ACE\n",
    "from GodotRLPettingZooWrapper import GodotRLPettingZooWrapper\n",
    "\n",
    "from CollectorMA import CollectorMA\n",
    "from MAParalellPolicy import MAParalellPolicy\n",
    "\n",
    "\n",
    "####---------------------------#######\n",
    "#Tianshou Adjustment\n",
    "import wandb\n",
    "os.environ[\"WANDB_NOTEBOOK_NAME\"] = \"Tianshow_Training_GoDot.ipybn\"\n",
    "from tianshou.utils import WandbLogger\n",
    "from tianshou.utils.logger.base import LOG_DATA_TYPE\n",
    "def new_write(self, step_type: str, step: int, data: LOG_DATA_TYPE) -> None:\n",
    "    data[step_type] = step\n",
    "    wandb.log(data)   \n",
    "WandbLogger.write = new_write \n",
    "####---------------------------#######\n",
    "\n",
    "\n",
    "model  =  \"PPO_DNN\"#\"SISL_Task_MultiHead\" #\"CNN_ATT_SISL\" #\"MultiHead_SISL\" \n",
    "test_num  =  \"_B_ACE01\"\n",
    "policyModel  =  \"PPO\"\n",
    "name = model + test_num\n",
    "\n",
    "train_env_num = 3\n",
    "test_env_num = 3\n",
    "\n",
    "now = datetime.datetime.now().strftime(\"%y%m%d-%H%M%S\")\n",
    "log_name = name + str(now)\n",
    "log_path = os.path.join('./', \"Logs\", \"dqn_sisl\", log_name)\n",
    "\n",
    "load_policy_name = f'policy_SISL_Task_MultiHead_Desk_NewExpCor231219-173711_44.pth'\n",
    "save_policy_name = f'policy_{log_name}'\n",
    "policy_path = \"ppo_B_ACE\"\n",
    "\n",
    "\n",
    "model_load_path = os.path.join(policy_path, load_policy_name)  \n",
    "model_save_path = os.path.join(policy_path, save_policy_name)        \n",
    "os.makedirs(os.path.join(policy_path), exist_ok=True)  \n",
    "os.makedirs(os.path.join(log_path), exist_ok=True)\n",
    "\n",
    "Policy_Config = {\n",
    "    \"same_policy\" : True,\n",
    "    \"load_model\" : False,\n",
    "    \"freeze_CNN\" : False     \n",
    "                }\n",
    "\n",
    "B_ACE_Config = {\n",
    "    \"max_cycles\" : 200\n",
    "}\n",
    "\n",
    "max_cycles = B_ACE_Config[\"max_cycles\"]\n",
    "n_agents = 4#B_ACE_Config[\"n_pursuers\"]\n",
    "\n",
    "dqn_params = {\"discount_factor\": 0.98, \n",
    "              \"estimation_step\": 20, \n",
    "              \"target_update_freq\": 1000,#max_cycles * n_agents,\n",
    "              \"optminizer\": \"Adam\",\n",
    "              \"lr\": 0.00016 }\n",
    "\n",
    "trainer_params = {\"max_epoch\": 500,\n",
    "                  \"step_per_epoch\": 20000,#5 * (150 * n_agents),\n",
    "                  \"step_per_collect\": 400,# * (10 * n_agents),\n",
    "                  \"episode_per_test\": 30,\n",
    "                  \"batch_size\" : 32 * n_agents,\n",
    "                  \"update_per_step\": 1 / 50, #Only run after close a Collect (run many times as necessary to meet the value)\n",
    "                  \"tn_eps_max\": 0.10,\n",
    "                  \"ts_eps_max\": 0.01,\n",
    "                  \"warmup_size\" : 1\n",
    "                  }\n",
    "\n",
    "runConfig = dqn_params\n",
    "runConfig.update(Policy_Config)\n",
    "runConfig.update(trainer_params) \n",
    "runConfig.update(B_ACE_Config)\n",
    "\n",
    "\n",
    "def _get_agents(\n",
    "    agent_learn: Optional[BasePolicy] = None,\n",
    "    agent_opponent: Optional[BasePolicy] = None,\n",
    "    optim: Optional[torch.optim.Optimizer] = None,\n",
    "    policy_load_path = None,\n",
    ") -> Tuple[BasePolicy, torch.optim.Optimizer, list]:\n",
    "    \n",
    "    env = _get_env()       \n",
    "    agent_observation_space = env.observation_space\n",
    "   \n",
    "    print(env.action_space)\n",
    "    action_shape = env.action_space.shape\n",
    "    \n",
    "    \n",
    "    device=\"cuda\" if torch.cuda.is_available() else \"cpu\"  \n",
    "\n",
    "    agents = []        \n",
    "    \n",
    "    if Policy_Config[\"same_policy\"]:\n",
    "        policies_number = 1\n",
    "    else:\n",
    "        policies_number = 4#len(env.agents)\n",
    "\n",
    "    for _ in range(policies_number):      \n",
    "        \n",
    "        if model == \"DNN_B_ACE\":\n",
    "            net = DNN_B_ACE(\n",
    "                obs_shape=agent_observation_space,                \n",
    "                action_shape=4,                \n",
    "                device=\"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "                \n",
    "            ).to(device)        \n",
    "                   \n",
    "        if policyModel == \"DQN\":\n",
    "            agent_learn = DQNPolicy(\n",
    "                model=net,\n",
    "                optim=optim,\n",
    "                action_space = action_shape,\n",
    "                discount_factor= dqn_params[\"discount_factor\"],\n",
    "                estimation_step=dqn_params[\"estimation_step\"],\n",
    "                target_update_freq=dqn_params[\"target_update_freq\"],\n",
    "                reward_normalization = False,\n",
    "                clip_loss_grad = False \n",
    "            ) \n",
    "         \n",
    "         \n",
    "        if model == \"PPO_DNN\":\n",
    "            \n",
    "            netActor = DNN_B_ACE(\n",
    "                obs_shape=agent_observation_space.shape,                \n",
    "                action_shape=4,                \n",
    "                device=\"cuda\" if torch.cuda.is_available() else \"cpu\"                \n",
    "            ).to(device)\n",
    "\n",
    "            netCritic = DNN_B_ACE(\n",
    "                obs_shape=agent_observation_space.shape,                \n",
    "                action_shape=4,                \n",
    "                device=\"cuda\" if torch.cuda.is_available() else \"cpu\"                \n",
    "            ).to(device)\n",
    "            \n",
    "                        \n",
    "            if False:#torch.cuda.is_available():\n",
    "                actor = DataParallelNet(Actor(netActor, action_shape, device=None).to(device))\n",
    "                critic = DataParallelNet(Critic(netCritic, device=None).to(device))\n",
    "            else:\n",
    "                actor = Actor(netActor, action_shape, device=device).to(device)\n",
    "                critic = Critic(netCritic, device=device).to(device)\n",
    "            \n",
    "            actor_critic = ActorCritic(actor, critic)\n",
    "        \n",
    "            # orthogonal initialization\n",
    "            # for m in actor_critic.modules():\n",
    "            #     if isinstance(m, torch.nn.Linear):\n",
    "            #         torch.nn.init.orthogonal_(m.weight)\n",
    "            #         torch.nn.init.zeros_(m.bias)\n",
    "            \n",
    "            \n",
    "            dist = torch.distributions.Categorical         \n",
    "                \n",
    "            #optim_actor  = torch.optim.Adam(netActor.parameters(),  lr=dqn_params[\"lr\"], weight_decay=0.0, amsgrad= True )\n",
    "            #optim_critic = torch.optim.Adam(netCritic.parameters(), lr=dqn_params[\"lr\"], weight_decay=0.0, amsgrad= True )\n",
    "            optim = torch.optim.Adam(actor_critic.parameters(), lr=dqn_params[\"lr\"])\n",
    "                    \n",
    "            agent_learn = PPOPolicy(\n",
    "                actor=actor,\n",
    "                critic=critic,\n",
    "                optim=optim,\n",
    "                dist_fn=dist,\n",
    "                # action_scaling=isinstance(env.action_space, Box),\n",
    "                discount_factor=0.99,\n",
    "                max_grad_norm=0.5,\n",
    "                eps_clip=0.2,\n",
    "                vf_coef=0.5,\n",
    "                ent_coef=0.0,\n",
    "                gae_lambda=0.95,\n",
    "                reward_normalization=0,\n",
    "                dual_clip=None,\n",
    "                value_clip=0,\n",
    "                action_space=env.action_space,\n",
    "                deterministic_eval=True,\n",
    "                advantage_normalization=0,\n",
    "                recompute_advantage=0,\n",
    "            )\n",
    "            \n",
    "            if Policy_Config[\"load_model\"] is True:\n",
    "                # Load the saved checkpoint             \n",
    "                agent_learn.load_state_dict(torch.load(model_load_path))\n",
    "                print(f'Loaded-> {model_load_path}')\n",
    "                   \n",
    "        \n",
    "        agents.append(agent_learn)\n",
    "\n",
    "    if Policy_Config[\"same_policy\"]:\n",
    "        agents = [agents[0] for _ in range(len(env.agents))]\n",
    "    else:\n",
    "        for _ in range(len(env.agents) - policies_number):\n",
    "            agents.append(agents[0])\n",
    "\n",
    "    # policy = MultiAgentPolicyManager(policies = agents, env=env)  \n",
    "    policy = MultiAgentPolicyManager(policies = agents, env=env )  \n",
    "        \n",
    "    return policy, optim, env.agents\n",
    "\n",
    "def _get_env():\n",
    "    \"\"\"This function is needed to provide callables for DummyVectorEnv.\"\"\"   \n",
    "    \n",
    "    # env_paralell = MultiUAVEnv()  \n",
    "    # env = pursuit_v4.env()   \n",
    "    env_path = \"BVR_AirCombat/bin/FlyBy.exe\"\n",
    "    env = GodotRLPettingZooWrapper( env_path, 2)\n",
    "    env = parallel_to_aec(env)\n",
    "    \n",
    "    env= PettingZooEnv(env)\n",
    "    \n",
    "    #env = PettingZooParallelEnv(env)  \n",
    "    \n",
    "    return env  \n",
    "   \n",
    "\n",
    "# print(json.dumps(runConfig, indent=4))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "waiting for remote GODOT connection on port 11008\n",
      "connection established\n",
      "action space {'act': {'size': [4], 'space': 'box'}}\n"
     ]
    },
    {
     "ename": "KeyError",
     "evalue": "'action_type'",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mKeyError\u001b[0m                                  Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[2], line 6\u001b[0m\n\u001b[0;32m      3\u001b[0m torch\u001b[38;5;241m.\u001b[39mset_grad_enabled(\u001b[38;5;28;01mTrue\u001b[39;00m) \n\u001b[0;32m      5\u001b[0m \u001b[38;5;66;03m# ======== Step 1: Environment setup =========\u001b[39;00m\n\u001b[1;32m----> 6\u001b[0m train_envs \u001b[38;5;241m=\u001b[39m \u001b[43mDummyVectorEnv\u001b[49m\u001b[43m(\u001b[49m\u001b[43m[\u001b[49m\u001b[43m_get_env\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mfor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43m_\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01min\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;28;43mrange\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mtrain_env_num\u001b[49m\u001b[43m)\u001b[49m\u001b[43m]\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m      7\u001b[0m test_envs \u001b[38;5;241m=\u001b[39m DummyVectorEnv([_get_env \u001b[38;5;28;01mfor\u001b[39;00m _ \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(test_env_num)]) \n\u001b[0;32m      9\u001b[0m \u001b[38;5;66;03m# seed\u001b[39;00m\n",
      "File \u001b[1;32md:\\Python310\\lib\\site-packages\\tianshou\\env\\venvs.py:444\u001b[0m, in \u001b[0;36mDummyVectorEnv.__init__\u001b[1;34m(self, env_fns, **kwargs)\u001b[0m\n\u001b[0;32m    443\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__init__\u001b[39m(\u001b[38;5;28mself\u001b[39m, env_fns: List[Callable[[], ENV_TYPE]], \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs: Any) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m--> 444\u001b[0m     \u001b[38;5;28msuper\u001b[39m()\u001b[38;5;241m.\u001b[39m\u001b[38;5;21m__init__\u001b[39m(env_fns, DummyEnvWorker, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n",
      "File \u001b[1;32md:\\Python310\\lib\\site-packages\\tianshou\\env\\venvs.py:150\u001b[0m, in \u001b[0;36mBaseVectorEnv.__init__\u001b[1;34m(self, env_fns, worker_fn, wait_num, timeout)\u001b[0m\n\u001b[0;32m    147\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_env_fns \u001b[38;5;241m=\u001b[39m env_fns\n\u001b[0;32m    148\u001b[0m \u001b[38;5;66;03m# A VectorEnv contains a pool of EnvWorkers, which corresponds to\u001b[39;00m\n\u001b[0;32m    149\u001b[0m \u001b[38;5;66;03m# interact with the given envs (one worker <-> one env).\u001b[39;00m\n\u001b[1;32m--> 150\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mworkers \u001b[38;5;241m=\u001b[39m [worker_fn(_patch_env_generator(fn)) \u001b[38;5;28;01mfor\u001b[39;00m fn \u001b[38;5;129;01min\u001b[39;00m env_fns]\n\u001b[0;32m    151\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mworker_class \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mtype\u001b[39m(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mworkers[\u001b[38;5;241m0\u001b[39m])\n\u001b[0;32m    152\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m \u001b[38;5;28missubclass\u001b[39m(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mworker_class, EnvWorker)\n",
      "File \u001b[1;32md:\\Python310\\lib\\site-packages\\tianshou\\env\\venvs.py:150\u001b[0m, in \u001b[0;36m<listcomp>\u001b[1;34m(.0)\u001b[0m\n\u001b[0;32m    147\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_env_fns \u001b[38;5;241m=\u001b[39m env_fns\n\u001b[0;32m    148\u001b[0m \u001b[38;5;66;03m# A VectorEnv contains a pool of EnvWorkers, which corresponds to\u001b[39;00m\n\u001b[0;32m    149\u001b[0m \u001b[38;5;66;03m# interact with the given envs (one worker <-> one env).\u001b[39;00m\n\u001b[1;32m--> 150\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mworkers \u001b[38;5;241m=\u001b[39m [\u001b[43mworker_fn\u001b[49m\u001b[43m(\u001b[49m\u001b[43m_patch_env_generator\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfn\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m \u001b[38;5;28;01mfor\u001b[39;00m fn \u001b[38;5;129;01min\u001b[39;00m env_fns]\n\u001b[0;32m    151\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mworker_class \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mtype\u001b[39m(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mworkers[\u001b[38;5;241m0\u001b[39m])\n\u001b[0;32m    152\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m \u001b[38;5;28missubclass\u001b[39m(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mworker_class, EnvWorker)\n",
      "File \u001b[1;32md:\\Python310\\lib\\site-packages\\tianshou\\env\\worker\\dummy.py:13\u001b[0m, in \u001b[0;36mDummyEnvWorker.__init__\u001b[1;34m(self, env_fn)\u001b[0m\n\u001b[0;32m     12\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__init__\u001b[39m(\u001b[38;5;28mself\u001b[39m, env_fn: Callable[[], gym\u001b[38;5;241m.\u001b[39mEnv]) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m---> 13\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39menv \u001b[38;5;241m=\u001b[39m \u001b[43menv_fn\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m     14\u001b[0m     \u001b[38;5;28msuper\u001b[39m()\u001b[38;5;241m.\u001b[39m\u001b[38;5;21m__init__\u001b[39m(env_fn)\n",
      "File \u001b[1;32md:\\Python310\\lib\\site-packages\\tianshou\\env\\venvs.py:45\u001b[0m, in \u001b[0;36m_patch_env_generator.<locals>.patched\u001b[1;34m()\u001b[0m\n\u001b[0;32m     40\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mpatched\u001b[39m() \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m gym\u001b[38;5;241m.\u001b[39mEnv:\n\u001b[0;32m     41\u001b[0m     \u001b[38;5;28;01massert\u001b[39;00m \u001b[38;5;28mcallable\u001b[39m(\n\u001b[0;32m     42\u001b[0m         fn\n\u001b[0;32m     43\u001b[0m     ), \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mEnv generators that are provided to vector environemnts must be callable.\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m---> 45\u001b[0m     env \u001b[38;5;241m=\u001b[39m \u001b[43mfn\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m     46\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(env, (gym\u001b[38;5;241m.\u001b[39mEnv, PettingZooEnv, PettingZooParallelEnv)):\n\u001b[0;32m     47\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m env\n",
      "Cell \u001b[1;32mIn[1], line 233\u001b[0m, in \u001b[0;36m_get_env\u001b[1;34m()\u001b[0m\n\u001b[0;32m    230\u001b[0m \u001b[38;5;66;03m# env_paralell = MultiUAVEnv()  \u001b[39;00m\n\u001b[0;32m    231\u001b[0m \u001b[38;5;66;03m# env = pursuit_v4.env()   \u001b[39;00m\n\u001b[0;32m    232\u001b[0m env_path \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mBVR_AirCombat/bin/FlyBy.exe\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m--> 233\u001b[0m env \u001b[38;5;241m=\u001b[39m \u001b[43mGodotRLPettingZooWrapper\u001b[49m\u001b[43m(\u001b[49m\u001b[43m \u001b[49m\u001b[43menv_path\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m2\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[0;32m    234\u001b[0m env \u001b[38;5;241m=\u001b[39m parallel_to_aec(env)\n\u001b[0;32m    236\u001b[0m env\u001b[38;5;241m=\u001b[39m PettingZooEnv(env)\n",
      "File \u001b[1;32md:\\GoDot\\godot_rl_envs-1\\GodotRLPettingZooWrapper.py:12\u001b[0m, in \u001b[0;36mGodotRLPettingZooWrapper.__init__\u001b[1;34m(self, env_path, num_agents, **kwargs)\u001b[0m\n\u001b[0;32m     10\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__init__\u001b[39m(\u001b[38;5;28mself\u001b[39m, env_path, num_agents, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs):\n\u001b[1;32m---> 12\u001b[0m     \u001b[38;5;28msuper\u001b[39m()\u001b[38;5;241m.\u001b[39m\u001b[38;5;21m__init__\u001b[39m(env_path\u001b[38;5;241m=\u001b[39menv_path, convert_action_space\u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[0;32m     14\u001b[0m     \u001b[38;5;66;03m# self.num_agents = num_agents\u001b[39;00m\n\u001b[0;32m     15\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39magents \u001b[38;5;241m=\u001b[39m [\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124magent_\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mi\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m'\u001b[39m \u001b[38;5;28;01mfor\u001b[39;00m i \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(num_agents)]  \u001b[38;5;66;03m# Initialize agents\u001b[39;00m\n",
      "File \u001b[1;32md:\\Python310\\lib\\site-packages\\godot_rl\\core\\godot_env.py:62\u001b[0m, in \u001b[0;36mGodotEnv.__init__\u001b[1;34m(self, env_path, port, show_window, seed, framerate, action_repeat, speedup, convert_action_space)\u001b[0m\n\u001b[0;32m     60\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mnum_envs \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[0;32m     61\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_handshake()\n\u001b[1;32m---> 62\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_get_env_info\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m     63\u001b[0m \u001b[38;5;66;03m# sf2 requires a tuple action space\u001b[39;00m\n\u001b[0;32m     64\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_tuple_action_space \u001b[38;5;241m=\u001b[39m spaces\u001b[38;5;241m.\u001b[39mTuple([v \u001b[38;5;28;01mfor\u001b[39;00m _, v \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_action_space\u001b[38;5;241m.\u001b[39mitems()])\n",
      "File \u001b[1;32md:\\Python310\\lib\\site-packages\\godot_rl\\core\\godot_env.py:326\u001b[0m, in \u001b[0;36mGodotEnv._get_env_info\u001b[1;34m(self)\u001b[0m\n\u001b[0;32m    324\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124maction space\u001b[39m\u001b[38;5;124m\"\u001b[39m, json_dict[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124maction_space\u001b[39m\u001b[38;5;124m\"\u001b[39m])\n\u001b[0;32m    325\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m k, v \u001b[38;5;129;01min\u001b[39;00m json_dict[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124maction_space\u001b[39m\u001b[38;5;124m\"\u001b[39m]\u001b[38;5;241m.\u001b[39mitems():\n\u001b[1;32m--> 326\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[43mv\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43maction_type\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m]\u001b[49m \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mdiscrete\u001b[39m\u001b[38;5;124m\"\u001b[39m:\n\u001b[0;32m    327\u001b[0m         action_spaces[k] \u001b[38;5;241m=\u001b[39m spaces\u001b[38;5;241m.\u001b[39mDiscrete(v[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msize\u001b[39m\u001b[38;5;124m\"\u001b[39m])\n\u001b[0;32m    328\u001b[0m     \u001b[38;5;28;01melif\u001b[39;00m v[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124maction_type\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mcontinuous\u001b[39m\u001b[38;5;124m\"\u001b[39m:\n",
      "\u001b[1;31mKeyError\u001b[0m: 'action_type'"
     ]
    }
   ],
   "source": [
    "if __name__ == \"__main__\":\n",
    "                        \n",
    "    torch.set_grad_enabled(True) \n",
    "   \n",
    "    # ======== Step 1: Environment setup =========\n",
    "    train_envs = DummyVectorEnv([_get_env for _ in range(train_env_num)])\n",
    "    test_envs = DummyVectorEnv([_get_env for _ in range(test_env_num)]) \n",
    "\n",
    "    # seed\n",
    "    seed = 100\n",
    "    np.random.seed(seed)\n",
    "    \n",
    "    torch.manual_seed(seed)\n",
    "\n",
    "    train_envs.seed(seed)\n",
    "    test_envs.seed(seed)\n",
    "\n",
    "    # ======== Step 2: Agent setup =========\n",
    "    policy, optim, agents = _get_agents()    \n",
    "\n",
    "    \n",
    "    if True:\n",
    "        # ======== Step 3: Collector setup =========\n",
    "        train_collector = Collector(\n",
    "            policy,\n",
    "            train_envs,\n",
    "            # VectorReplayBuffer(300_000, len(train_envs)),\n",
    "            PrioritizedVectorReplayBuffer( 300_000, len(train_envs), alpha=0.6, beta=0.4) , \n",
    "            #ListReplayBuffer(100000)       \n",
    "            # buffer = StateMemoryVectorReplayBuffer(\n",
    "            #         300_000,\n",
    "            #         len(train_envs),  # Assuming train_envs is your vectorized environment\n",
    "            #         memory_size=10,                \n",
    "            #     ),\n",
    "            exploration_noise=True             \n",
    "        )\n",
    "        test_collector = Collector(policy, test_envs, exploration_noise=True)\n",
    "        \n",
    "\n",
    "    if False:\n",
    "        agents_buffers_training = {agent : \n",
    "                        PrioritizedVectorReplayBuffer( 30_000, \n",
    "                                                        len(train_envs), \n",
    "                                                        alpha=0.6, \n",
    "                                                        beta=0.4) \n",
    "                                                        for agent in agents\n",
    "                        }\n",
    "        agents_buffers_test = {agent : \n",
    "                        PrioritizedVectorReplayBuffer( 30_000, \n",
    "                                                        len(train_envs), \n",
    "                                                        alpha=0.6, \n",
    "                                                        beta=0.4) \n",
    "                                                        for agent in agents\n",
    "                        }\n",
    "    \n",
    "        # ======== Step 3: Collector setup =========\n",
    "        train_collector = CollectorMA(\n",
    "            policy,\n",
    "            train_envs,\n",
    "            agents_buffers_training,                        \n",
    "            exploration_noise=True             \n",
    "        )\n",
    "        test_collector = CollectorMA(policy, test_envs, agents_buffers_test, exploration_noise=True)\n",
    "\n",
    "    \n",
    "        \n",
    "    print(\"Buffer Warming Up \")    \n",
    "    for i in range(trainer_params[\"warmup_size\"]):#int(trainer_params['batch_size'] / (300 * 10 ) )):\n",
    "        \n",
    "         train_collector.collect(n_episode=train_env_num)#,random=True) #trainer_params['batch_size'] * train_env_num))\n",
    "         #train_collector.collect(n_step=300 * 10)\n",
    "         print(\".\", end=\"\") \n",
    "    \n",
    "    # len_buffer = len(train_collector.buffer) / (B_ACE_Config[\"max_cycles\"] * SISL_Config[\"n_pursuers\"])\n",
    "    # print(\"\\nBuffer Lenght: \", len_buffer ) \n",
    "    len_buffer = 0\n",
    "    \n",
    "    info = { \"Buffer\"  : \"PriorizedReplayBuffer\", \" Warmup_ep\" : len_buffer}\n",
    "    \n",
    "    # ======== tensorboard logging setup =========                       \n",
    "    logger = WandbLogger(\n",
    "        train_interval = runConfig[\"max_cycles\"] ,\n",
    "        test_interval = 1,#runConfig[\"max_cycles\"] * runConfig[\"n_pursuers\"],\n",
    "        update_interval = runConfig[\"max_cycles\"],\n",
    "        save_interval = 1,\n",
    "        write_flush = True,\n",
    "        project = \"B_ACE01\",\n",
    "        name = log_name,\n",
    "        entity = None,\n",
    "        run_id = log_name,\n",
    "        config = runConfig,\n",
    "        monitor_gym = True )\n",
    "    \n",
    "    writer = SummaryWriter(log_path)    \n",
    "    writer.add_text(\"args\", str(runConfig))    \n",
    "    logger.load(writer)\n",
    "\n",
    "    global_step_holder = [0] \n",
    "        \n",
    "    # ======== Step 4: Callback functions setup =========\n",
    "    def save_best_fn(policy):                \n",
    "        \n",
    "        if Policy_Config[\"same_policy\"]:\n",
    "            torch.save(policy.policies[agents[0]].state_dict(), model_save_path + \"_\" + str(global_step_holder[0]) + \"_BestRew.pth\")\n",
    "            print(\"Best Saved Rew\" , str(global_step_holder[0]))\n",
    "        \n",
    "        else:\n",
    "            for n,agent in enumerate(agents):\n",
    "                torch.save(policy.policies[agent].state_dict(), model_save_path + \"_\" + str(global_step_holder[0]) + \"_\" + agent + \".pth\")\n",
    "            \n",
    "            print(\"Bests Saved Rew\" , str(global_step_holder[0]))\n",
    "        \n",
    "    def save_test_best_fn(policy):                \n",
    "        \n",
    "        if Policy_Config[\"same_policy\"]:\n",
    "            torch.save(policy.policies[agents[0]].state_dict(), model_save_path + \"_\" + str(global_step_holder[0]) + \"_BestLen.pth\")\n",
    "            print(\"Best Saved Length\" , str(global_step_holder[0]))\n",
    "        \n",
    "        else:\n",
    "            for n,agent in enumerate(agents):\n",
    "                torch.save(policy.policies[agent].state_dict(), model_save_path + \"_\" + str(global_step_holder[0]) + \"_\" + agent + \".pth\")\n",
    "            \n",
    "            print(\"Best Saved Length\" , str(global_step_holder[0]))\n",
    "        \n",
    "\n",
    "    def stop_fn(mean_rewards):\n",
    "        return mean_rewards >= 99999939.0\n",
    "\n",
    "    def train_fn(epoch, env_step):\n",
    "        epsilon = trainer_params['tn_eps_max'] - (trainer_params['tn_eps_max'] - trainer_params['tn_eps_max']/100)*(epoch/trainer_params['max_epoch'])          \n",
    "        if Policy_Config[\"same_policy\"]:\n",
    "            policy.policies[agents[0]].set_eps(epsilon)\n",
    "        else:\n",
    "            for agent in agents:\n",
    "                policy.policies[agent].set_eps(epsilon)\n",
    "                \n",
    "        \n",
    "        # if env_step % 500 == 0:\n",
    "            # logger.write(\"train/env_step\", env_step, {\"train/eps\": eps})\n",
    "\n",
    "\n",
    "    def test_fn(epoch, env_step):\n",
    "               \n",
    "        epsilon = trainer_params['ts_eps_max']#0.01#max(0.001, 0.1 - epoch * 0.001)\n",
    "        if Policy_Config[\"same_policy\"]:\n",
    "            policy.policies[agents[0]].set_eps(epsilon)\n",
    "        else:            \n",
    "            for agent in agents:                             \n",
    "                 policy.policies[agent].set_eps(epsilon)\n",
    "                \n",
    "        \n",
    "        if global_step_holder[0] % 10 == 0:\n",
    "            \n",
    "            if Policy_Config[\"same_policy\"]:\n",
    "                torch.save(policy.policies[agents[0]].state_dict(), model_save_path + \"_\" + str(global_step_holder[0]) + \"_Step.pth\")\n",
    "                print(\"Steps Policy Saved \" , str(global_step_holder[0]))\n",
    "            \n",
    "            else:\n",
    "                for n,agent in enumerate(agents):\n",
    "                    torch.save(policy.policies[agent].state_dict(), model_save_path + \"_\" + str(global_step_holder[0]) + \"_\" + agent + \"Step\" + str(global_step_holder[0]) + \".pth\")\n",
    "                \n",
    "                print(\"Steps Policy Saved \" , str(global_step_holder[0]))\n",
    "\n",
    "        \n",
    "    def reward_metric(rews):       \n",
    "                \n",
    "        global_step_holder[0] +=1 \n",
    "        return np.sum(rews, axis = 1)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "    # # ======== Step 5: Run the trainer =========   \n",
    "    onPolicyTrainer = OnpolicyTrainer(\n",
    "        policy=policy,\n",
    "        train_collector=train_collector,\n",
    "        test_collector=test_collector,\n",
    "        max_epoch=trainer_params['max_epoch'],\n",
    "        step_per_epoch=trainer_params['step_per_epoch'],\n",
    "        repeat_per_collect=10, #TODO: understand\n",
    "        episode_per_test=trainer_params['episode_per_test'],\n",
    "        batch_size=trainer_params['batch_size'],\n",
    "        step_per_collect=trainer_params['step_per_collect'],\n",
    "        stop_fn=stop_fn,\n",
    "        save_best_fn=save_best_fn,\n",
    "        logger=logger,\n",
    "    )\n",
    "    \n",
    "\n",
    "    writer.close()\n",
    "    \n",
    "    # # ======== Step 5: Run the trainer =========\n",
    "    # offPolicyTrainer = OffpolicyTrainer(\n",
    "    #     policy=policy,\n",
    "    #     train_collector=train_collector,\n",
    "    #     test_collector=test_collector,        \n",
    "    #     max_epoch=trainer_params['max_epoch'],\n",
    "    #     step_per_epoch=trainer_params['step_per_epoch'],\n",
    "    #     step_per_collect=trainer_params['step_per_collect'],        \n",
    "    #     episode_per_test= trainer_params['episode_per_test'],\n",
    "    #     batch_size=trainer_params['batch_size'],\n",
    "    #     train_fn=train_fn,\n",
    "    #     test_fn=test_fn,\n",
    "    #     stop_fn=stop_fn,\n",
    "    #     save_best_fn=save_best_fn,\n",
    "    #     # save_test_best_fn=save_test_best_fn,\n",
    "    #     update_per_step=trainer_params['update_per_step'],\n",
    "    #     logger=logger,\n",
    "    #     test_in_train=True,\n",
    "    #     reward_metric=reward_metric,\n",
    "    #     show_progress = True \n",
    "               \n",
    "    #     )\n",
    "    \n",
    "    result = onPolicyTrainer.run()\n",
    "    writer.close()\n",
    "    # return result, policy.policies[agents[1]]\n",
    "    print(f\"\\n==========Result==========\\n{result}\")\n",
    "    print(\"\\n(the trained policy can be accessed via policy.policies[agents[0]])\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}