[
  {
    "path": ".gitignore",
    "content": "**.tfevents.**\n\n# Byte-compiled / optimized / DLL files\n__pycache__/\n*.py[cod]\n*$py.class\n\n# C extensions\n*.so\n\n# Distribution / packaging\n.Python\nbuild/\ndevelop-eggs/\ndist/\ndownloads/\neggs/\n.eggs/\nlib/\nlib64/\nparts/\nsdist/\nvar/\nwheels/\npip-wheel-metadata/\nshare/python-wheels/\n*.egg-info/\n.installed.cfg\n*.egg\nMANIFEST\n\n# PyInstaller\n#  Usually these files are written by a python script from a template\n#  before PyInstaller builds the exe, so as to inject date/other infos into it.\n*.manifest\n*.spec\n\n# Installer logs\npip-log.txt\npip-delete-this-directory.txt\n\n# Unit test / coverage reports\nhtmlcov/\n.tox/\n.nox/\n.coverage\n.coverage.*\n.cache\nnosetests.xml\ncoverage.xml\n*.cover\n*.py,cover\n.hypothesis/\n.pytest_cache/\n\n# Translations\n*.mo\n*.pot\n\n# Django stuff:\n*.log\nlocal_settings.py\ndb.sqlite3\ndb.sqlite3-journal\n\n# Flask stuff:\ninstance/\n.webassets-cache\n\n# Scrapy stuff:\n.scrapy\n\n# Sphinx documentation\ndocs/_build/\n\n# PyBuilder\ntarget/\n\n# Jupyter Notebook\n.ipynb_checkpoints\n\n# IPython\nprofile_default/\nipython_config.py\n\n# pipenv\n#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.\n#   However, in case of collaboration, if having platform-specific dependencies or dependencies\n#   having no cross-platform support, pipenv may install dependencies that don't work, or not\n#   install all needed dependencies.\n#Pipfile.lock\n\n# PEP 582; used by e.g. github.com/David-OConnor/pyflow\n__pypackages__/\n\n# Celery stuff\ncelerybeat-schedule\ncelerybeat.pid\n\n# SageMath parsed files\n*.sage.py\n\n# Environments\n.env\n.venv\nenv/\nvenv/\nENV/\nenv.bak/\nvenv.bak/\n\n# Spyder project settings\n.spyderproject\n.spyproject\n\n# Rope project settings\n.ropeproject\n\n# mkdocs documentation\n/site\n\n# mypy\n.mypy_cache/\n.dmypy.json\ndmypy.json\n\n# Pyre type checker\n.pyre/\n"
  },
  {
    "path": ".python-version",
    "content": "3.9.5/envs/invalid-action-masking"
  },
  {
    "path": "LICENSE",
    "content": "MIT License\n\nCopyright (c) 2020 neurips2020submission\n\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in all\ncopies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\nSOFTWARE.\n"
  },
  {
    "path": "README.MD",
    "content": "# A Closer Look at Invalid Action Masking in Policy Gradient Algorithms\n\nThis repo contains the source code to reproduce the results in the paper [*A Closer Look at Invalid Action Masking in Policy Gradient Algorithms*](https://arxiv.org/abs/2006.14171). \n\n## Get started\n\nIf you have pyenv or poetry:\n```bash\npoetry install\n\nrm -rf ~/microrts && mkdir ~/microrts && \\\n    wget -O ~/microrts/microrts.zip http://microrts.s3.amazonaws.com/microrts/artifacts/202004222224.microrts.zip && \\\n    unzip ~/microrts/microrts.zip -d ~/microrts/ && \\\n    rm ~/microrts/microrts.zip\n```\n\nElse, you can also install dependencies via `pip install -r requirements.txt`.\n\n## 10x10 Experiments\n```\npoetry run python invalid_action_masking/ppo_10x10.py\npoetry run python invalid_action_masking/ppo_no_adj_10x10.py\npoetry run python invalid_action_masking/ppo_no_mask_10x10.py\npoetry run python ppo.py # newer & recommended PPO implementation that matches implementation details in `openai/baselines`\n```\n\n## Citation\n\n```bibtex\n@inproceedings{huang2020closer,\n  author    = {Shengyi Huang and\n               Santiago Onta{\\~{n}}{\\'{o}}n},\n  editor    = {Roman Bart{\\'{a}}k and\n               Fazel Keshtkar and\n               Michael Franklin},\n  title     = {A Closer Look at Invalid Action Masking in Policy Gradient Algorithms},\n  booktitle = {Proceedings of the Thirty-Fifth International Florida Artificial Intelligence\n               Research Society Conference, {FLAIRS} 2022, Hutchinson Island, Jensen\n               Beach, Florida, USA, May 15-18, 2022},\n  year      = {2022},\n  url       = {https://doi.org/10.32473/flairs.v35i.130584},\n  doi       = {10.32473/flairs.v35i.130584},\n  timestamp = {Thu, 09 Jun 2022 16:44:11 +0200},\n  biburl    = {https://dblp.org/rec/conf/flairs/HuangO22.bib},\n  bibsource = {dblp computer science bibliography, https://dblp.org}\n}\n```\n"
  },
  {
    "path": "build.sh",
    "content": "docker build -t invalid_action_masking:latest -f sharedmemory.Dockerfile .\n"
  },
  {
    "path": "gym_vec_api/ppo_multidiscrete.py",
    "content": "import argparse\nimport os\nimport random\nimport time\nfrom distutils.util import strtobool\n\nimport gym\nimport gym_microrts # fmt: off\nimport numpy as np\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.distributions.categorical import Categorical\nfrom torch.utils.tensorboard import SummaryWriter\n\n\ndef parse_args():\n    # fmt: off\n    parser = argparse.ArgumentParser()\n    parser.add_argument('--exp-name', type=str, default=os.path.basename(__file__).rstrip(\".py\"),\n        help='the name of this experiment')\n    parser.add_argument('--gym-id', type=str, default=\"MicrortsMining10x10F9-v0\",\n        help='the id of the gym environment')\n    parser.add_argument('--learning-rate', type=float, default=2.5e-4,\n        help='the learning rate of the optimizer')\n    parser.add_argument('--seed', type=int, default=1,\n        help='seed of the experiment')\n    parser.add_argument('--total-timesteps', type=int, default=10000000,\n        help='total timesteps of the experiments')\n    parser.add_argument('--torch-deterministic', type=lambda x:bool(strtobool(x)), default=True, nargs='?', const=True,\n        help='if toggled, `torch.backends.cudnn.deterministic=False`')\n    parser.add_argument('--cuda', type=lambda x:bool(strtobool(x)), default=True, nargs='?', const=True,\n        help='if toggled, cuda will be enabled by default')\n    parser.add_argument('--track', type=lambda x:bool(strtobool(x)), default=False, nargs='?', const=True,\n        help='if toggled, this experiment will be tracked with Weights and Biases')\n    parser.add_argument('--wandb-project-name', type=str, default=\"cleanRL\",\n        help=\"the wandb's project name\")\n    parser.add_argument('--wandb-entity', type=str, default=None,\n        help=\"the entity (team) of wandb's project\")\n    parser.add_argument('--capture-video', type=lambda x:bool(strtobool(x)), default=False, nargs='?', const=True,\n        help='weather to capture videos of the agent performances (check out `videos` folder)')\n\n    # Algorithm specific arguments\n    parser.add_argument('--num-envs', type=int, default=4,\n        help='the number of parallel game environments')\n    parser.add_argument('--num-steps', type=int, default=128,\n        help='the number of steps to run in each environment per policy rollout')\n    parser.add_argument('--anneal-lr', type=lambda x:bool(strtobool(x)), default=True, nargs='?', const=True,\n        help=\"Toggle learning rate annealing for policy and value networks\")\n    parser.add_argument('--gae', type=lambda x:bool(strtobool(x)), default=True, nargs='?', const=True,\n        help='Use GAE for advantage computation')\n    parser.add_argument('--gamma', type=float, default=0.99,\n        help='the discount factor gamma')\n    parser.add_argument('--gae-lambda', type=float, default=0.95,\n        help='the lambda for the general advantage estimation')\n    parser.add_argument('--num-minibatches', type=int, default=4,\n        help='the number of mini-batches')\n    parser.add_argument('--update-epochs', type=int, default=4,\n        help=\"the K epochs to update the policy\")\n    parser.add_argument('--norm-adv', type=lambda x:bool(strtobool(x)), default=True, nargs='?', const=True,\n        help=\"Toggles advantages normalization\")\n    parser.add_argument('--clip-coef', type=float, default=0.1,\n        help=\"the surrogate clipping coefficient\")\n    parser.add_argument('--clip-vloss', type=lambda x:bool(strtobool(x)), default=True, nargs='?', const=True,\n        help='Toggles wheter or not to use a clipped loss for the value function, as per the paper.')\n    parser.add_argument('--ent-coef', type=float, default=0.01,\n        help=\"coefficient of the entropy\")\n    parser.add_argument('--vf-coef', type=float, default=0.5,\n        help=\"coefficient of the value function\")\n    parser.add_argument('--max-grad-norm', type=float, default=0.5,\n        help='the maximum norm for the gradient clipping')\n    parser.add_argument('--target-kl', type=float, default=None,\n        help='the target KL divergence threshold')\n    args = parser.parse_args()\n    args.batch_size = int(args.num_envs * args.num_steps)\n    args.minibatch_size = int(args.batch_size // args.num_minibatches)\n    # fmt: on\n    return args\n\n\ndef make_env(gym_id, seed, idx, capture_video, run_name):\n    def thunk():\n        env = gym.make(gym_id)\n        env = gym.wrappers.RecordEpisodeStatistics(env)\n        if capture_video:\n            if idx == 0:\n                env = gym.wrappers.RecordVideo(env, f\"videos/{run_name}\")\n        env.seed(seed)\n        env.action_space.seed(seed)\n        env.observation_space.seed(seed)\n        return env\n\n    return thunk\n\n\ndef 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\nclass Transpose(nn.Module):\n    def __init__(self, permutation):\n        super().__init__()\n        self.permutation = permutation\n\n    def forward(self, x):\n        return x.permute(self.permutation)\n\n\nclass Agent(nn.Module):\n    def __init__(self, envs):\n        super(Agent, self).__init__()\n        self.network = nn.Sequential(\n            Transpose((0, 3, 1, 2)),\n            layer_init(nn.Conv2d(27, 16, kernel_size=3, stride=2)),\n            nn.ReLU(),\n            layer_init(nn.Conv2d(16, 32, kernel_size=2)),\n            nn.ReLU(),\n            nn.Flatten(),\n            layer_init(nn.Linear(32*3*3, 128)),\n            nn.ReLU(),\n        )\n        self.nvec = envs.single_action_space.nvec\n        self.actor = layer_init(nn.Linear(128, self.nvec.sum()), std=0.01)\n        self.critic = layer_init(nn.Linear(128, 1), std=1)\n\n    def get_value(self, x):\n        return self.critic(self.network(x))\n\n    def get_action_and_value(self, x, action=None):\n        hidden = self.network(x)\n        logits = self.actor(hidden)\n        split_logits = torch.split(logits, self.nvec.tolist(), dim=1)\n        multi_categoricals = [Categorical(logits=logits) for logits in split_logits]\n        if action is None:\n            action = torch.stack([categorical.sample() for categorical in multi_categoricals])\n        logprob = torch.stack([categorical.log_prob(a) for a, categorical in zip(action, multi_categoricals)])\n        entropy = torch.stack([categorical.entropy() for categorical in multi_categoricals])\n        return action.T, logprob.sum(0), entropy.sum(0), self.critic(hidden)\n\n\nif __name__ == \"__main__\":\n    args = parse_args()\n    run_name = f\"{args.gym_id}__{args.exp_name}__{args.seed}__{int(time.time())}\"\n    if args.track:\n        import wandb\n\n        wandb.init(\n            project=args.wandb_project_name,\n            entity=args.wandb_entity,\n            sync_tensorboard=True,\n            config=vars(args),\n            name=run_name,\n            monitor_gym=True,\n            save_code=True,\n        )\n    writer = SummaryWriter(f\"runs/{run_name}\")\n    writer.add_text(\n        \"hyperparameters\",\n        \"|param|value|\\n|-|-|\\n%s\" % (\"\\n\".join([f\"|{key}|{value}|\" for key, value in vars(args).items()])),\n    )\n\n    # TRY NOT TO MODIFY: seeding\n    random.seed(args.seed)\n    np.random.seed(args.seed)\n    torch.manual_seed(args.seed)\n    torch.backends.cudnn.deterministic = args.torch_deterministic\n\n    device = torch.device(\"cuda\" if torch.cuda.is_available() and args.cuda else \"cpu\")\n\n    # env setup\n    envs = gym.vector.SyncVectorEnv(\n        [make_env(args.gym_id, args.seed + i, i, args.capture_video, run_name) for i in range(args.num_envs)]\n    )\n    assert isinstance(envs.single_action_space, gym.spaces.MultiDiscrete), \"only MultiDiscrete action space is supported\"\n\n    agent = Agent(envs).to(device)\n    optimizer = optim.Adam(agent.parameters(), lr=args.learning_rate, eps=1e-5)\n\n    # ALGO Logic: Storage setup\n    obs = torch.zeros((args.num_steps, args.num_envs) + envs.single_observation_space.shape).to(device)\n    actions = torch.zeros((args.num_steps, args.num_envs) + envs.single_action_space.shape).to(device)\n    logprobs = torch.zeros((args.num_steps, args.num_envs)).to(device)\n    rewards = torch.zeros((args.num_steps, args.num_envs)).to(device)\n    dones = torch.zeros((args.num_steps, args.num_envs)).to(device)\n    values = torch.zeros((args.num_steps, args.num_envs)).to(device)\n\n    # TRY NOT TO MODIFY: start the game\n    global_step = 0\n    start_time = time.time()\n    next_obs = torch.Tensor(envs.reset()).to(device)\n    next_done = torch.zeros(args.num_envs).to(device)\n    num_updates = args.total_timesteps // args.batch_size\n\n    for update in range(1, num_updates + 1):\n        # Annealing the rate if instructed to do so.\n        if args.anneal_lr:\n            frac = 1.0 - (update - 1.0) / num_updates\n            lrnow = frac * args.learning_rate\n            optimizer.param_groups[0][\"lr\"] = lrnow\n\n        for step in range(0, args.num_steps):\n            global_step += 1 * args.num_envs\n            obs[step] = next_obs\n            dones[step] = next_done\n\n            # ALGO LOGIC: action logic\n            with torch.no_grad():\n                action, logprob, _, value = agent.get_action_and_value(next_obs)\n                values[step] = value.flatten()\n            actions[step] = action\n            logprobs[step] = logprob\n\n            # TRY NOT TO MODIFY: execute the game and log data.\n            next_obs, reward, done, info = envs.step(action.cpu().numpy())\n            rewards[step] = torch.tensor(reward).to(device).view(-1)\n            next_obs, next_done = torch.Tensor(next_obs).to(device), torch.Tensor(done).to(device)\n\n            for item in info:\n                if \"episode\" in item.keys():\n                    print(f\"global_step={global_step}, episodic_return={item['episode']['r']}\")\n                    writer.add_scalar(\"charts/episodic_return\", item[\"episode\"][\"r\"], global_step)\n                    writer.add_scalar(\"charts/episodic_length\", item[\"episode\"][\"l\"], global_step)\n                    break\n\n        # bootstrap value if not done\n        with torch.no_grad():\n            next_value = agent.get_value(next_obs).reshape(1, -1)\n            if args.gae:\n                advantages = torch.zeros_like(rewards).to(device)\n                lastgaelam = 0\n                for t in reversed(range(args.num_steps)):\n                    if t == args.num_steps - 1:\n                        nextnonterminal = 1.0 - next_done\n                        nextvalues = next_value\n                    else:\n                        nextnonterminal = 1.0 - dones[t + 1]\n                        nextvalues = values[t + 1]\n                    delta = rewards[t] + args.gamma * nextvalues * nextnonterminal - values[t]\n                    advantages[t] = lastgaelam = delta + args.gamma * args.gae_lambda * nextnonterminal * lastgaelam\n                returns = advantages + values\n            else:\n                returns = torch.zeros_like(rewards).to(device)\n                for t in reversed(range(args.num_steps)):\n                    if t == args.num_steps - 1:\n                        nextnonterminal = 1.0 - next_done\n                        next_return = next_value\n                    else:\n                        nextnonterminal = 1.0 - dones[t + 1]\n                        next_return = returns[t + 1]\n                    returns[t] = rewards[t] + args.gamma * nextnonterminal * next_return\n                advantages = returns - values\n\n        # flatten the batch\n        b_obs = obs.reshape((-1,) + envs.single_observation_space.shape)\n        b_logprobs = logprobs.reshape(-1)\n        b_actions = actions.reshape((-1,) + envs.single_action_space.shape)\n        b_advantages = advantages.reshape(-1)\n        b_returns = returns.reshape(-1)\n        b_values = values.reshape(-1)\n\n        # Optimizaing the policy and value network\n        b_inds = np.arange(args.batch_size)\n        clipfracs = []\n        for epoch in range(args.update_epochs):\n            np.random.shuffle(b_inds)\n            for start in range(0, args.batch_size, args.minibatch_size):\n                end = start + args.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.long()[mb_inds].T)\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() > args.clip_coef).float().mean().item()]\n\n                mb_advantages = b_advantages[mb_inds]\n                if args.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 - args.clip_coef, 1 + args.clip_coef)\n                pg_loss = torch.max(pg_loss1, pg_loss2).mean()\n\n                # Value loss\n                newvalue = newvalue.view(-1)\n                if args.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                        -args.clip_coef,\n                        args.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 - args.ent_coef * entropy_loss + v_loss * args.vf_coef\n\n                optimizer.zero_grad()\n                loss.backward()\n                nn.utils.clip_grad_norm_(agent.parameters(), args.max_grad_norm)\n                optimizer.step()\n\n            if args.target_kl is not None:\n                if approx_kl > args.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\n\n        # TRY NOT TO MODIFY: record rewards for plotting purposes\n        writer.add_scalar(\"charts/learning_rate\", optimizer.param_groups[0][\"lr\"], global_step)\n        writer.add_scalar(\"losses/value_loss\", v_loss.item(), global_step)\n        writer.add_scalar(\"losses/policy_loss\", pg_loss.item(), global_step)\n        writer.add_scalar(\"losses/entropy\", entropy_loss.item(), global_step)\n        writer.add_scalar(\"losses/approx_kl\", approx_kl.item(), global_step)\n        writer.add_scalar(\"losses/clipfrac\", np.mean(clipfracs), global_step)\n        writer.add_scalar(\"losses/explained_variance\", explained_var, global_step)\n        print(\"SPS:\", int(global_step / (time.time() - start_time)))\n        writer.add_scalar(\"charts/SPS\", int(global_step / (time.time() - start_time)), global_step)\n\n    envs.close()\n    writer.close()\n"
  },
  {
    "path": "gym_vec_api/ppo_multidiscrete_mask.py",
    "content": "import argparse\nimport os\nimport random\nimport time\nfrom distutils.util import strtobool\n\nimport gym\nimport gym_microrts # fmt: off\nimport numpy as np\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.distributions.categorical import Categorical\nfrom torch.utils.tensorboard import SummaryWriter\n\n\ndef parse_args():\n    # fmt: off\n    parser = argparse.ArgumentParser()\n    parser.add_argument('--exp-name', type=str, default=os.path.basename(__file__).rstrip(\".py\"),\n        help='the name of this experiment')\n    parser.add_argument('--gym-id', type=str, default=\"MicrortsMining10x10F9-v0\",\n        help='the id of the gym environment')\n    parser.add_argument('--learning-rate', type=float, default=2.5e-4,\n        help='the learning rate of the optimizer')\n    parser.add_argument('--seed', type=int, default=1,\n        help='seed of the experiment')\n    parser.add_argument('--total-timesteps', type=int, default=10000000,\n        help='total timesteps of the experiments')\n    parser.add_argument('--torch-deterministic', type=lambda x:bool(strtobool(x)), default=True, nargs='?', const=True,\n        help='if toggled, `torch.backends.cudnn.deterministic=False`')\n    parser.add_argument('--cuda', type=lambda x:bool(strtobool(x)), default=True, nargs='?', const=True,\n        help='if toggled, cuda will be enabled by default')\n    parser.add_argument('--track', type=lambda x:bool(strtobool(x)), default=False, nargs='?', const=True,\n        help='if toggled, this experiment will be tracked with Weights and Biases')\n    parser.add_argument('--wandb-project-name', type=str, default=\"cleanRL\",\n        help=\"the wandb's project name\")\n    parser.add_argument('--wandb-entity', type=str, default=None,\n        help=\"the entity (team) of wandb's project\")\n    parser.add_argument('--capture-video', type=lambda x:bool(strtobool(x)), default=False, nargs='?', const=True,\n        help='weather to capture videos of the agent performances (check out `videos` folder)')\n\n    # Algorithm specific arguments\n    parser.add_argument('--num-envs', type=int, default=4,\n        help='the number of parallel game environments')\n    parser.add_argument('--num-steps', type=int, default=128,\n        help='the number of steps to run in each environment per policy rollout')\n    parser.add_argument('--anneal-lr', type=lambda x:bool(strtobool(x)), default=True, nargs='?', const=True,\n        help=\"Toggle learning rate annealing for policy and value networks\")\n    parser.add_argument('--gae', type=lambda x:bool(strtobool(x)), default=True, nargs='?', const=True,\n        help='Use GAE for advantage computation')\n    parser.add_argument('--gamma', type=float, default=0.99,\n        help='the discount factor gamma')\n    parser.add_argument('--gae-lambda', type=float, default=0.95,\n        help='the lambda for the general advantage estimation')\n    parser.add_argument('--num-minibatches', type=int, default=4,\n        help='the number of mini-batches')\n    parser.add_argument('--update-epochs', type=int, default=4,\n        help=\"the K epochs to update the policy\")\n    parser.add_argument('--norm-adv', type=lambda x:bool(strtobool(x)), default=True, nargs='?', const=True,\n        help=\"Toggles advantages normalization\")\n    parser.add_argument('--clip-coef', type=float, default=0.1,\n        help=\"the surrogate clipping coefficient\")\n    parser.add_argument('--clip-vloss', type=lambda x:bool(strtobool(x)), default=True, nargs='?', const=True,\n        help='Toggles wheter or not to use a clipped loss for the value function, as per the paper.')\n    parser.add_argument('--ent-coef', type=float, default=0.01,\n        help=\"coefficient of the entropy\")\n    parser.add_argument('--vf-coef', type=float, default=0.5,\n        help=\"coefficient of the value function\")\n    parser.add_argument('--max-grad-norm', type=float, default=0.5,\n        help='the maximum norm for the gradient clipping')\n    parser.add_argument('--target-kl', type=float, default=None,\n        help='the target KL divergence threshold')\n    args = parser.parse_args()\n    args.batch_size = int(args.num_envs * args.num_steps)\n    args.minibatch_size = int(args.batch_size // args.num_minibatches)\n    # fmt: on\n    return args\n\n\ndef make_env(gym_id, seed, idx, capture_video, run_name):\n    def thunk():\n        env = gym.make(gym_id)\n        env = gym.wrappers.RecordEpisodeStatistics(env)\n        if capture_video:\n            if idx == 0:\n                env = gym.wrappers.RecordVideo(env, f\"videos/{run_name}\")\n        env.seed(seed)\n        env.action_space.seed(seed)\n        env.observation_space.seed(seed)\n        return env\n\n    return thunk\n\n\ndef 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\nclass Transpose(nn.Module):\n    def __init__(self, permutation):\n        super().__init__()\n        self.permutation = permutation\n\n    def forward(self, x):\n        return x.permute(self.permutation)\n\n\nclass CategoricalMasked(Categorical):\n    def __init__(self, probs=None, logits=None, validate_args=None, masks=[]):\n        self.masks = masks\n        if len(self.masks) == 0:\n            super(CategoricalMasked, self).__init__(probs, logits, validate_args)\n        else:\n            self.masks = masks.type(torch.BoolTensor).to(device)\n            logits = torch.where(self.masks, logits, torch.tensor(-1e+8).to(device))\n            super(CategoricalMasked, self).__init__(probs, logits, validate_args)\n    \n    def entropy(self):\n        if len(self.masks) == 0:\n            return super(CategoricalMasked, self).entropy()\n        p_log_p = self.logits * self.probs\n        p_log_p = torch.where(self.masks, p_log_p, torch.tensor(0.).to(device))\n        return -p_log_p.sum(-1)\n\n\nclass Agent(nn.Module):\n    def __init__(self, envs):\n        super(Agent, self).__init__()\n        self.network = nn.Sequential(\n            Transpose((0, 3, 1, 2)),\n            layer_init(nn.Conv2d(27, 16, kernel_size=3, stride=2)),\n            nn.ReLU(),\n            layer_init(nn.Conv2d(16, 32, kernel_size=2)),\n            nn.ReLU(),\n            nn.Flatten(),\n            layer_init(nn.Linear(32*3*3, 128)),\n            nn.ReLU(),\n        )\n        self.nvec = envs.single_action_space.nvec\n        self.actor = layer_init(nn.Linear(128, self.nvec.sum()), std=0.01)\n        self.critic = layer_init(nn.Linear(128, 1), std=1)\n\n    def get_value(self, x):\n        return self.critic(self.network(x))\n\n    def get_action_and_value(self, x, action_mask, action=None):\n        hidden = self.network(x)\n        logits = self.actor(hidden)\n        split_logits = torch.split(logits, self.nvec.tolist(), dim=1)\n        split_action_masks = torch.split(action_mask, self.nvec.tolist(), dim=1)\n        multi_categoricals = [\n            CategoricalMasked(logits=logits, masks=iam)\n            for (logits, iam) in zip(split_logits, split_action_masks)\n        ]\n        if action is None:\n            action = torch.stack([categorical.sample() for categorical in multi_categoricals])\n        logprob = torch.stack([categorical.log_prob(a) for a, categorical in zip(action, multi_categoricals)])\n        entropy = torch.stack([categorical.entropy() for categorical in multi_categoricals])\n        return action.T, logprob.sum(0), entropy.sum(0), self.critic(hidden)\n\n\nif __name__ == \"__main__\":\n    args = parse_args()\n    run_name = f\"{args.gym_id}__{args.exp_name}__{args.seed}__{int(time.time())}\"\n    if args.track:\n        import wandb\n\n        wandb.init(\n            project=args.wandb_project_name,\n            entity=args.wandb_entity,\n            sync_tensorboard=True,\n            config=vars(args),\n            name=run_name,\n            monitor_gym=True,\n            save_code=True,\n        )\n    writer = SummaryWriter(f\"runs/{run_name}\")\n    writer.add_text(\n        \"hyperparameters\",\n        \"|param|value|\\n|-|-|\\n%s\" % (\"\\n\".join([f\"|{key}|{value}|\" for key, value in vars(args).items()])),\n    )\n\n    # TRY NOT TO MODIFY: seeding\n    random.seed(args.seed)\n    np.random.seed(args.seed)\n    torch.manual_seed(args.seed)\n    torch.backends.cudnn.deterministic = args.torch_deterministic\n\n    device = torch.device(\"cuda\" if torch.cuda.is_available() and args.cuda else \"cpu\")\n\n    # env setup\n    envs = gym.vector.SyncVectorEnv(\n        [make_env(args.gym_id, args.seed + i, i, args.capture_video, run_name) for i in range(args.num_envs)]\n    )\n    assert isinstance(envs.single_action_space, gym.spaces.MultiDiscrete), \"only MultiDiscrete action space is supported\"\n\n    agent = Agent(envs).to(device)\n    optimizer = optim.Adam(agent.parameters(), lr=args.learning_rate, eps=1e-5)\n\n    # ALGO Logic: Storage setup\n    obs = torch.zeros((args.num_steps, args.num_envs) + envs.single_observation_space.shape).to(device)\n    actions = torch.zeros((args.num_steps, args.num_envs) + envs.single_action_space.shape).to(device)\n    logprobs = torch.zeros((args.num_steps, args.num_envs)).to(device)\n    rewards = torch.zeros((args.num_steps, args.num_envs)).to(device)\n    dones = torch.zeros((args.num_steps, args.num_envs)).to(device)\n    values = torch.zeros((args.num_steps, args.num_envs)).to(device)\n    action_masks = torch.zeros((args.num_steps, args.num_envs) + (envs.single_action_space.nvec.sum(),)).to(device)\n\n    # TRY NOT TO MODIFY: start the game\n    global_step = 0\n    start_time = time.time()\n    next_obs = torch.Tensor(envs.reset()).to(device)\n    next_done = torch.zeros(args.num_envs).to(device)\n    num_updates = args.total_timesteps // args.batch_size\n\n    for update in range(1, num_updates + 1):\n        # Annealing the rate if instructed to do so.\n        if args.anneal_lr:\n            frac = 1.0 - (update - 1.0) / num_updates\n            lrnow = frac * args.learning_rate\n            optimizer.param_groups[0][\"lr\"] = lrnow\n\n        for step in range(0, args.num_steps):\n            global_step += 1 * args.num_envs\n            obs[step] = next_obs\n            dones[step] = next_done\n            action_masks[step] = torch.Tensor(\n                np.array([env.action_mask for env in envs.envs])\n            )\n\n            # ALGO LOGIC: action logic\n            with torch.no_grad():\n                action, logprob, _, value = agent.get_action_and_value(next_obs, action_masks[step])\n                values[step] = value.flatten()\n            actions[step] = action\n            logprobs[step] = logprob\n\n            # TRY NOT TO MODIFY: execute the game and log data.\n            next_obs, reward, done, info = envs.step(action.cpu().numpy())\n            rewards[step] = torch.tensor(reward).to(device).view(-1)\n            next_obs, next_done = torch.Tensor(next_obs).to(device), torch.Tensor(done).to(device)\n\n            for item in info:\n                if \"episode\" in item.keys():\n                    print(f\"global_step={global_step}, episodic_return={item['episode']['r']}\")\n                    writer.add_scalar(\"charts/episodic_return\", item[\"episode\"][\"r\"], global_step)\n                    writer.add_scalar(\"charts/episodic_length\", item[\"episode\"][\"l\"], global_step)\n                    break\n\n        # bootstrap value if not done\n        with torch.no_grad():\n            next_value = agent.get_value(next_obs).reshape(1, -1)\n            if args.gae:\n                advantages = torch.zeros_like(rewards).to(device)\n                lastgaelam = 0\n                for t in reversed(range(args.num_steps)):\n                    if t == args.num_steps - 1:\n                        nextnonterminal = 1.0 - next_done\n                        nextvalues = next_value\n                    else:\n                        nextnonterminal = 1.0 - dones[t + 1]\n                        nextvalues = values[t + 1]\n                    delta = rewards[t] + args.gamma * nextvalues * nextnonterminal - values[t]\n                    advantages[t] = lastgaelam = delta + args.gamma * args.gae_lambda * nextnonterminal * lastgaelam\n                returns = advantages + values\n            else:\n                returns = torch.zeros_like(rewards).to(device)\n                for t in reversed(range(args.num_steps)):\n                    if t == args.num_steps - 1:\n                        nextnonterminal = 1.0 - next_done\n                        next_return = next_value\n                    else:\n                        nextnonterminal = 1.0 - dones[t + 1]\n                        next_return = returns[t + 1]\n                    returns[t] = rewards[t] + args.gamma * nextnonterminal * next_return\n                advantages = returns - values\n\n        # flatten the batch\n        b_obs = obs.reshape((-1,) + envs.single_observation_space.shape)\n        b_logprobs = logprobs.reshape(-1)\n        b_actions = actions.reshape((-1,) + envs.single_action_space.shape)\n        b_advantages = advantages.reshape(-1)\n        b_returns = returns.reshape(-1)\n        b_values = values.reshape(-1)\n        b_action_masks = action_masks.reshape((-1, action_masks.shape[-1]))\n\n        # Optimizaing the policy and value network\n        b_inds = np.arange(args.batch_size)\n        clipfracs = []\n        for epoch in range(args.update_epochs):\n            np.random.shuffle(b_inds)\n            for start in range(0, args.batch_size, args.minibatch_size):\n                end = start + args.minibatch_size\n                mb_inds = b_inds[start:end]\n\n                _, newlogprob, entropy, newvalue = agent.get_action_and_value(\n                    b_obs[mb_inds],\n                    b_action_masks[mb_inds],\n                    b_actions.long()[mb_inds].T,\n                )\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() > args.clip_coef).float().mean().item()]\n\n                mb_advantages = b_advantages[mb_inds]\n                if args.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 - args.clip_coef, 1 + args.clip_coef)\n                pg_loss = torch.max(pg_loss1, pg_loss2).mean()\n\n                # Value loss\n                newvalue = newvalue.view(-1)\n                if args.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                        -args.clip_coef,\n                        args.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 - args.ent_coef * entropy_loss + v_loss * args.vf_coef\n\n                optimizer.zero_grad()\n                loss.backward()\n                nn.utils.clip_grad_norm_(agent.parameters(), args.max_grad_norm)\n                optimizer.step()\n\n            if args.target_kl is not None:\n                if approx_kl > args.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\n\n        # TRY NOT TO MODIFY: record rewards for plotting purposes\n        writer.add_scalar(\"charts/learning_rate\", optimizer.param_groups[0][\"lr\"], global_step)\n        writer.add_scalar(\"losses/value_loss\", v_loss.item(), global_step)\n        writer.add_scalar(\"losses/policy_loss\", pg_loss.item(), global_step)\n        writer.add_scalar(\"losses/entropy\", entropy_loss.item(), global_step)\n        writer.add_scalar(\"losses/approx_kl\", approx_kl.item(), global_step)\n        writer.add_scalar(\"losses/clipfrac\", np.mean(clipfracs), global_step)\n        writer.add_scalar(\"losses/explained_variance\", explained_var, global_step)\n        print(\"SPS:\", int(global_step / (time.time() - start_time)))\n        writer.add_scalar(\"charts/SPS\", int(global_step / (time.time() - start_time)), global_step)\n\n    envs.close()\n    writer.close()\n"
  },
  {
    "path": "invalid_action_masking/ppo_10x10.py",
    "content": "import torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport torch.nn.functional as F\nfrom torch.distributions.categorical import Categorical\nfrom torch.utils.tensorboard import SummaryWriter\n\nfrom cleanrl.common import preprocess_obs_space, preprocess_ac_space\nimport argparse\nimport numpy as np\nimport gym\nimport gym_microrts\nfrom gym.wrappers import TimeLimit, Monitor\nfrom gym.spaces import Discrete, Box, MultiBinary, MultiDiscrete, Space\nimport time\nimport random\nimport os\nimport pandas as pd\n\n# taken from https://github.com/openai/baselines/blob/master/baselines/common/vec_env/vec_normalize.py\nclass RunningMeanStd(object):\n    def __init__(self, epsilon=1e-4, shape=()):\n        self.mean = np.zeros(shape, 'float64')\n        self.var = np.ones(shape, 'float64')\n        self.count = epsilon\n\n    def update(self, x):\n        batch_mean = np.mean([x], axis=0)\n        batch_var = np.var([x], axis=0)\n        batch_count = 1\n        self.update_from_moments(batch_mean, batch_var, batch_count)\n\n    def update_from_moments(self, batch_mean, batch_var, batch_count):\n        self.mean, self.var, self.count = update_mean_var_count_from_moments(\n            self.mean, self.var, self.count, batch_mean, batch_var, batch_count)\n\ndef update_mean_var_count_from_moments(mean, var, count, batch_mean, batch_var, batch_count):\n    delta = batch_mean - mean\n    tot_count = count + batch_count\n\n    new_mean = mean + delta * batch_count / tot_count\n    m_a = var * count\n    m_b = batch_var * batch_count\n    M2 = m_a + m_b + np.square(delta) * count * batch_count / tot_count\n    new_var = M2 / tot_count\n    new_count = tot_count\n\n    return new_mean, new_var, new_count\nclass NormalizedEnv(gym.core.Wrapper):\n    def __init__(self, env, ob=True, ret=True, clipob=10., cliprew=10., gamma=0.99, epsilon=1e-8):\n        super(NormalizedEnv, self).__init__(env)\n        self.ob_rms = RunningMeanStd(shape=self.observation_space.shape) if ob else None\n        self.ret_rms = RunningMeanStd(shape=(1,)) if ret else None\n        self.clipob = clipob\n        self.cliprew = cliprew\n        self.ret = np.zeros(())\n        self.gamma = gamma\n        self.epsilon = epsilon\n\n    def step(self, action):\n        obs, rews, news, infos = self.env.step(action)\n        infos['real_reward'] = rews\n        # print(\"before\", self.ret)\n        self.ret = self.ret * self.gamma + rews\n        # print(\"after\", self.ret)\n        obs = self._obfilt(obs)\n        if self.ret_rms:\n            self.ret_rms.update(np.array([self.ret].copy()))\n            rews = np.clip(rews / np.sqrt(self.ret_rms.var + self.epsilon), -self.cliprew, self.cliprew)\n        self.ret = self.ret * (1-float(news))\n        return obs, rews, news, infos\n\n    def _obfilt(self, obs):\n        if self.ob_rms:\n            self.ob_rms.update(obs)\n            obs = np.clip((obs - self.ob_rms.mean) / np.sqrt(self.ob_rms.var + self.epsilon), -self.clipob, self.clipob)\n            return obs\n        else:\n            return obs\n\n    def reset(self):\n        self.ret = np.zeros(())\n        obs = self.env.reset()\n        return self._obfilt(obs)\n\nif __name__ == \"__main__\":\n    parser = argparse.ArgumentParser(description='PPO agent')\n    # Common arguments\n    parser.add_argument('--exp-name', type=str, default=os.path.basename(__file__).rstrip(\".py\"),\n                       help='the name of this experiment')\n    parser.add_argument('--gym-id', type=str, default=\"MicrortsMining10x10F9-v0\",\n                       help='the id of the gym environment')\n    parser.add_argument('--seed', type=int, default=1,\n                       help='seed of the experiment')\n    parser.add_argument('--episode-length', type=int, default=0,\n                       help='the maximum length of each episode')\n    parser.add_argument('--total-timesteps', type=int, default=100000,\n                       help='total timesteps of the experiments')\n    parser.add_argument('--no-torch-deterministic', action='store_false', dest=\"torch_deterministic\", default=True,\n                       help='if toggled, `torch.backends.cudnn.deterministic=False`')\n    parser.add_argument('--no-cuda', action='store_false', dest=\"cuda\", default=True,\n                       help='if toggled, cuda will not be enabled by default')\n    parser.add_argument('--prod-mode', action='store_true', default=False,\n                       help='run the script in production mode and use wandb to log outputs')\n    parser.add_argument('--capture-video', action='store_true', default=False,\n                       help='weather to capture videos of the agent performances (check out `videos` folder)')\n    parser.add_argument('--wandb-project-name', type=str, default=\"cleanRL\",\n                       help=\"the wandb's project name\")\n    parser.add_argument('--wandb-entity', type=str, default=None,\n                       help=\"the entity (team) of wandb's project\")\n\n    # Algorithm specific arguments\n    parser.add_argument('--batch-size', type=int, default=2048,\n                       help='the batch size of ppo')\n    parser.add_argument('--minibatch-size', type=int, default=256,\n                       help='the mini batch size of ppo')\n    parser.add_argument('--gamma', type=float, default=0.99,\n                       help='the discount factor gamma')\n    parser.add_argument('--gae-lambda', type=float, default=0.97,\n                       help='the lambda for the general advantage estimation')\n    parser.add_argument('--ent-coef', type=float, default=0.01,\n                       help=\"coefficient of the entropy\")\n    parser.add_argument('--max-grad-norm', type=float, default=0.5,\n                       help='the maximum norm for the gradient clipping')\n    parser.add_argument('--clip-coef', type=float, default=0.2,\n                       help=\"the surrogate clipping coefficient\")\n    parser.add_argument('--update-epochs', type=int, default=10,\n                        help=\"the K epochs to update the policy\")\n    parser.add_argument('--kle-stop', action='store_true', default=False,\n                        help='If toggled, the policy updates will be early stopped w.r.t target-kl')\n    parser.add_argument('--kle-rollback', action='store_true', default=False,\n                        help='If toggled, the policy updates will roll back to previous policy if KL exceeds target-kl')\n    parser.add_argument('--target-kl', type=float, default=0.015,\n                        help='the target-kl variable that is referred by --kl')\n    parser.add_argument('--gae', action='store_true', default=True,\n                        help='Use GAE for advantage computation')\n    parser.add_argument('--policy-lr', type=float, default=3e-4,\n                        help=\"the learning rate of the policy optimizer\")\n    parser.add_argument('--value-lr', type=float, default=3e-4,\n                        help=\"the learning rate of the critic optimizer\")\n    parser.add_argument('--norm-obs', action='store_true', default=False,\n                        help=\"Toggles observation normalization\")\n    parser.add_argument('--norm-returns', action='store_true', default=False,\n                        help=\"Toggles returns normalization\")\n    parser.add_argument('--norm-adv', action='store_true', default=False,\n                        help=\"Toggles advantages normalization\")\n    parser.add_argument('--obs-clip', type=float, default=10.0,\n                        help=\"Value for reward clipping, as per the paper\")\n    parser.add_argument('--rew-clip', type=float, default=10.0,\n                        help=\"Value for observation clipping, as per the paper\")\n    parser.add_argument('--anneal-lr', action='store_true', default=False,\n                        help=\"Toggle learning rate annealing for policy and value networks\")\n    parser.add_argument('--weights-init', default=\"orthogonal\", choices=[\"xavier\", 'orthogonal'],\n                        help='Selects the scheme to be used for weights initialization'),\n    parser.add_argument('--clip-vloss', action=\"store_true\", default=False,\n                        help='Toggles wheter or not to use a clipped loss for the value function, as per the paper.')\n    parser.add_argument('--pol-layer-norm', action='store_true', default=False,\n                       help='Enables layer normalization in the policy network')\n    args = parser.parse_args()\n    if not args.seed:\n        args.seed = int(time.time())\n\nargs.features_turned_on = sum([args.kle_stop, args.kle_rollback, args.gae, args.norm_obs, args.norm_returns, args.norm_adv, args.anneal_lr, args.clip_vloss, args.pol_layer_norm])\n# TRY NOT TO MODIFY: setup the environment\nexperiment_name = f\"{args.gym_id}__{args.exp_name}__{args.seed}__{int(time.time())}\"\nwriter = SummaryWriter(f\"runs/{experiment_name}\")\nwriter.add_text('hyperparameters', \"|param|value|\\n|-|-|\\n%s\" % (\n        '\\n'.join([f\"|{key}|{value}|\" for key, value in vars(args).items()])))\n\nif args.prod_mode:\n    import wandb\n    wandb.init(project=args.wandb_project_name, entity=args.wandb_entity, tensorboard=True, config=vars(args), name=experiment_name, monitor_gym=True)\n    writer = SummaryWriter(f\"/tmp/{experiment_name}\")\n    wandb.save(os.path.abspath(__file__))\n\n# TRY NOT TO MODIFY: seeding\ndevice = torch.device('cuda' if torch.cuda.is_available() and args.cuda else 'cpu')\nenv = gym.make(args.gym_id)\n# respect the default timelimit\nassert isinstance(env.action_space, MultiDiscrete), \"only MultiDiscrete action space is supported\"\nassert isinstance(env, TimeLimit) or int(args.episode_length), \"the gym env does not have a built in TimeLimit, please specify by using --episode-length\"\nif isinstance(env, TimeLimit):\n    if int(args.episode_length):\n        env._max_episode_steps = int(args.episode_length)\n    args.episode_length = env._max_episode_steps\nelse:\n    env = TimeLimit(env, int(args.episode_length))\nenv = NormalizedEnv(env.env, ob=args.norm_obs, ret=args.norm_returns, clipob=args.obs_clip, cliprew=args.rew_clip, gamma=args.gamma)\nenv = TimeLimit(env, int(args.episode_length))\nrandom.seed(args.seed)\nnp.random.seed(args.seed)\ntorch.manual_seed(args.seed)\ntorch.backends.cudnn.deterministic = args.torch_deterministic\nenv.seed(args.seed)\nenv.action_space.seed(args.seed)\nenv.observation_space.seed(args.seed)\nif args.capture_video:\n    env = Monitor(env, f'videos/{experiment_name}')\n\n# ALGO LOGIC: initialize agent here:\nclass CategoricalMasked(Categorical):\n    def __init__(self, probs=None, logits=None, validate_args=None, masks=[]):\n        self.masks = masks\n        if len(self.masks) == 0:\n            super(CategoricalMasked, self).__init__(probs, logits, validate_args)\n        else:\n            self.masks = masks.type(torch.BoolTensor).to(device)\n            logits = torch.where(self.masks, logits, torch.tensor(-1e+8).to(device))\n            super(CategoricalMasked, self).__init__(probs, logits, validate_args)\n    \n    def entropy(self):\n        if len(self.masks) == 0:\n            return super(CategoricalMasked, self).entropy()\n        p_log_p = self.logits * self.probs\n        p_log_p = torch.where(self.masks, p_log_p, torch.tensor(0.).to(device))\n        return -p_log_p.sum(-1)\n\nclass Policy(nn.Module):\n    def __init__(self):\n        super(Policy, self).__init__()\n        self.features = nn.Sequential(\n            nn.Conv2d(27, 16, kernel_size=3,),\n            nn.MaxPool2d(1),\n            nn.ReLU(),\n            nn.Conv2d(16, 32, kernel_size=3),\n            nn.MaxPool2d(1),\n            nn.ReLU())\n        self.fc = nn.Sequential(\n            nn.Linear(32*6*6, 128),\n            nn.ReLU(),\n            nn.Linear(128, env.action_space.nvec.sum())\n        )\n\n    def forward(self, x):\n        x = torch.Tensor(np.moveaxis(x, -1, 1)).to(device)\n        x = self.features(x)\n        x = x.reshape(x.size(0), -1)\n        x = self.fc(x)\n        return x\n\n    def get_action(self, x, action=None, invalid_action_masks=None):\n        logits = self.forward(x)\n        split_logits = torch.split(logits, env.action_space.nvec.tolist(), dim=1)\n        \n        if invalid_action_masks is not None:\n            split_invalid_action_masks = torch.split(invalid_action_masks, env.action_space.nvec.tolist(), dim=1)\n            multi_categoricals = [CategoricalMasked(logits=logits, masks=iam) for (logits, iam) in zip(split_logits, split_invalid_action_masks)]\n        else:\n            multi_categoricals = [Categorical(logits=logits) for logits in split_logits]\n        \n        if action is None:\n            action = torch.stack([categorical.sample() for categorical in multi_categoricals])\n        logprob = torch.stack([categorical.log_prob(a) for a, categorical in zip(action, multi_categoricals)])\n        return action, logprob, [], multi_categoricals\n\nclass Value(nn.Module):\n    def __init__(self):\n        super(Value, self).__init__()\n        self.features = nn.Sequential(\n            nn.Conv2d(27, 16, kernel_size=3,),\n            nn.MaxPool2d(1),\n            nn.ReLU(),\n            nn.Conv2d(16, 32, kernel_size=3),\n            nn.MaxPool2d(1),\n            nn.ReLU())\n        self.fc = nn.Sequential(\n            nn.Linear(32*6*6, 128),\n            nn.ReLU(),\n            nn.Linear(128, 1)\n        )\n\n    def forward(self, x):\n        x = torch.Tensor(np.moveaxis(x, -1, 1)).to(device)\n        x = self.features(x)\n        x = x.reshape(x.size(0), -1)\n        x = self.fc(x)\n        return x\n\ndef discount_cumsum(x, dones, gamma):\n    \"\"\"\n    computing discounted cumulative sums of vectors that resets with dones\n    input:\n        vector x,  vector dones,\n        [x0,       [0,\n         x1,        0,\n         x2         1,\n         x3         0, \n         x4]        0]\n    output:\n        [x0 + discount * x1 + discount^2 * x2,\n         x1 + discount * x2,\n         x2,\n         x3 + discount * x4,\n         x4]\n    \"\"\"\n    discount_cumsum = np.zeros_like(x)\n    discount_cumsum[-1] = x[-1]\n    for t in reversed(range(x.shape[0]-1)):\n        discount_cumsum[t] = x[t] + gamma * discount_cumsum[t+1] * (1-dones[t])\n    return discount_cumsum\n\npg = Policy().to(device)\nvf = Value().to(device)\n\n# MODIFIED: Separate optimizer and learning rates\npg_optimizer = optim.Adam(list(pg.parameters()), lr=args.policy_lr)\nv_optimizer = optim.Adam(list(vf.parameters()), lr=args.value_lr)\n\n# MODIFIED: Initializing learning rate anneal scheduler when need\nif args.anneal_lr:\n    anneal_fn = lambda f: max(0, 1-f / args.total_timesteps)\n    pg_lr_scheduler = optim.lr_scheduler.LambdaLR(pg_optimizer, lr_lambda=anneal_fn)\n    vf_lr_scheduler = optim.lr_scheduler.LambdaLR(v_optimizer, lr_lambda=anneal_fn)\n\nloss_fn = nn.MSELoss()\n\ndef evaluate_with_no_mask():\n    evaluate_rewards = []\n    evaluate_invalid_action_stats = []\n    if args.capture_video:\n        env.stats_recorder.done=True\n    next_obs = np.array(env.reset())\n\n    # ALGO Logic: Storage for epoch data\n    obs = np.empty((args.batch_size,) + env.observation_space.shape)\n\n    actions = np.empty((args.batch_size,) + env.action_space.shape)\n    logprobs = torch.zeros((env.action_space.nvec.shape[0], args.batch_size,)).to(device)\n\n    rewards = np.zeros((args.batch_size,))\n    raw_rewards = np.zeros((len(env.rfs),args.batch_size,))\n    \n    real_rewards = []\n    invalid_action_stats = []\n\n    dones = np.zeros((args.batch_size,))\n    values = torch.zeros((args.batch_size,)).to(device)\n\n    # TRY NOT TO MODIFY: prepare the execution of the game.\n    for step in range(args.batch_size):\n        env.render()\n        obs[step] = next_obs.copy()\n\n        # ALGO LOGIC: put action logic here\n        with torch.no_grad():\n            values[step] = vf.forward(obs[step:step+1])\n            action, logproba, _, probs = pg.get_action(obs[step:step+1])\n        \n        actions[step] = action[:,0].data.cpu().numpy()\n        logprobs[:,[step]] = logproba\n\n        # TRY NOT TO MODIFY: execute the game and log data.\n        next_obs, rewards[step], dones[step], info = env.step(action[:,0].data.cpu().numpy())\n        raw_rewards[:,step] = info[\"rewards\"]\n        real_rewards += [info['real_reward']]\n        invalid_action_stats += [info['invalid_action_stats']]\n        next_obs = np.array(next_obs)\n\n        if dones[step]:\n            # Computing the discounted returns:\n            writer.add_scalar(\"charts/episode_reward\", np.sum(real_rewards), global_step)\n            evaluate_rewards += [np.sum(real_rewards)]\n            real_rewards = []\n            evaluate_invalid_action_stats += [pd.DataFrame(invalid_action_stats).sum(0)]\n            invalid_action_stats = []\n            next_obs = np.array(env.reset())\n    \n    return np.average(evaluate_rewards), pd.DataFrame(evaluate_invalid_action_stats).mean(0)\n\n# TRY NOT TO MODIFY: start the game\nglobal_step = 0\nwhile global_step < args.total_timesteps:\n    if args.capture_video:\n        env.stats_recorder.done=True\n    next_obs = np.array(env.reset())\n\n    # ALGO Logic: Storage for epoch data\n    obs = np.empty((args.batch_size,) + env.observation_space.shape)\n\n    actions = np.empty((args.batch_size,) + env.action_space.shape)\n    logprobs = torch.zeros((env.action_space.nvec.shape[0], args.batch_size,)).to(device)\n\n    rewards = np.zeros((args.batch_size,))\n    raw_rewards = np.zeros((len(env.rfs),args.batch_size,))\n    \n    real_rewards = []\n    invalid_action_stats = []\n\n    dones = np.zeros((args.batch_size,))\n    values = torch.zeros((args.batch_size,)).to(device)\n\n    invalid_action_masks = torch.zeros((args.batch_size, env.action_space.nvec.sum()))\n    # TRY NOT TO MODIFY: prepare the execution of the game.\n    for step in range(args.batch_size):\n        env.render()\n        global_step += 1\n        obs[step] = next_obs.copy()\n\n        # ALGO LOGIC: put action logic here\n        invalid_action_mask = torch.ones(env.action_space.nvec.sum())\n        invalid_action_mask[0:env.action_space.nvec[0]] = torch.tensor(env.unit_location_mask)\n        invalid_action_mask[-env.action_space.nvec[-1]:] = torch.tensor(env.target_unit_location_mask)\n        invalid_action_masks[step] = invalid_action_mask\n        with torch.no_grad():\n            values[step] = vf.forward(obs[step:step+1])\n            action, logproba, _, probs = pg.get_action(obs[step:step+1], invalid_action_masks=invalid_action_masks[step:step+1])\n        \n        actions[step] = action[:,0].data.cpu().numpy()\n        logprobs[:,[step]] = logproba\n\n        # TRY NOT TO MODIFY: execute the game and log data.\n        next_obs, rewards[step], dones[step], info = env.step(action[:,0].data.cpu().numpy())\n        raw_rewards[:,step] = info[\"rewards\"]\n        real_rewards += [info['real_reward']]\n        invalid_action_stats += [info['invalid_action_stats']]\n        next_obs = np.array(next_obs)\n\n        # Annealing the rate if instructed to do so.\n        if args.anneal_lr:\n            pg_lr_scheduler.step()\n            vf_lr_scheduler.step()\n\n        if dones[step]:\n            # Computing the discounted returns:\n            writer.add_scalar(\"charts/episode_reward\", np.sum(real_rewards), global_step)\n            print(f\"global_step={global_step}, episode_reward={np.sum(real_rewards)}\")\n            for i in range(len(env.rfs)):\n                writer.add_scalar(f\"charts/episode_reward/{str(env.rfs[i])}\", raw_rewards.sum(1)[i], global_step)\n            real_rewards = []\n            for key, idx in zip(info['invalid_action_stats'], range(len(info['invalid_action_stats']))):\n                writer.add_scalar(f\"stats/{key}\", pd.DataFrame(invalid_action_stats).sum(0)[idx], global_step)\n            invalid_action_stats = []\n            next_obs = np.array(env.reset())\n\n    # bootstrap reward if not done. reached the batch limit\n    last_value = 0\n    if not dones[step]:\n        last_value = vf.forward(next_obs.reshape((1,)+next_obs.shape))[0].detach().cpu().numpy()[0]\n    bootstrapped_rewards = np.append(rewards, last_value)\n\n    # calculate the returns and advantages\n    if args.gae:\n        bootstrapped_values = np.append(values.detach().cpu().numpy(), last_value)\n        deltas = bootstrapped_rewards[:-1] + args.gamma * bootstrapped_values[1:] * (1-dones) - bootstrapped_values[:-1]\n        advantages = discount_cumsum(deltas, dones, args.gamma * args.gae_lambda)\n        advantages = torch.Tensor(advantages).to(device)\n        returns = advantages + values\n    else:\n        returns = discount_cumsum(bootstrapped_rewards, dones, args.gamma)[:-1]\n        advantages = returns - values.detach().cpu().numpy()\n        advantages = torch.Tensor(advantages).to(device)\n        returns = torch.Tensor(returns).to(device)\n\n    # Advantage normalization\n    if args.norm_adv:\n        EPS = 1e-10\n        advantages = (advantages - advantages.mean()) / (advantages.std() + EPS)\n\n    # Optimizaing policy network\n    entropys = []\n    target_pg = Policy().to(device)\n    inds = np.arange(args.batch_size,)\n    for i_epoch_pi in range(args.update_epochs):\n        np.random.shuffle(inds)\n        for start in range(0, args.batch_size, args.minibatch_size):\n            end = start + args.minibatch_size\n            minibatch_ind = inds[start:end]\n            target_pg.load_state_dict(pg.state_dict())\n            \n            _, newlogproba, _, _ = pg.get_action(\n                obs[minibatch_ind],\n                torch.LongTensor(actions[minibatch_ind].astype(np.int)).to(device).T,\n                invalid_action_masks[minibatch_ind])\n            ratio = (newlogproba - logprobs[:,minibatch_ind]).exp()\n\n            # Policy loss as in OpenAI SpinUp\n            clip_adv = torch.where(advantages[minibatch_ind] > 0,\n                                    (1.+args.clip_coef) * advantages[minibatch_ind],\n                                    (1.-args.clip_coef) * advantages[minibatch_ind]).to(device)\n\n            # Entropy computation with resampled actions\n            entropy = -(newlogproba.exp() * newlogproba).mean()\n            entropys.append(entropy.item())\n\n            policy_loss = -torch.min(ratio * advantages[minibatch_ind], clip_adv) + args.ent_coef * entropy\n            policy_loss = policy_loss.mean()\n            \n            pg_optimizer.zero_grad()\n            policy_loss.backward()\n            nn.utils.clip_grad_norm_(pg.parameters(), args.max_grad_norm)\n            pg_optimizer.step()\n\n            approx_kl = (logprobs[:,minibatch_ind] - newlogproba).mean()\n            # Optimizing value network\n            new_values = vf.forward(obs[minibatch_ind]).view(-1)\n\n            # Value loss clipping\n            if args.clip_vloss:\n                v_loss_unclipped = ((new_values - returns[minibatch_ind]) ** 2)\n                v_clipped = values[minibatch_ind] + torch.clamp(new_values - values[minibatch_ind], -args.clip_coef, args.clip_coef)\n                v_loss_clipped = (v_clipped - returns[minibatch_ind])**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 = torch.mean((returns[minibatch_ind]- new_values).pow(2))\n\n            v_optimizer.zero_grad()\n            v_loss.backward()\n            nn.utils.clip_grad_norm_(vf.parameters(), args.max_grad_norm)\n            v_optimizer.step()\n\n        if args.kle_stop:\n            if approx_kl > args.target_kl:\n                break\n        if args.kle_rollback:\n            if (logprobs[:,minibatch_ind] - \n                pg.get_action(\n                    obs[minibatch_ind],\n                    torch.LongTensor(actions[minibatch_ind].astype(np.int)).to(device).T,\n                    invalid_action_masks[minibatch_ind])[1]).mean() > args.target_kl:\n                pg.load_state_dict(target_pg.state_dict())\n                break\n\n    # TRY NOT TO MODIFY: record rewards for plotting purposes\n    writer.add_scalar(\"losses/value_loss\", v_loss.item(), global_step)\n    writer.add_scalar(\"charts/policy_learning_rate\", pg_optimizer.param_groups[0]['lr'], global_step)\n    writer.add_scalar(\"charts/value_learning_rate\", v_optimizer.param_groups[0]['lr'], global_step)\n    writer.add_scalar(\"losses/policy_loss\", policy_loss.item(), global_step)\n    writer.add_scalar(\"losses/entropy\", np.mean(entropys), global_step)\n    writer.add_scalar(\"losses/approx_kl\", approx_kl.item(), global_step)\n    if args.kle_stop or args.kle_rollback:\n        writer.add_scalar(\"debug/pg_stop_iter\", i_epoch_pi, global_step)\n        \n    # evaluate no mask\n    average_reward, average_invalid_action_stats = evaluate_with_no_mask()\n    writer.add_scalar(\"evals/charts/episode_reward\", average_reward, global_step)\n    print(f\"global_step={global_step}, eval_reward={average_reward}\")\n    for key, idx in zip(info['invalid_action_stats'], range(len(info['invalid_action_stats']))):\n        writer.add_scalar(f\"evals/stats/{key}\", average_invalid_action_stats[idx], global_step)\n\nenv.close()\nwriter.close()\n"
  },
  {
    "path": "invalid_action_masking/ppo_16x16.py",
    "content": "import torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport torch.nn.functional as F\nfrom torch.distributions.categorical import Categorical\nfrom torch.utils.tensorboard import SummaryWriter\n\nfrom cleanrl.common import preprocess_obs_space, preprocess_ac_space\nimport argparse\nimport numpy as np\nimport gym\nimport gym_microrts\nfrom gym.wrappers import TimeLimit, Monitor\nfrom gym.spaces import Discrete, Box, MultiBinary, MultiDiscrete, Space\nimport time\nimport random\nimport os\nimport pandas as pd\n\n# taken from https://github.com/openai/baselines/blob/master/baselines/common/vec_env/vec_normalize.py\nclass RunningMeanStd(object):\n    def __init__(self, epsilon=1e-4, shape=()):\n        self.mean = np.zeros(shape, 'float64')\n        self.var = np.ones(shape, 'float64')\n        self.count = epsilon\n\n    def update(self, x):\n        batch_mean = np.mean([x], axis=0)\n        batch_var = np.var([x], axis=0)\n        batch_count = 1\n        self.update_from_moments(batch_mean, batch_var, batch_count)\n\n    def update_from_moments(self, batch_mean, batch_var, batch_count):\n        self.mean, self.var, self.count = update_mean_var_count_from_moments(\n            self.mean, self.var, self.count, batch_mean, batch_var, batch_count)\n\ndef update_mean_var_count_from_moments(mean, var, count, batch_mean, batch_var, batch_count):\n    delta = batch_mean - mean\n    tot_count = count + batch_count\n\n    new_mean = mean + delta * batch_count / tot_count\n    m_a = var * count\n    m_b = batch_var * batch_count\n    M2 = m_a + m_b + np.square(delta) * count * batch_count / tot_count\n    new_var = M2 / tot_count\n    new_count = tot_count\n\n    return new_mean, new_var, new_count\nclass NormalizedEnv(gym.core.Wrapper):\n    def __init__(self, env, ob=True, ret=True, clipob=10., cliprew=10., gamma=0.99, epsilon=1e-8):\n        super(NormalizedEnv, self).__init__(env)\n        self.ob_rms = RunningMeanStd(shape=self.observation_space.shape) if ob else None\n        self.ret_rms = RunningMeanStd(shape=(1,)) if ret else None\n        self.clipob = clipob\n        self.cliprew = cliprew\n        self.ret = np.zeros(())\n        self.gamma = gamma\n        self.epsilon = epsilon\n\n    def step(self, action):\n        obs, rews, news, infos = self.env.step(action)\n        infos['real_reward'] = rews\n        # print(\"before\", self.ret)\n        self.ret = self.ret * self.gamma + rews\n        # print(\"after\", self.ret)\n        obs = self._obfilt(obs)\n        if self.ret_rms:\n            self.ret_rms.update(np.array([self.ret].copy()))\n            rews = np.clip(rews / np.sqrt(self.ret_rms.var + self.epsilon), -self.cliprew, self.cliprew)\n        self.ret = self.ret * (1-float(news))\n        return obs, rews, news, infos\n\n    def _obfilt(self, obs):\n        if self.ob_rms:\n            self.ob_rms.update(obs)\n            obs = np.clip((obs - self.ob_rms.mean) / np.sqrt(self.ob_rms.var + self.epsilon), -self.clipob, self.clipob)\n            return obs\n        else:\n            return obs\n\n    def reset(self):\n        self.ret = np.zeros(())\n        obs = self.env.reset()\n        return self._obfilt(obs)\n\nif __name__ == \"__main__\":\n    parser = argparse.ArgumentParser(description='PPO agent')\n    # Common arguments\n    parser.add_argument('--exp-name', type=str, default=os.path.basename(__file__).rstrip(\".py\"),\n                       help='the name of this experiment')\n    parser.add_argument('--gym-id', type=str, default=\"MicrortsMining16x16F9-v0\",\n                       help='the id of the gym environment')\n    parser.add_argument('--seed', type=int, default=1,\n                       help='seed of the experiment')\n    parser.add_argument('--episode-length', type=int, default=0,\n                       help='the maximum length of each episode')\n    parser.add_argument('--total-timesteps', type=int, default=100000,\n                       help='total timesteps of the experiments')\n    parser.add_argument('--no-torch-deterministic', action='store_false', dest=\"torch_deterministic\", default=True,\n                       help='if toggled, `torch.backends.cudnn.deterministic=False`')\n    parser.add_argument('--no-cuda', action='store_false', dest=\"cuda\", default=True,\n                       help='if toggled, cuda will not be enabled by default')\n    parser.add_argument('--prod-mode', action='store_true', default=False,\n                       help='run the script in production mode and use wandb to log outputs')\n    parser.add_argument('--capture-video', action='store_true', default=False,\n                       help='weather to capture videos of the agent performances (check out `videos` folder)')\n    parser.add_argument('--wandb-project-name', type=str, default=\"cleanRL\",\n                       help=\"the wandb's project name\")\n    parser.add_argument('--wandb-entity', type=str, default=None,\n                       help=\"the entity (team) of wandb's project\")\n\n    # Algorithm specific arguments\n    parser.add_argument('--batch-size', type=int, default=2048,\n                       help='the batch size of ppo')\n    parser.add_argument('--minibatch-size', type=int, default=256,\n                       help='the mini batch size of ppo')\n    parser.add_argument('--gamma', type=float, default=0.99,\n                       help='the discount factor gamma')\n    parser.add_argument('--gae-lambda', type=float, default=0.97,\n                       help='the lambda for the general advantage estimation')\n    parser.add_argument('--ent-coef', type=float, default=0.01,\n                       help=\"coefficient of the entropy\")\n    parser.add_argument('--max-grad-norm', type=float, default=0.5,\n                       help='the maximum norm for the gradient clipping')\n    parser.add_argument('--clip-coef', type=float, default=0.2,\n                       help=\"the surrogate clipping coefficient\")\n    parser.add_argument('--update-epochs', type=int, default=10,\n                        help=\"the K epochs to update the policy\")\n    parser.add_argument('--kle-stop', action='store_true', default=False,\n                        help='If toggled, the policy updates will be early stopped w.r.t target-kl')\n    parser.add_argument('--kle-rollback', action='store_true', default=False,\n                        help='If toggled, the policy updates will roll back to previous policy if KL exceeds target-kl')\n    parser.add_argument('--target-kl', type=float, default=0.015,\n                        help='the target-kl variable that is referred by --kl')\n    parser.add_argument('--gae', action='store_true', default=True,\n                        help='Use GAE for advantage computation')\n    parser.add_argument('--policy-lr', type=float, default=3e-4,\n                        help=\"the learning rate of the policy optimizer\")\n    parser.add_argument('--value-lr', type=float, default=3e-4,\n                        help=\"the learning rate of the critic optimizer\")\n    parser.add_argument('--norm-obs', action='store_true', default=True,\n                        help=\"Toggles observation normalization\")\n    parser.add_argument('--norm-returns', action='store_true', default=False,\n                        help=\"Toggles returns normalization\")\n    parser.add_argument('--norm-adv', action='store_true', default=True,\n                        help=\"Toggles advantages normalization\")\n    parser.add_argument('--obs-clip', type=float, default=10.0,\n                        help=\"Value for reward clipping, as per the paper\")\n    parser.add_argument('--rew-clip', type=float, default=10.0,\n                        help=\"Value for observation clipping, as per the paper\")\n    parser.add_argument('--anneal-lr', action='store_true', default=True,\n                        help=\"Toggle learning rate annealing for policy and value networks\")\n    parser.add_argument('--weights-init', default=\"orthogonal\", choices=[\"xavier\", 'orthogonal'],\n                        help='Selects the scheme to be used for weights initialization'),\n    parser.add_argument('--clip-vloss', action=\"store_true\", default=True,\n                        help='Toggles wheter or not to use a clipped loss for the value function, as per the paper.')\n    parser.add_argument('--pol-layer-norm', action='store_true', default=False,\n                       help='Enables layer normalization in the policy network')\n    args = parser.parse_args()\n    if not args.seed:\n        args.seed = int(time.time())\n\nargs.features_turned_on = sum([args.kle_stop, args.kle_rollback, args.gae, args.norm_obs, args.norm_returns, args.norm_adv, args.anneal_lr, args.clip_vloss, args.pol_layer_norm])\n# TRY NOT TO MODIFY: setup the environment\nexperiment_name = f\"{args.gym_id}__{args.exp_name}__{args.seed}__{int(time.time())}\"\nwriter = SummaryWriter(f\"runs/{experiment_name}\")\nwriter.add_text('hyperparameters', \"|param|value|\\n|-|-|\\n%s\" % (\n        '\\n'.join([f\"|{key}|{value}|\" for key, value in vars(args).items()])))\n\nif args.prod_mode:\n    import wandb\n    wandb.init(project=args.wandb_project_name, entity=args.wandb_entity, tensorboard=True, config=vars(args), name=experiment_name, monitor_gym=True)\n    writer = SummaryWriter(f\"/tmp/{experiment_name}\")\n    wandb.save(os.path.abspath(__file__))\n\n# TRY NOT TO MODIFY: seeding\ndevice = torch.device('cuda' if torch.cuda.is_available() and args.cuda else 'cpu')\nenv = gym.make(args.gym_id)\n# respect the default timelimit\nassert isinstance(env.action_space, MultiDiscrete), \"only MultiDiscrete action space is supported\"\nassert isinstance(env, TimeLimit) or int(args.episode_length), \"the gym env does not have a built in TimeLimit, please specify by using --episode-length\"\nif isinstance(env, TimeLimit):\n    if int(args.episode_length):\n        env._max_episode_steps = int(args.episode_length)\n    args.episode_length = env._max_episode_steps\nelse:\n    env = TimeLimit(env, int(args.episode_length))\nenv = NormalizedEnv(env.env, ob=args.norm_obs, ret=args.norm_returns, clipob=args.obs_clip, cliprew=args.rew_clip, gamma=args.gamma)\nenv = TimeLimit(env, int(args.episode_length))\nrandom.seed(args.seed)\nnp.random.seed(args.seed)\ntorch.manual_seed(args.seed)\ntorch.backends.cudnn.deterministic = args.torch_deterministic\nenv.seed(args.seed)\nenv.action_space.seed(args.seed)\nenv.observation_space.seed(args.seed)\nif args.capture_video:\n    env = Monitor(env, f'videos/{experiment_name}')\n\n# ALGO LOGIC: initialize agent here:\nclass CategoricalMasked(Categorical):\n    def __init__(self, probs=None, logits=None, validate_args=None, masks=[]):\n        self.masks = masks\n        if len(self.masks) == 0:\n            super(CategoricalMasked, self).__init__(probs, logits, validate_args)\n        else:\n            self.masks = masks.type(torch.BoolTensor).to(device)\n            logits = torch.where(self.masks, logits, torch.tensor(-1e+8).to(device))\n            super(CategoricalMasked, self).__init__(probs, logits, validate_args)\n    \n    def entropy(self):\n        if len(self.masks) == 0:\n            return super(CategoricalMasked, self).entropy()\n        p_log_p = self.logits * self.probs\n        p_log_p = torch.where(self.masks, p_log_p, torch.tensor(0.).to(device))\n        return -p_log_p.sum(-1)\n\nclass Policy(nn.Module):\n    def __init__(self):\n        super(Policy, self).__init__()\n        self.features = nn.Sequential(\n            nn.Conv2d(27, 16, kernel_size=3),\n            nn.MaxPool2d(1),\n            nn.ReLU(),\n            nn.Conv2d(16, 32, kernel_size=3),\n            nn.MaxPool2d(1),\n            nn.ReLU())\n        self.fc = nn.Sequential(\n            nn.Linear(32*12*12, 128),\n            nn.ReLU(),\n            nn.Linear(128, env.action_space.nvec.sum())\n        )\n\n    def forward(self, x):\n        x = torch.Tensor(np.moveaxis(x, -1, 1)).to(device)\n        x = self.features(x)\n        x = x.reshape(x.size(0), -1)\n        x = self.fc(x)\n        return x\n\n    def get_action(self, x, action=None, invalid_action_masks=None):\n        logits = self.forward(x)\n        split_logits = torch.split(logits, env.action_space.nvec.tolist(), dim=1)\n        \n        if invalid_action_masks is not None:\n            split_invalid_action_masks = torch.split(invalid_action_masks, env.action_space.nvec.tolist(), dim=1)\n            multi_categoricals = [CategoricalMasked(logits=logits, masks=iam) for (logits, iam) in zip(split_logits, split_invalid_action_masks)]\n        else:\n            multi_categoricals = [Categorical(logits=logits) for logits in split_logits]\n        \n        if action is None:\n            action = torch.stack([categorical.sample() for categorical in multi_categoricals])\n        logprob = torch.stack([categorical.log_prob(a) for a, categorical in zip(action, multi_categoricals)])\n        return action, logprob, [], multi_categoricals\n\nclass Value(nn.Module):\n    def __init__(self):\n        super(Value, self).__init__()\n        self.features = nn.Sequential(\n            nn.Conv2d(27, 16, kernel_size=3),\n            nn.MaxPool2d(1),\n            nn.ReLU(),\n            nn.Conv2d(16, 32, kernel_size=3),\n            nn.MaxPool2d(1),\n            nn.ReLU())\n        self.fc = nn.Sequential(\n            nn.Linear(32*12*12, 128),\n            nn.ReLU(),\n            nn.Linear(128, 1)\n        )\n\n    def forward(self, x):\n        x = torch.Tensor(np.moveaxis(x, -1, 1)).to(device)\n        x = self.features(x)\n        x = x.reshape(x.size(0), -1)\n        x = self.fc(x)\n        return x\n\ndef discount_cumsum(x, dones, gamma):\n    \"\"\"\n    computing discounted cumulative sums of vectors that resets with dones\n    input:\n        vector x,  vector dones,\n        [x0,       [0,\n         x1,        0,\n         x2         1,\n         x3         0, \n         x4]        0]\n    output:\n        [x0 + discount * x1 + discount^2 * x2,\n         x1 + discount * x2,\n         x2,\n         x3 + discount * x4,\n         x4]\n    \"\"\"\n    discount_cumsum = np.zeros_like(x)\n    discount_cumsum[-1] = x[-1]\n    for t in reversed(range(x.shape[0]-1)):\n        discount_cumsum[t] = x[t] + gamma * discount_cumsum[t+1] * (1-dones[t])\n    return discount_cumsum\n\npg = Policy().to(device)\nvf = Value().to(device)\n\n# MODIFIED: Separate optimizer and learning rates\npg_optimizer = optim.Adam(list(pg.parameters()), lr=args.policy_lr)\nv_optimizer = optim.Adam(list(vf.parameters()), lr=args.value_lr)\n\n# MODIFIED: Initializing learning rate anneal scheduler when need\nif args.anneal_lr:\n    anneal_fn = lambda f: max(0, 1-f / args.total_timesteps)\n    pg_lr_scheduler = optim.lr_scheduler.LambdaLR(pg_optimizer, lr_lambda=anneal_fn)\n    vf_lr_scheduler = optim.lr_scheduler.LambdaLR(v_optimizer, lr_lambda=anneal_fn)\n\nloss_fn = nn.MSELoss()\n\ndef evaluate_with_no_mask():\n    evaluate_rewards = []\n    evaluate_invalid_action_stats = []\n    if args.capture_video:\n        env.stats_recorder.done=True\n    next_obs = np.array(env.reset())\n\n    # ALGO Logic: Storage for epoch data\n    obs = np.empty((args.batch_size,) + env.observation_space.shape)\n\n    actions = np.empty((args.batch_size,) + env.action_space.shape)\n    logprobs = torch.zeros((env.action_space.nvec.shape[0], args.batch_size,)).to(device)\n\n    rewards = np.zeros((args.batch_size,))\n    raw_rewards = np.zeros((len(env.rfs),args.batch_size,))\n    \n    real_rewards = []\n    invalid_action_stats = []\n\n    dones = np.zeros((args.batch_size,))\n    values = torch.zeros((args.batch_size,)).to(device)\n\n    # TRY NOT TO MODIFY: prepare the execution of the game.\n    for step in range(args.batch_size):\n        env.render()\n        obs[step] = next_obs.copy()\n\n        # ALGO LOGIC: put action logic here\n        with torch.no_grad():\n            values[step] = vf.forward(obs[step:step+1])\n            action, logproba, _, probs = pg.get_action(obs[step:step+1])\n        \n        actions[step] = action[:,0].data.cpu().numpy()\n        logprobs[:,[step]] = logproba\n\n        # TRY NOT TO MODIFY: execute the game and log data.\n        next_obs, rewards[step], dones[step], info = env.step(action[:,0].data.cpu().numpy())\n        raw_rewards[:,step] = info[\"rewards\"]\n        real_rewards += [info['real_reward']]\n        invalid_action_stats += [info['invalid_action_stats']]\n        next_obs = np.array(next_obs)\n\n        if dones[step]:\n            # Computing the discounted returns:\n            writer.add_scalar(\"charts/episode_reward\", np.sum(real_rewards), global_step)\n            evaluate_rewards += [np.sum(real_rewards)]\n            real_rewards = []\n            evaluate_invalid_action_stats += [pd.DataFrame(invalid_action_stats).sum(0)]\n            invalid_action_stats = []\n            next_obs = np.array(env.reset())\n    \n    return np.average(evaluate_rewards), pd.DataFrame(evaluate_invalid_action_stats).mean(0)\n\n# TRY NOT TO MODIFY: start the game\nglobal_step = 0\nwhile global_step < args.total_timesteps:\n    if args.capture_video:\n        env.stats_recorder.done=True\n    next_obs = np.array(env.reset())\n\n    # ALGO Logic: Storage for epoch data\n    obs = np.empty((args.batch_size,) + env.observation_space.shape)\n\n    actions = np.empty((args.batch_size,) + env.action_space.shape)\n    logprobs = torch.zeros((env.action_space.nvec.shape[0], args.batch_size,)).to(device)\n\n    rewards = np.zeros((args.batch_size,))\n    raw_rewards = np.zeros((len(env.rfs),args.batch_size,))\n    \n    real_rewards = []\n    invalid_action_stats = []\n\n    dones = np.zeros((args.batch_size,))\n    values = torch.zeros((args.batch_size,)).to(device)\n\n    invalid_action_masks = torch.zeros((args.batch_size, env.action_space.nvec.sum()))\n    # TRY NOT TO MODIFY: prepare the execution of the game.\n    for step in range(args.batch_size):\n        env.render()\n        global_step += 1\n        obs[step] = next_obs.copy()\n\n        # ALGO LOGIC: put action logic here\n        invalid_action_mask = torch.ones(env.action_space.nvec.sum())\n        invalid_action_mask[0:env.action_space.nvec[0]] = torch.tensor(env.unit_location_mask)\n        invalid_action_mask[-env.action_space.nvec[-1]:] = torch.tensor(env.target_unit_location_mask)\n        invalid_action_masks[step] = invalid_action_mask\n        with torch.no_grad():\n            values[step] = vf.forward(obs[step:step+1])\n            action, logproba, _, probs = pg.get_action(obs[step:step+1], invalid_action_masks=invalid_action_masks[step:step+1])\n        \n        actions[step] = action[:,0].data.cpu().numpy()\n        logprobs[:,[step]] = logproba\n\n        # TRY NOT TO MODIFY: execute the game and log data.\n        next_obs, rewards[step], dones[step], info = env.step(action[:,0].data.cpu().numpy())\n        raw_rewards[:,step] = info[\"rewards\"]\n        real_rewards += [info['real_reward']]\n        invalid_action_stats += [info['invalid_action_stats']]\n        next_obs = np.array(next_obs)\n\n        # Annealing the rate if instructed to do so.\n        if args.anneal_lr:\n            pg_lr_scheduler.step()\n            vf_lr_scheduler.step()\n\n        if dones[step]:\n            # Computing the discounted returns:\n            writer.add_scalar(\"charts/episode_reward\", np.sum(real_rewards), global_step)\n            print(f\"global_step={global_step}, episode_reward={np.sum(real_rewards)}\")\n            for i in range(len(env.rfs)):\n                writer.add_scalar(f\"charts/episode_reward/{str(env.rfs[i])}\", raw_rewards.sum(1)[i], global_step)\n            real_rewards = []\n            for key, idx in zip(info['invalid_action_stats'], range(len(info['invalid_action_stats']))):\n                writer.add_scalar(f\"stats/{key}\", pd.DataFrame(invalid_action_stats).sum(0)[idx], global_step)\n            invalid_action_stats = []\n            next_obs = np.array(env.reset())\n\n    # bootstrap reward if not done. reached the batch limit\n    last_value = 0\n    if not dones[step]:\n        last_value = vf.forward(next_obs.reshape((1,)+next_obs.shape))[0].detach().cpu().numpy()[0]\n    bootstrapped_rewards = np.append(rewards, last_value)\n\n    # calculate the returns and advantages\n    if args.gae:\n        bootstrapped_values = np.append(values.detach().cpu().numpy(), last_value)\n        deltas = bootstrapped_rewards[:-1] + args.gamma * bootstrapped_values[1:] * (1-dones) - bootstrapped_values[:-1]\n        advantages = discount_cumsum(deltas, dones, args.gamma * args.gae_lambda)\n        advantages = torch.Tensor(advantages).to(device)\n        returns = advantages + values\n    else:\n        returns = discount_cumsum(bootstrapped_rewards, dones, args.gamma)[:-1]\n        advantages = returns - values.detach().cpu().numpy()\n        advantages = torch.Tensor(advantages).to(device)\n        returns = torch.Tensor(returns).to(device)\n\n    # Advantage normalization\n    if args.norm_adv:\n        EPS = 1e-10\n        advantages = (advantages - advantages.mean()) / (advantages.std() + EPS)\n\n    # Optimizaing policy network\n    entropys = []\n    target_pg = Policy().to(device)\n    inds = np.arange(args.batch_size,)\n    for i_epoch_pi in range(args.update_epochs):\n        np.random.shuffle(inds)\n        for start in range(0, args.batch_size, args.minibatch_size):\n            end = start + args.minibatch_size\n            minibatch_ind = inds[start:end]\n            target_pg.load_state_dict(pg.state_dict())\n            \n            _, newlogproba, _, _ = pg.get_action(\n                obs[minibatch_ind],\n                torch.LongTensor(actions[minibatch_ind].astype(np.int)).to(device).T,\n                invalid_action_masks[minibatch_ind])\n            ratio = (newlogproba - logprobs[:,minibatch_ind]).exp()\n\n            # Policy loss as in OpenAI SpinUp\n            clip_adv = torch.where(advantages[minibatch_ind] > 0,\n                                    (1.+args.clip_coef) * advantages[minibatch_ind],\n                                    (1.-args.clip_coef) * advantages[minibatch_ind]).to(device)\n\n            # Entropy computation with resampled actions\n            entropy = -(newlogproba.exp() * newlogproba).mean()\n            entropys.append(entropy.item())\n\n            policy_loss = -torch.min(ratio * advantages[minibatch_ind], clip_adv) + args.ent_coef * entropy\n            policy_loss = policy_loss.mean()\n            \n            pg_optimizer.zero_grad()\n            policy_loss.backward()\n            nn.utils.clip_grad_norm_(pg.parameters(), args.max_grad_norm)\n            pg_optimizer.step()\n\n            approx_kl = (logprobs[:,minibatch_ind] - newlogproba).mean()\n            # Optimizing value network\n            new_values = vf.forward(obs[minibatch_ind]).view(-1)\n\n            # Value loss clipping\n            if args.clip_vloss:\n                v_loss_unclipped = ((new_values - returns[minibatch_ind]) ** 2)\n                v_clipped = values[minibatch_ind] + torch.clamp(new_values - values[minibatch_ind], -args.clip_coef, args.clip_coef)\n                v_loss_clipped = (v_clipped - returns[minibatch_ind])**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 = torch.mean((returns[minibatch_ind]- new_values).pow(2))\n\n            v_optimizer.zero_grad()\n            v_loss.backward()\n            nn.utils.clip_grad_norm_(vf.parameters(), args.max_grad_norm)\n            v_optimizer.step()\n\n        if args.kle_stop:\n            if approx_kl > args.target_kl:\n                break\n        if args.kle_rollback:\n            if (logprobs[:,minibatch_ind] - \n                pg.get_action(\n                    obs[minibatch_ind],\n                    torch.LongTensor(actions[minibatch_ind].astype(np.int)).to(device).T,\n                    invalid_action_masks[minibatch_ind])[1]).mean() > args.target_kl:\n                pg.load_state_dict(target_pg.state_dict())\n                break\n\n    # TRY NOT TO MODIFY: record rewards for plotting purposes\n    writer.add_scalar(\"losses/value_loss\", v_loss.item(), global_step)\n    writer.add_scalar(\"charts/policy_learning_rate\", pg_optimizer.param_groups[0]['lr'], global_step)\n    writer.add_scalar(\"charts/value_learning_rate\", v_optimizer.param_groups[0]['lr'], global_step)\n    writer.add_scalar(\"losses/policy_loss\", policy_loss.item(), global_step)\n    writer.add_scalar(\"losses/entropy\", np.mean(entropys), global_step)\n    writer.add_scalar(\"losses/approx_kl\", approx_kl.item(), global_step)\n    if args.kle_stop or args.kle_rollback:\n        writer.add_scalar(\"debug/pg_stop_iter\", i_epoch_pi, global_step)\n        \n    # evaluate no mask\n    average_reward, average_invalid_action_stats = evaluate_with_no_mask()\n    writer.add_scalar(\"evals/charts/episode_reward\", average_reward, global_step)\n    print(f\"global_step={global_step}, eval_reward={average_reward}\")\n    for key, idx in zip(info['invalid_action_stats'], range(len(info['invalid_action_stats']))):\n        writer.add_scalar(f\"evals/stats/{key}\", average_invalid_action_stats[idx], global_step)\n\nenv.close()\nwriter.close()\n"
  },
  {
    "path": "invalid_action_masking/ppo_24x24.py",
    "content": "import torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport torch.nn.functional as F\nfrom torch.distributions.categorical import Categorical\nfrom torch.utils.tensorboard import SummaryWriter\n\nfrom cleanrl.common import preprocess_obs_space, preprocess_ac_space\nimport argparse\nimport numpy as np\nimport gym\nimport gym_microrts\nfrom gym.wrappers import TimeLimit, Monitor\nfrom gym.spaces import Discrete, Box, MultiBinary, MultiDiscrete, Space\nimport time\nimport random\nimport os\nimport pandas as pd\n\n# taken from https://github.com/openai/baselines/blob/master/baselines/common/vec_env/vec_normalize.py\nclass RunningMeanStd(object):\n    def __init__(self, epsilon=1e-4, shape=()):\n        self.mean = np.zeros(shape, 'float64')\n        self.var = np.ones(shape, 'float64')\n        self.count = epsilon\n\n    def update(self, x):\n        batch_mean = np.mean([x], axis=0)\n        batch_var = np.var([x], axis=0)\n        batch_count = 1\n        self.update_from_moments(batch_mean, batch_var, batch_count)\n\n    def update_from_moments(self, batch_mean, batch_var, batch_count):\n        self.mean, self.var, self.count = update_mean_var_count_from_moments(\n            self.mean, self.var, self.count, batch_mean, batch_var, batch_count)\n\ndef update_mean_var_count_from_moments(mean, var, count, batch_mean, batch_var, batch_count):\n    delta = batch_mean - mean\n    tot_count = count + batch_count\n\n    new_mean = mean + delta * batch_count / tot_count\n    m_a = var * count\n    m_b = batch_var * batch_count\n    M2 = m_a + m_b + np.square(delta) * count * batch_count / tot_count\n    new_var = M2 / tot_count\n    new_count = tot_count\n\n    return new_mean, new_var, new_count\nclass NormalizedEnv(gym.core.Wrapper):\n    def __init__(self, env, ob=True, ret=True, clipob=10., cliprew=10., gamma=0.99, epsilon=1e-8):\n        super(NormalizedEnv, self).__init__(env)\n        self.ob_rms = RunningMeanStd(shape=self.observation_space.shape) if ob else None\n        self.ret_rms = RunningMeanStd(shape=(1,)) if ret else None\n        self.clipob = clipob\n        self.cliprew = cliprew\n        self.ret = np.zeros(())\n        self.gamma = gamma\n        self.epsilon = epsilon\n\n    def step(self, action):\n        obs, rews, news, infos = self.env.step(action)\n        infos['real_reward'] = rews\n        # print(\"before\", self.ret)\n        self.ret = self.ret * self.gamma + rews\n        # print(\"after\", self.ret)\n        obs = self._obfilt(obs)\n        if self.ret_rms:\n            self.ret_rms.update(np.array([self.ret].copy()))\n            rews = np.clip(rews / np.sqrt(self.ret_rms.var + self.epsilon), -self.cliprew, self.cliprew)\n        self.ret = self.ret * (1-float(news))\n        return obs, rews, news, infos\n\n    def _obfilt(self, obs):\n        if self.ob_rms:\n            self.ob_rms.update(obs)\n            obs = np.clip((obs - self.ob_rms.mean) / np.sqrt(self.ob_rms.var + self.epsilon), -self.clipob, self.clipob)\n            return obs\n        else:\n            return obs\n\n    def reset(self):\n        self.ret = np.zeros(())\n        obs = self.env.reset()\n        return self._obfilt(obs)\n\nif __name__ == \"__main__\":\n    parser = argparse.ArgumentParser(description='PPO agent')\n    # Common arguments\n    parser.add_argument('--exp-name', type=str, default=os.path.basename(__file__).rstrip(\".py\"),\n                       help='the name of this experiment')\n    parser.add_argument('--gym-id', type=str, default=\"MicrortsMining24x24F9-v0\",\n                       help='the id of the gym environment')\n    parser.add_argument('--seed', type=int, default=1,\n                       help='seed of the experiment')\n    parser.add_argument('--episode-length', type=int, default=0,\n                       help='the maximum length of each episode')\n    parser.add_argument('--total-timesteps', type=int, default=100000,\n                       help='total timesteps of the experiments')\n    parser.add_argument('--no-torch-deterministic', action='store_false', dest=\"torch_deterministic\", default=True,\n                       help='if toggled, `torch.backends.cudnn.deterministic=False`')\n    parser.add_argument('--no-cuda', action='store_false', dest=\"cuda\", default=True,\n                       help='if toggled, cuda will not be enabled by default')\n    parser.add_argument('--prod-mode', action='store_true', default=False,\n                       help='run the script in production mode and use wandb to log outputs')\n    parser.add_argument('--capture-video', action='store_true', default=False,\n                       help='weather to capture videos of the agent performances (check out `videos` folder)')\n    parser.add_argument('--wandb-project-name', type=str, default=\"cleanRL\",\n                       help=\"the wandb's project name\")\n    parser.add_argument('--wandb-entity', type=str, default=None,\n                       help=\"the entity (team) of wandb's project\")\n\n    # Algorithm specific arguments\n    parser.add_argument('--batch-size', type=int, default=2048,\n                       help='the batch size of ppo')\n    parser.add_argument('--minibatch-size', type=int, default=256,\n                       help='the mini batch size of ppo')\n    parser.add_argument('--gamma', type=float, default=0.99,\n                       help='the discount factor gamma')\n    parser.add_argument('--gae-lambda', type=float, default=0.97,\n                       help='the lambda for the general advantage estimation')\n    parser.add_argument('--ent-coef', type=float, default=0.01,\n                       help=\"coefficient of the entropy\")\n    parser.add_argument('--max-grad-norm', type=float, default=0.5,\n                       help='the maximum norm for the gradient clipping')\n    parser.add_argument('--clip-coef', type=float, default=0.2,\n                       help=\"the surrogate clipping coefficient\")\n    parser.add_argument('--update-epochs', type=int, default=10,\n                        help=\"the K epochs to update the policy\")\n    parser.add_argument('--kle-stop', action='store_true', default=False,\n                        help='If toggled, the policy updates will be early stopped w.r.t target-kl')\n    parser.add_argument('--kle-rollback', action='store_true', default=False,\n                        help='If toggled, the policy updates will roll back to previous policy if KL exceeds target-kl')\n    parser.add_argument('--target-kl', type=float, default=0.015,\n                        help='the target-kl variable that is referred by --kl')\n    parser.add_argument('--gae', action='store_true', default=True,\n                        help='Use GAE for advantage computation')\n    parser.add_argument('--policy-lr', type=float, default=3e-4,\n                        help=\"the learning rate of the policy optimizer\")\n    parser.add_argument('--value-lr', type=float, default=3e-4,\n                        help=\"the learning rate of the critic optimizer\")\n    parser.add_argument('--norm-obs', action='store_true', default=True,\n                        help=\"Toggles observation normalization\")\n    parser.add_argument('--norm-returns', action='store_true', default=False,\n                        help=\"Toggles returns normalization\")\n    parser.add_argument('--norm-adv', action='store_true', default=True,\n                        help=\"Toggles advantages normalization\")\n    parser.add_argument('--obs-clip', type=float, default=10.0,\n                        help=\"Value for reward clipping, as per the paper\")\n    parser.add_argument('--rew-clip', type=float, default=10.0,\n                        help=\"Value for observation clipping, as per the paper\")\n    parser.add_argument('--anneal-lr', action='store_true', default=True,\n                        help=\"Toggle learning rate annealing for policy and value networks\")\n    parser.add_argument('--weights-init', default=\"orthogonal\", choices=[\"xavier\", 'orthogonal'],\n                        help='Selects the scheme to be used for weights initialization'),\n    parser.add_argument('--clip-vloss', action=\"store_true\", default=True,\n                        help='Toggles wheter or not to use a clipped loss for the value function, as per the paper.')\n    parser.add_argument('--pol-layer-norm', action='store_true', default=False,\n                       help='Enables layer normalization in the policy network')\n    args = parser.parse_args()\n    if not args.seed:\n        args.seed = int(time.time())\n\nargs.features_turned_on = sum([args.kle_stop, args.kle_rollback, args.gae, args.norm_obs, args.norm_returns, args.norm_adv, args.anneal_lr, args.clip_vloss, args.pol_layer_norm])\n# TRY NOT TO MODIFY: setup the environment\nexperiment_name = f\"{args.gym_id}__{args.exp_name}__{args.seed}__{int(time.time())}\"\nwriter = SummaryWriter(f\"runs/{experiment_name}\")\nwriter.add_text('hyperparameters', \"|param|value|\\n|-|-|\\n%s\" % (\n        '\\n'.join([f\"|{key}|{value}|\" for key, value in vars(args).items()])))\n\nif args.prod_mode:\n    import wandb\n    wandb.init(project=args.wandb_project_name, entity=args.wandb_entity, tensorboard=True, config=vars(args), name=experiment_name, monitor_gym=True)\n    writer = SummaryWriter(f\"/tmp/{experiment_name}\")\n    wandb.save(os.path.abspath(__file__))\n\n# TRY NOT TO MODIFY: seeding\ndevice = torch.device('cuda' if torch.cuda.is_available() and args.cuda else 'cpu')\nenv = gym.make(args.gym_id)\n# respect the default timelimit\nassert isinstance(env.action_space, MultiDiscrete), \"only MultiDiscrete action space is supported\"\nassert isinstance(env, TimeLimit) or int(args.episode_length), \"the gym env does not have a built in TimeLimit, please specify by using --episode-length\"\nif isinstance(env, TimeLimit):\n    if int(args.episode_length):\n        env._max_episode_steps = int(args.episode_length)\n    args.episode_length = env._max_episode_steps\nelse:\n    env = TimeLimit(env, int(args.episode_length))\nenv = NormalizedEnv(env.env, ob=args.norm_obs, ret=args.norm_returns, clipob=args.obs_clip, cliprew=args.rew_clip, gamma=args.gamma)\nenv = TimeLimit(env, int(args.episode_length))\nrandom.seed(args.seed)\nnp.random.seed(args.seed)\ntorch.manual_seed(args.seed)\ntorch.backends.cudnn.deterministic = args.torch_deterministic\nenv.seed(args.seed)\nenv.action_space.seed(args.seed)\nenv.observation_space.seed(args.seed)\nif args.capture_video:\n    env = Monitor(env, f'videos/{experiment_name}')\n\n# ALGO LOGIC: initialize agent here:\nclass CategoricalMasked(Categorical):\n    def __init__(self, probs=None, logits=None, validate_args=None, masks=[]):\n        self.masks = masks\n        if len(self.masks) == 0:\n            super(CategoricalMasked, self).__init__(probs, logits, validate_args)\n        else:\n            self.masks = masks.type(torch.BoolTensor).to(device)\n            logits = torch.where(self.masks, logits, torch.tensor(-1e+8).to(device))\n            super(CategoricalMasked, self).__init__(probs, logits, validate_args)\n    \n    def entropy(self):\n        if len(self.masks) == 0:\n            return super(CategoricalMasked, self).entropy()\n        p_log_p = self.logits * self.probs\n        p_log_p = torch.where(self.masks, p_log_p, torch.tensor(0.).to(device))\n        return -p_log_p.sum(-1)\n\nclass Policy(nn.Module):\n    def __init__(self):\n        super(Policy, self).__init__()\n        self.features = nn.Sequential(\n            nn.Conv2d(27, 16, kernel_size=3, stride=1),\n            nn.MaxPool2d(2),\n            nn.ReLU(),\n            nn.Conv2d(16, 32, kernel_size=2, stride=1),\n            nn.MaxPool2d(2),\n            nn.ReLU())\n        self.fc = nn.Sequential(\n            nn.Linear(32*5*5, 128),\n            nn.ReLU(),\n            nn.Linear(128, env.action_space.nvec.sum())\n        )\n\n    def forward(self, x):\n        x = torch.Tensor(np.moveaxis(x, -1, 1)).to(device)\n        x = self.features(x)\n        x = x.reshape(x.size(0), -1)\n        x = self.fc(x)\n        return x\n\n    def get_action(self, x, action=None, invalid_action_masks=None):\n        logits = self.forward(x)\n        split_logits = torch.split(logits, env.action_space.nvec.tolist(), dim=1)\n        \n        if invalid_action_masks is not None:\n            split_invalid_action_masks = torch.split(invalid_action_masks, env.action_space.nvec.tolist(), dim=1)\n            multi_categoricals = [CategoricalMasked(logits=logits, masks=iam) for (logits, iam) in zip(split_logits, split_invalid_action_masks)]\n        else:\n            multi_categoricals = [Categorical(logits=logits) for logits in split_logits]\n        \n        if action is None:\n            action = torch.stack([categorical.sample() for categorical in multi_categoricals])\n        logprob = torch.stack([categorical.log_prob(a) for a, categorical in zip(action, multi_categoricals)])\n        return action, logprob, [], multi_categoricals\n\nclass Value(nn.Module):\n    def __init__(self):\n        super(Value, self).__init__()\n        self.features = nn.Sequential(\n            nn.Conv2d(27, 16, kernel_size=3, stride=1),\n            nn.MaxPool2d(2),\n            nn.ReLU(),\n            nn.Conv2d(16, 32, kernel_size=2, stride=1),\n            nn.MaxPool2d(2),\n            nn.ReLU())\n        self.fc = nn.Sequential(\n            nn.Linear(32*5*5, 128),\n            nn.ReLU(),\n            nn.Linear(128, 1)\n        )\n\n    def forward(self, x):\n        x = torch.Tensor(np.moveaxis(x, -1, 1)).to(device)\n        x = self.features(x)\n        x = x.reshape(x.size(0), -1)\n        x = self.fc(x)\n        return x\n\ndef discount_cumsum(x, dones, gamma):\n    \"\"\"\n    computing discounted cumulative sums of vectors that resets with dones\n    input:\n        vector x,  vector dones,\n        [x0,       [0,\n         x1,        0,\n         x2         1,\n         x3         0, \n         x4]        0]\n    output:\n        [x0 + discount * x1 + discount^2 * x2,\n         x1 + discount * x2,\n         x2,\n         x3 + discount * x4,\n         x4]\n    \"\"\"\n    discount_cumsum = np.zeros_like(x)\n    discount_cumsum[-1] = x[-1]\n    for t in reversed(range(x.shape[0]-1)):\n        discount_cumsum[t] = x[t] + gamma * discount_cumsum[t+1] * (1-dones[t])\n    return discount_cumsum\n\npg = Policy().to(device)\nvf = Value().to(device)\n\n# MODIFIED: Separate optimizer and learning rates\npg_optimizer = optim.Adam(list(pg.parameters()), lr=args.policy_lr)\nv_optimizer = optim.Adam(list(vf.parameters()), lr=args.value_lr)\n\n# MODIFIED: Initializing learning rate anneal scheduler when need\nif args.anneal_lr:\n    anneal_fn = lambda f: max(0, 1-f / args.total_timesteps)\n    pg_lr_scheduler = optim.lr_scheduler.LambdaLR(pg_optimizer, lr_lambda=anneal_fn)\n    vf_lr_scheduler = optim.lr_scheduler.LambdaLR(v_optimizer, lr_lambda=anneal_fn)\n\nloss_fn = nn.MSELoss()\n\ndef evaluate_with_no_mask():\n    evaluate_rewards = []\n    evaluate_invalid_action_stats = []\n    if args.capture_video:\n        env.stats_recorder.done=True\n    next_obs = np.array(env.reset())\n\n    # ALGO Logic: Storage for epoch data\n    obs = np.empty((args.batch_size,) + env.observation_space.shape)\n\n    actions = np.empty((args.batch_size,) + env.action_space.shape)\n    logprobs = torch.zeros((env.action_space.nvec.shape[0], args.batch_size,)).to(device)\n\n    rewards = np.zeros((args.batch_size,))\n    raw_rewards = np.zeros((len(env.rfs),args.batch_size,))\n    \n    real_rewards = []\n    invalid_action_stats = []\n\n    dones = np.zeros((args.batch_size,))\n    values = torch.zeros((args.batch_size,)).to(device)\n\n    # TRY NOT TO MODIFY: prepare the execution of the game.\n    for step in range(args.batch_size):\n        env.render()\n        obs[step] = next_obs.copy()\n\n        # ALGO LOGIC: put action logic here\n        with torch.no_grad():\n            values[step] = vf.forward(obs[step:step+1])\n            action, logproba, _, probs = pg.get_action(obs[step:step+1])\n        \n        actions[step] = action[:,0].data.cpu().numpy()\n        logprobs[:,[step]] = logproba\n\n        # TRY NOT TO MODIFY: execute the game and log data.\n        next_obs, rewards[step], dones[step], info = env.step(action[:,0].data.cpu().numpy())\n        raw_rewards[:,step] = info[\"rewards\"]\n        real_rewards += [info['real_reward']]\n        invalid_action_stats += [info['invalid_action_stats']]\n        next_obs = np.array(next_obs)\n\n        if dones[step]:\n            # Computing the discounted returns:\n            writer.add_scalar(\"charts/episode_reward\", np.sum(real_rewards), global_step)\n            evaluate_rewards += [np.sum(real_rewards)]\n            real_rewards = []\n            evaluate_invalid_action_stats += [pd.DataFrame(invalid_action_stats).sum(0)]\n            invalid_action_stats = []\n            next_obs = np.array(env.reset())\n    \n    return np.average(evaluate_rewards), pd.DataFrame(evaluate_invalid_action_stats).mean(0)\n\n# TRY NOT TO MODIFY: start the game\nglobal_step = 0\nwhile global_step < args.total_timesteps:\n    if args.capture_video:\n        env.stats_recorder.done=True\n    next_obs = np.array(env.reset())\n\n    # ALGO Logic: Storage for epoch data\n    obs = np.empty((args.batch_size,) + env.observation_space.shape)\n\n    actions = np.empty((args.batch_size,) + env.action_space.shape)\n    logprobs = torch.zeros((env.action_space.nvec.shape[0], args.batch_size,)).to(device)\n\n    rewards = np.zeros((args.batch_size,))\n    raw_rewards = np.zeros((len(env.rfs),args.batch_size,))\n    \n    real_rewards = []\n    invalid_action_stats = []\n\n    dones = np.zeros((args.batch_size,))\n    values = torch.zeros((args.batch_size,)).to(device)\n\n    invalid_action_masks = torch.zeros((args.batch_size, env.action_space.nvec.sum()))\n    # TRY NOT TO MODIFY: prepare the execution of the game.\n    for step in range(args.batch_size):\n        env.render()\n        global_step += 1\n        obs[step] = next_obs.copy()\n\n        # ALGO LOGIC: put action logic here\n        invalid_action_mask = torch.ones(env.action_space.nvec.sum())\n        invalid_action_mask[0:env.action_space.nvec[0]] = torch.tensor(env.unit_location_mask)\n        invalid_action_mask[-env.action_space.nvec[-1]:] = torch.tensor(env.target_unit_location_mask)\n        invalid_action_masks[step] = invalid_action_mask\n        with torch.no_grad():\n            values[step] = vf.forward(obs[step:step+1])\n            action, logproba, _, probs = pg.get_action(obs[step:step+1], invalid_action_masks=invalid_action_masks[step:step+1])\n        \n        actions[step] = action[:,0].data.cpu().numpy()\n        logprobs[:,[step]] = logproba\n\n        # TRY NOT TO MODIFY: execute the game and log data.\n        next_obs, rewards[step], dones[step], info = env.step(action[:,0].data.cpu().numpy())\n        raw_rewards[:,step] = info[\"rewards\"]\n        real_rewards += [info['real_reward']]\n        invalid_action_stats += [info['invalid_action_stats']]\n        next_obs = np.array(next_obs)\n\n        # Annealing the rate if instructed to do so.\n        if args.anneal_lr:\n            pg_lr_scheduler.step()\n            vf_lr_scheduler.step()\n\n        if dones[step]:\n            # Computing the discounted returns:\n            writer.add_scalar(\"charts/episode_reward\", np.sum(real_rewards), global_step)\n            print(f\"global_step={global_step}, episode_reward={np.sum(real_rewards)}\")\n            for i in range(len(env.rfs)):\n                writer.add_scalar(f\"charts/episode_reward/{str(env.rfs[i])}\", raw_rewards.sum(1)[i], global_step)\n            real_rewards = []\n            for key, idx in zip(info['invalid_action_stats'], range(len(info['invalid_action_stats']))):\n                writer.add_scalar(f\"stats/{key}\", pd.DataFrame(invalid_action_stats).sum(0)[idx], global_step)\n            invalid_action_stats = []\n            next_obs = np.array(env.reset())\n\n    # bootstrap reward if not done. reached the batch limit\n    last_value = 0\n    if not dones[step]:\n        last_value = vf.forward(next_obs.reshape((1,)+next_obs.shape))[0].detach().cpu().numpy()[0]\n    bootstrapped_rewards = np.append(rewards, last_value)\n\n    # calculate the returns and advantages\n    if args.gae:\n        bootstrapped_values = np.append(values.detach().cpu().numpy(), last_value)\n        deltas = bootstrapped_rewards[:-1] + args.gamma * bootstrapped_values[1:] * (1-dones) - bootstrapped_values[:-1]\n        advantages = discount_cumsum(deltas, dones, args.gamma * args.gae_lambda)\n        advantages = torch.Tensor(advantages).to(device)\n        returns = advantages + values\n    else:\n        returns = discount_cumsum(bootstrapped_rewards, dones, args.gamma)[:-1]\n        advantages = returns - values.detach().cpu().numpy()\n        advantages = torch.Tensor(advantages).to(device)\n        returns = torch.Tensor(returns).to(device)\n\n    # Advantage normalization\n    if args.norm_adv:\n        EPS = 1e-10\n        advantages = (advantages - advantages.mean()) / (advantages.std() + EPS)\n\n    # Optimizaing policy network\n    entropys = []\n    target_pg = Policy().to(device)\n    inds = np.arange(args.batch_size,)\n    for i_epoch_pi in range(args.update_epochs):\n        np.random.shuffle(inds)\n        for start in range(0, args.batch_size, args.minibatch_size):\n            end = start + args.minibatch_size\n            minibatch_ind = inds[start:end]\n            target_pg.load_state_dict(pg.state_dict())\n            \n            _, newlogproba, _, _ = pg.get_action(\n                obs[minibatch_ind],\n                torch.LongTensor(actions[minibatch_ind].astype(np.int)).to(device).T,\n                invalid_action_masks[minibatch_ind])\n            ratio = (newlogproba - logprobs[:,minibatch_ind]).exp()\n\n            # Policy loss as in OpenAI SpinUp\n            clip_adv = torch.where(advantages[minibatch_ind] > 0,\n                                    (1.+args.clip_coef) * advantages[minibatch_ind],\n                                    (1.-args.clip_coef) * advantages[minibatch_ind]).to(device)\n\n            # Entropy computation with resampled actions\n            entropy = -(newlogproba.exp() * newlogproba).mean()\n            entropys.append(entropy.item())\n\n            policy_loss = -torch.min(ratio * advantages[minibatch_ind], clip_adv) + args.ent_coef * entropy\n            policy_loss = policy_loss.mean()\n            \n            pg_optimizer.zero_grad()\n            policy_loss.backward()\n            nn.utils.clip_grad_norm_(pg.parameters(), args.max_grad_norm)\n            pg_optimizer.step()\n\n            approx_kl = (logprobs[:,minibatch_ind] - newlogproba).mean()\n            # Optimizing value network\n            new_values = vf.forward(obs[minibatch_ind]).view(-1)\n\n            # Value loss clipping\n            if args.clip_vloss:\n                v_loss_unclipped = ((new_values - returns[minibatch_ind]) ** 2)\n                v_clipped = values[minibatch_ind] + torch.clamp(new_values - values[minibatch_ind], -args.clip_coef, args.clip_coef)\n                v_loss_clipped = (v_clipped - returns[minibatch_ind])**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 = torch.mean((returns[minibatch_ind]- new_values).pow(2))\n\n            v_optimizer.zero_grad()\n            v_loss.backward()\n            nn.utils.clip_grad_norm_(vf.parameters(), args.max_grad_norm)\n            v_optimizer.step()\n\n        if args.kle_stop:\n            if approx_kl > args.target_kl:\n                break\n        if args.kle_rollback:\n            if (logprobs[:,minibatch_ind] - \n                pg.get_action(\n                    obs[minibatch_ind],\n                    torch.LongTensor(actions[minibatch_ind].astype(np.int)).to(device).T,\n                    invalid_action_masks[minibatch_ind])[1]).mean() > args.target_kl:\n                pg.load_state_dict(target_pg.state_dict())\n                break\n\n    # TRY NOT TO MODIFY: record rewards for plotting purposes\n    writer.add_scalar(\"losses/value_loss\", v_loss.item(), global_step)\n    writer.add_scalar(\"charts/policy_learning_rate\", pg_optimizer.param_groups[0]['lr'], global_step)\n    writer.add_scalar(\"charts/value_learning_rate\", v_optimizer.param_groups[0]['lr'], global_step)\n    writer.add_scalar(\"losses/policy_loss\", policy_loss.item(), global_step)\n    writer.add_scalar(\"losses/entropy\", np.mean(entropys), global_step)\n    writer.add_scalar(\"losses/approx_kl\", approx_kl.item(), global_step)\n    if args.kle_stop or args.kle_rollback:\n        writer.add_scalar(\"debug/pg_stop_iter\", i_epoch_pi, global_step)\n        \n    # evaluate no mask\n    average_reward, average_invalid_action_stats = evaluate_with_no_mask()\n    writer.add_scalar(\"evals/charts/episode_reward\", average_reward, global_step)\n    print(f\"global_step={global_step}, eval_reward={average_reward}\")\n    for key, idx in zip(info['invalid_action_stats'], range(len(info['invalid_action_stats']))):\n        writer.add_scalar(f\"evals/stats/{key}\", average_invalid_action_stats[idx], global_step)\n\nenv.close()\nwriter.close()\n"
  },
  {
    "path": "invalid_action_masking/ppo_4x4.py",
    "content": "import torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport torch.nn.functional as F\nfrom torch.distributions.categorical import Categorical\nfrom torch.utils.tensorboard import SummaryWriter\n\nfrom cleanrl.common import preprocess_obs_space, preprocess_ac_space\nimport argparse\nimport numpy as np\nimport gym\nimport gym_microrts\nfrom gym.wrappers import TimeLimit, Monitor\nfrom gym.spaces import Discrete, Box, MultiBinary, MultiDiscrete, Space\nimport time\nimport random\nimport os\nimport pandas as pd\n\n# taken from https://github.com/openai/baselines/blob/master/baselines/common/vec_env/vec_normalize.py\nclass RunningMeanStd(object):\n    def __init__(self, epsilon=1e-4, shape=()):\n        self.mean = np.zeros(shape, 'float64')\n        self.var = np.ones(shape, 'float64')\n        self.count = epsilon\n\n    def update(self, x):\n        batch_mean = np.mean([x], axis=0)\n        batch_var = np.var([x], axis=0)\n        batch_count = 1\n        self.update_from_moments(batch_mean, batch_var, batch_count)\n\n    def update_from_moments(self, batch_mean, batch_var, batch_count):\n        self.mean, self.var, self.count = update_mean_var_count_from_moments(\n            self.mean, self.var, self.count, batch_mean, batch_var, batch_count)\n\ndef update_mean_var_count_from_moments(mean, var, count, batch_mean, batch_var, batch_count):\n    delta = batch_mean - mean\n    tot_count = count + batch_count\n\n    new_mean = mean + delta * batch_count / tot_count\n    m_a = var * count\n    m_b = batch_var * batch_count\n    M2 = m_a + m_b + np.square(delta) * count * batch_count / tot_count\n    new_var = M2 / tot_count\n    new_count = tot_count\n\n    return new_mean, new_var, new_count\nclass NormalizedEnv(gym.core.Wrapper):\n    def __init__(self, env, ob=True, ret=True, clipob=10., cliprew=10., gamma=0.99, epsilon=1e-8):\n        super(NormalizedEnv, self).__init__(env)\n        self.ob_rms = RunningMeanStd(shape=self.observation_space.shape) if ob else None\n        self.ret_rms = RunningMeanStd(shape=(1,)) if ret else None\n        self.clipob = clipob\n        self.cliprew = cliprew\n        self.ret = np.zeros(())\n        self.gamma = gamma\n        self.epsilon = epsilon\n\n    def step(self, action):\n        obs, rews, news, infos = self.env.step(action)\n        infos['real_reward'] = rews\n        # print(\"before\", self.ret)\n        self.ret = self.ret * self.gamma + rews\n        # print(\"after\", self.ret)\n        obs = self._obfilt(obs)\n        if self.ret_rms:\n            self.ret_rms.update(np.array([self.ret].copy()))\n            rews = np.clip(rews / np.sqrt(self.ret_rms.var + self.epsilon), -self.cliprew, self.cliprew)\n        self.ret = self.ret * (1-float(news))\n        return obs, rews, news, infos\n\n    def _obfilt(self, obs):\n        if self.ob_rms:\n            self.ob_rms.update(obs)\n            obs = np.clip((obs - self.ob_rms.mean) / np.sqrt(self.ob_rms.var + self.epsilon), -self.clipob, self.clipob)\n            return obs\n        else:\n            return obs\n\n    def reset(self):\n        self.ret = np.zeros(())\n        obs = self.env.reset()\n        return self._obfilt(obs)\n\nif __name__ == \"__main__\":\n    parser = argparse.ArgumentParser(description='PPO agent')\n    # Common arguments\n    parser.add_argument('--exp-name', type=str, default=os.path.basename(__file__).rstrip(\".py\"),\n                       help='the name of this experiment')\n    parser.add_argument('--gym-id', type=str, default=\"MicrortsMining4x4F9-v0\",\n                       help='the id of the gym environment')\n    parser.add_argument('--seed', type=int, default=1,\n                       help='seed of the experiment')\n    parser.add_argument('--episode-length', type=int, default=0,\n                       help='the maximum length of each episode')\n    parser.add_argument('--total-timesteps', type=int, default=100000,\n                       help='total timesteps of the experiments')\n    parser.add_argument('--no-torch-deterministic', action='store_false', dest=\"torch_deterministic\", default=True,\n                       help='if toggled, `torch.backends.cudnn.deterministic=False`')\n    parser.add_argument('--no-cuda', action='store_false', dest=\"cuda\", default=True,\n                       help='if toggled, cuda will not be enabled by default')\n    parser.add_argument('--prod-mode', action='store_true', default=False,\n                       help='run the script in production mode and use wandb to log outputs')\n    parser.add_argument('--capture-video', action='store_true', default=False,\n                       help='weather to capture videos of the agent performances (check out `videos` folder)')\n    parser.add_argument('--wandb-project-name', type=str, default=\"cleanRL\",\n                       help=\"the wandb's project name\")\n    parser.add_argument('--wandb-entity', type=str, default=None,\n                       help=\"the entity (team) of wandb's project\")\n\n    # Algorithm specific arguments\n    parser.add_argument('--batch-size', type=int, default=2048,\n                       help='the batch size of ppo')\n    parser.add_argument('--minibatch-size', type=int, default=256,\n                       help='the mini batch size of ppo')\n    parser.add_argument('--gamma', type=float, default=0.99,\n                       help='the discount factor gamma')\n    parser.add_argument('--gae-lambda', type=float, default=0.97,\n                       help='the lambda for the general advantage estimation')\n    parser.add_argument('--ent-coef', type=float, default=0.01,\n                       help=\"coefficient of the entropy\")\n    parser.add_argument('--max-grad-norm', type=float, default=0.5,\n                       help='the maximum norm for the gradient clipping')\n    parser.add_argument('--clip-coef', type=float, default=0.2,\n                       help=\"the surrogate clipping coefficient\")\n    parser.add_argument('--update-epochs', type=int, default=10,\n                        help=\"the K epochs to update the policy\")\n    parser.add_argument('--kle-stop', action='store_true', default=False,\n                        help='If toggled, the policy updates will be early stopped w.r.t target-kl')\n    parser.add_argument('--kle-rollback', action='store_true', default=False,\n                        help='If toggled, the policy updates will roll back to previous policy if KL exceeds target-kl')\n    parser.add_argument('--target-kl', type=float, default=0.015,\n                        help='the target-kl variable that is referred by --kl')\n    parser.add_argument('--gae', action='store_true', default=True,\n                        help='Use GAE for advantage computation')\n    parser.add_argument('--policy-lr', type=float, default=3e-4,\n                        help=\"the learning rate of the policy optimizer\")\n    parser.add_argument('--value-lr', type=float, default=3e-4,\n                        help=\"the learning rate of the critic optimizer\")\n    parser.add_argument('--norm-obs', action='store_true', default=True,\n                        help=\"Toggles observation normalization\")\n    parser.add_argument('--norm-returns', action='store_true', default=False,\n                        help=\"Toggles returns normalization\")\n    parser.add_argument('--norm-adv', action='store_true', default=True,\n                        help=\"Toggles advantages normalization\")\n    parser.add_argument('--obs-clip', type=float, default=10.0,\n                        help=\"Value for reward clipping, as per the paper\")\n    parser.add_argument('--rew-clip', type=float, default=10.0,\n                        help=\"Value for observation clipping, as per the paper\")\n    parser.add_argument('--anneal-lr', action='store_true', default=True,\n                        help=\"Toggle learning rate annealing for policy and value networks\")\n    parser.add_argument('--weights-init', default=\"orthogonal\", choices=[\"xavier\", 'orthogonal'],\n                        help='Selects the scheme to be used for weights initialization'),\n    parser.add_argument('--clip-vloss', action=\"store_true\", default=True,\n                        help='Toggles wheter or not to use a clipped loss for the value function, as per the paper.')\n    parser.add_argument('--pol-layer-norm', action='store_true', default=False,\n                       help='Enables layer normalization in the policy network')\n    args = parser.parse_args()\n    if not args.seed:\n        args.seed = int(time.time())\n\nargs.features_turned_on = sum([args.kle_stop, args.kle_rollback, args.gae, args.norm_obs, args.norm_returns, args.norm_adv, args.anneal_lr, args.clip_vloss, args.pol_layer_norm])\n# TRY NOT TO MODIFY: setup the environment\nexperiment_name = f\"{args.gym_id}__{args.exp_name}__{args.seed}__{int(time.time())}\"\nwriter = SummaryWriter(f\"runs/{experiment_name}\")\nwriter.add_text('hyperparameters', \"|param|value|\\n|-|-|\\n%s\" % (\n        '\\n'.join([f\"|{key}|{value}|\" for key, value in vars(args).items()])))\n\nif args.prod_mode:\n    import wandb\n    wandb.init(project=args.wandb_project_name, entity=args.wandb_entity, tensorboard=True, config=vars(args), name=experiment_name, monitor_gym=True)\n    writer = SummaryWriter(f\"/tmp/{experiment_name}\")\n    wandb.save(os.path.abspath(__file__))\n\n# TRY NOT TO MODIFY: seeding\ndevice = torch.device('cuda' if torch.cuda.is_available() and args.cuda else 'cpu')\nenv = gym.make(args.gym_id)\n# respect the default timelimit\nassert isinstance(env.action_space, MultiDiscrete), \"only MultiDiscrete action space is supported\"\nassert isinstance(env, TimeLimit) or int(args.episode_length), \"the gym env does not have a built in TimeLimit, please specify by using --episode-length\"\nif isinstance(env, TimeLimit):\n    if int(args.episode_length):\n        env._max_episode_steps = int(args.episode_length)\n    args.episode_length = env._max_episode_steps\nelse:\n    env = TimeLimit(env, int(args.episode_length))\nenv = NormalizedEnv(env.env, ob=args.norm_obs, ret=args.norm_returns, clipob=args.obs_clip, cliprew=args.rew_clip, gamma=args.gamma)\nenv = TimeLimit(env, int(args.episode_length))\nrandom.seed(args.seed)\nnp.random.seed(args.seed)\ntorch.manual_seed(args.seed)\ntorch.backends.cudnn.deterministic = args.torch_deterministic\nenv.seed(args.seed)\nenv.action_space.seed(args.seed)\nenv.observation_space.seed(args.seed)\nif args.capture_video:\n    env = Monitor(env, f'videos/{experiment_name}')\n\n# ALGO LOGIC: initialize agent here:\nclass CategoricalMasked(Categorical):\n    def __init__(self, probs=None, logits=None, validate_args=None, masks=[]):\n        self.masks = masks\n        if len(self.masks) == 0:\n            super(CategoricalMasked, self).__init__(probs, logits, validate_args)\n        else:\n            self.masks = masks.type(torch.BoolTensor).to(device)\n            logits = torch.where(self.masks, logits, torch.tensor(-1e+8).to(device))\n            super(CategoricalMasked, self).__init__(probs, logits, validate_args)\n    \n    def entropy(self):\n        if len(self.masks) == 0:\n            return super(CategoricalMasked, self).entropy()\n        p_log_p = self.logits * self.probs\n        p_log_p = torch.where(self.masks, p_log_p, torch.tensor(0.).to(device))\n        return -p_log_p.sum(-1)\n\nclass Policy(nn.Module):\n    def __init__(self):\n        super(Policy, self).__init__()\n        self.features = nn.Sequential(\n            nn.Conv2d(27, 16, kernel_size=2,),\n            nn.MaxPool2d(1),\n            nn.ReLU())\n        self.fc = nn.Sequential(\n            nn.Linear(16*3*3, 128),\n            nn.ReLU(),\n            nn.Linear(128, env.action_space.nvec.sum())\n        )\n    def forward(self, x):\n        x = torch.Tensor(np.moveaxis(x, -1, 1)).to(device)\n        x = self.features(x)\n        x = x.reshape(x.size(0), -1)\n        x = self.fc(x)\n        return x\n\n    def get_action(self, x, action=None, invalid_action_masks=None):\n        logits = self.forward(x)\n        split_logits = torch.split(logits, env.action_space.nvec.tolist(), dim=1)\n        \n        if invalid_action_masks is not None:\n            split_invalid_action_masks = torch.split(invalid_action_masks, env.action_space.nvec.tolist(), dim=1)\n            multi_categoricals = [CategoricalMasked(logits=logits, masks=iam) for (logits, iam) in zip(split_logits, split_invalid_action_masks)]\n        else:\n            multi_categoricals = [Categorical(logits=logits) for logits in split_logits]\n        \n        if action is None:\n            action = torch.stack([categorical.sample() for categorical in multi_categoricals])\n        logprob = torch.stack([categorical.log_prob(a) for a, categorical in zip(action, multi_categoricals)])\n        return action, logprob, [], multi_categoricals\n\nclass Value(nn.Module):\n    def __init__(self):\n        super(Value, self).__init__()\n        self.features = nn.Sequential(\n            nn.Conv2d(27, 16, kernel_size=2,),\n            nn.MaxPool2d(1),\n            nn.ReLU())\n        self.fc = nn.Sequential(\n            nn.Linear(16*3*3, 128),\n            nn.ReLU(),\n            nn.Linear(128, 1)\n        )\n\n    def forward(self, x):\n        x = torch.Tensor(np.moveaxis(x, -1, 1)).to(device)\n        x = self.features(x)\n        x = x.reshape(x.size(0), -1)\n        x = self.fc(x)\n        return x\n\ndef discount_cumsum(x, dones, gamma):\n    \"\"\"\n    computing discounted cumulative sums of vectors that resets with dones\n    input:\n        vector x,  vector dones,\n        [x0,       [0,\n         x1,        0,\n         x2         1,\n         x3         0, \n         x4]        0]\n    output:\n        [x0 + discount * x1 + discount^2 * x2,\n         x1 + discount * x2,\n         x2,\n         x3 + discount * x4,\n         x4]\n    \"\"\"\n    discount_cumsum = np.zeros_like(x)\n    discount_cumsum[-1] = x[-1]\n    for t in reversed(range(x.shape[0]-1)):\n        discount_cumsum[t] = x[t] + gamma * discount_cumsum[t+1] * (1-dones[t])\n    return discount_cumsum\n\npg = Policy().to(device)\nvf = Value().to(device)\n\n# MODIFIED: Separate optimizer and learning rates\npg_optimizer = optim.Adam(list(pg.parameters()), lr=args.policy_lr)\nv_optimizer = optim.Adam(list(vf.parameters()), lr=args.value_lr)\n\n# MODIFIED: Initializing learning rate anneal scheduler when need\nif args.anneal_lr:\n    anneal_fn = lambda f: max(0, 1-f / args.total_timesteps)\n    pg_lr_scheduler = optim.lr_scheduler.LambdaLR(pg_optimizer, lr_lambda=anneal_fn)\n    vf_lr_scheduler = optim.lr_scheduler.LambdaLR(v_optimizer, lr_lambda=anneal_fn)\n\nloss_fn = nn.MSELoss()\n\ndef evaluate_with_no_mask():\n    evaluate_rewards = []\n    evaluate_invalid_action_stats = []\n    if args.capture_video:\n        env.stats_recorder.done=True\n    next_obs = np.array(env.reset())\n\n    # ALGO Logic: Storage for epoch data\n    obs = np.empty((args.batch_size,) + env.observation_space.shape)\n\n    actions = np.empty((args.batch_size,) + env.action_space.shape)\n    logprobs = torch.zeros((env.action_space.nvec.shape[0], args.batch_size,)).to(device)\n\n    rewards = np.zeros((args.batch_size,))\n    raw_rewards = np.zeros((len(env.rfs),args.batch_size,))\n    \n    real_rewards = []\n    invalid_action_stats = []\n\n    dones = np.zeros((args.batch_size,))\n    values = torch.zeros((args.batch_size,)).to(device)\n\n    # TRY NOT TO MODIFY: prepare the execution of the game.\n    for step in range(args.batch_size):\n        env.render()\n        obs[step] = next_obs.copy()\n\n        # ALGO LOGIC: put action logic here\n        with torch.no_grad():\n            values[step] = vf.forward(obs[step:step+1])\n            action, logproba, _, probs = pg.get_action(obs[step:step+1])\n        \n        actions[step] = action[:,0].data.cpu().numpy()\n        logprobs[:,[step]] = logproba\n\n        # TRY NOT TO MODIFY: execute the game and log data.\n        next_obs, rewards[step], dones[step], info = env.step(action[:,0].data.cpu().numpy())\n        raw_rewards[:,step] = info[\"rewards\"]\n        real_rewards += [info['real_reward']]\n        invalid_action_stats += [info['invalid_action_stats']]\n        next_obs = np.array(next_obs)\n\n        if dones[step]:\n            # Computing the discounted returns:\n            writer.add_scalar(\"charts/episode_reward\", np.sum(real_rewards), global_step)\n            evaluate_rewards += [np.sum(real_rewards)]\n            real_rewards = []\n            evaluate_invalid_action_stats += [pd.DataFrame(invalid_action_stats).sum(0)]\n            invalid_action_stats = []\n            next_obs = np.array(env.reset())\n    \n    return np.average(evaluate_rewards), pd.DataFrame(evaluate_invalid_action_stats).mean(0)\n\n# TRY NOT TO MODIFY: start the game\nglobal_step = 0\nwhile global_step < args.total_timesteps:\n    if args.capture_video:\n        env.stats_recorder.done=True\n    next_obs = np.array(env.reset())\n\n    # ALGO Logic: Storage for epoch data\n    obs = np.empty((args.batch_size,) + env.observation_space.shape)\n\n    actions = np.empty((args.batch_size,) + env.action_space.shape)\n    logprobs = torch.zeros((env.action_space.nvec.shape[0], args.batch_size,)).to(device)\n\n    rewards = np.zeros((args.batch_size,))\n    raw_rewards = np.zeros((len(env.rfs),args.batch_size,))\n    \n    real_rewards = []\n    invalid_action_stats = []\n\n    dones = np.zeros((args.batch_size,))\n    values = torch.zeros((args.batch_size,)).to(device)\n\n    invalid_action_masks = torch.zeros((args.batch_size, env.action_space.nvec.sum()))\n    # TRY NOT TO MODIFY: prepare the execution of the game.\n    for step in range(args.batch_size):\n        env.render()\n        global_step += 1\n        obs[step] = next_obs.copy()\n\n        # ALGO LOGIC: put action logic here\n        invalid_action_mask = torch.ones(env.action_space.nvec.sum())\n        invalid_action_mask[0:env.action_space.nvec[0]] = torch.tensor(env.unit_location_mask)\n        invalid_action_mask[-env.action_space.nvec[-1]:] = torch.tensor(env.target_unit_location_mask)\n        invalid_action_masks[step] = invalid_action_mask\n        with torch.no_grad():\n            values[step] = vf.forward(obs[step:step+1])\n            action, logproba, _, probs = pg.get_action(obs[step:step+1], invalid_action_masks=invalid_action_masks[step:step+1])\n        \n        actions[step] = action[:,0].data.cpu().numpy()\n        logprobs[:,[step]] = logproba\n\n        # TRY NOT TO MODIFY: execute the game and log data.\n        next_obs, rewards[step], dones[step], info = env.step(action[:,0].data.cpu().numpy())\n        raw_rewards[:,step] = info[\"rewards\"]\n        real_rewards += [info['real_reward']]\n        invalid_action_stats += [info['invalid_action_stats']]\n        next_obs = np.array(next_obs)\n\n        # Annealing the rate if instructed to do so.\n        if args.anneal_lr:\n            pg_lr_scheduler.step()\n            vf_lr_scheduler.step()\n\n        if dones[step]:\n            # Computing the discounted returns:\n            writer.add_scalar(\"charts/episode_reward\", np.sum(real_rewards), global_step)\n            print(f\"global_step={global_step}, episode_reward={np.sum(real_rewards)}\")\n            for i in range(len(env.rfs)):\n                writer.add_scalar(f\"charts/episode_reward/{str(env.rfs[i])}\", raw_rewards.sum(1)[i], global_step)\n            real_rewards = []\n            for key, idx in zip(info['invalid_action_stats'], range(len(info['invalid_action_stats']))):\n                writer.add_scalar(f\"stats/{key}\", pd.DataFrame(invalid_action_stats).sum(0)[idx], global_step)\n            invalid_action_stats = []\n            next_obs = np.array(env.reset())\n\n    # bootstrap reward if not done. reached the batch limit\n    last_value = 0\n    if not dones[step]:\n        last_value = vf.forward(next_obs.reshape((1,)+next_obs.shape))[0].detach().cpu().numpy()[0]\n    bootstrapped_rewards = np.append(rewards, last_value)\n\n    # calculate the returns and advantages\n    if args.gae:\n        bootstrapped_values = np.append(values.detach().cpu().numpy(), last_value)\n        deltas = bootstrapped_rewards[:-1] + args.gamma * bootstrapped_values[1:] * (1-dones) - bootstrapped_values[:-1]\n        advantages = discount_cumsum(deltas, dones, args.gamma * args.gae_lambda)\n        advantages = torch.Tensor(advantages).to(device)\n        returns = advantages + values\n    else:\n        returns = discount_cumsum(bootstrapped_rewards, dones, args.gamma)[:-1]\n        advantages = returns - values.detach().cpu().numpy()\n        advantages = torch.Tensor(advantages).to(device)\n        returns = torch.Tensor(returns).to(device)\n\n    # Advantage normalization\n    if args.norm_adv:\n        EPS = 1e-10\n        advantages = (advantages - advantages.mean()) / (advantages.std() + EPS)\n\n    # Optimizaing policy network\n    entropys = []\n    target_pg = Policy().to(device)\n    inds = np.arange(args.batch_size,)\n    for i_epoch_pi in range(args.update_epochs):\n        np.random.shuffle(inds)\n        for start in range(0, args.batch_size, args.minibatch_size):\n            end = start + args.minibatch_size\n            minibatch_ind = inds[start:end]\n            target_pg.load_state_dict(pg.state_dict())\n            \n            _, newlogproba, _, _ = pg.get_action(\n                obs[minibatch_ind],\n                torch.LongTensor(actions[minibatch_ind].astype(np.int)).to(device).T,\n                invalid_action_masks[minibatch_ind])\n            ratio = (newlogproba - logprobs[:,minibatch_ind]).exp()\n\n            # Policy loss as in OpenAI SpinUp\n            clip_adv = torch.where(advantages[minibatch_ind] > 0,\n                                    (1.+args.clip_coef) * advantages[minibatch_ind],\n                                    (1.-args.clip_coef) * advantages[minibatch_ind]).to(device)\n\n            # Entropy computation with resampled actions\n            entropy = -(newlogproba.exp() * newlogproba).mean()\n            entropys.append(entropy.item())\n\n            policy_loss = -torch.min(ratio * advantages[minibatch_ind], clip_adv) + args.ent_coef * entropy\n            policy_loss = policy_loss.mean()\n            \n            pg_optimizer.zero_grad()\n            policy_loss.backward()\n            nn.utils.clip_grad_norm_(pg.parameters(), args.max_grad_norm)\n            pg_optimizer.step()\n\n            approx_kl = (logprobs[:,minibatch_ind] - newlogproba).mean()\n            # Optimizing value network\n            new_values = vf.forward(obs[minibatch_ind]).view(-1)\n\n            # Value loss clipping\n            if args.clip_vloss:\n                v_loss_unclipped = ((new_values - returns[minibatch_ind]) ** 2)\n                v_clipped = values[minibatch_ind] + torch.clamp(new_values - values[minibatch_ind], -args.clip_coef, args.clip_coef)\n                v_loss_clipped = (v_clipped - returns[minibatch_ind])**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 = torch.mean((returns[minibatch_ind]- new_values).pow(2))\n\n            v_optimizer.zero_grad()\n            v_loss.backward()\n            nn.utils.clip_grad_norm_(vf.parameters(), args.max_grad_norm)\n            v_optimizer.step()\n\n        if args.kle_stop:\n            if approx_kl > args.target_kl:\n                break\n        if args.kle_rollback:\n            if (logprobs[:,minibatch_ind] - \n                pg.get_action(\n                    obs[minibatch_ind],\n                    torch.LongTensor(actions[minibatch_ind].astype(np.int)).to(device).T,\n                    invalid_action_masks[minibatch_ind])[1]).mean() > args.target_kl:\n                pg.load_state_dict(target_pg.state_dict())\n                break\n\n    # TRY NOT TO MODIFY: record rewards for plotting purposes\n    writer.add_scalar(\"losses/value_loss\", v_loss.item(), global_step)\n    writer.add_scalar(\"charts/policy_learning_rate\", pg_optimizer.param_groups[0]['lr'], global_step)\n    writer.add_scalar(\"charts/value_learning_rate\", v_optimizer.param_groups[0]['lr'], global_step)\n    writer.add_scalar(\"losses/policy_loss\", policy_loss.item(), global_step)\n    writer.add_scalar(\"losses/entropy\", np.mean(entropys), global_step)\n    writer.add_scalar(\"losses/approx_kl\", approx_kl.item(), global_step)\n    if args.kle_stop or args.kle_rollback:\n        writer.add_scalar(\"debug/pg_stop_iter\", i_epoch_pi, global_step)\n        \n    # evaluate no mask\n    average_reward, average_invalid_action_stats = evaluate_with_no_mask()\n    writer.add_scalar(\"evals/charts/episode_reward\", average_reward, global_step)\n    print(f\"global_step={global_step}, eval_reward={average_reward}\")\n    for key, idx in zip(info['invalid_action_stats'], range(len(info['invalid_action_stats']))):\n        writer.add_scalar(f\"evals/stats/{key}\", average_invalid_action_stats[idx], global_step)\n\nenv.close()\nwriter.close()\n"
  },
  {
    "path": "invalid_action_masking/ppo_no_adj_10x10.py",
    "content": "import torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport torch.nn.functional as F\nfrom torch.distributions.categorical import Categorical\nfrom torch.utils.tensorboard import SummaryWriter\n\nfrom cleanrl.common import preprocess_obs_space, preprocess_ac_space\nimport argparse\nimport numpy as np\nimport gym\nimport gym_microrts\nfrom gym.wrappers import TimeLimit, Monitor\nfrom gym.spaces import Discrete, Box, MultiBinary, MultiDiscrete, Space\nimport time\nimport random\nimport os\nimport pandas as pd\n\n# taken from https://github.com/openai/baselines/blob/master/baselines/common/vec_env/vec_normalize.py\nclass RunningMeanStd(object):\n    def __init__(self, epsilon=1e-4, shape=()):\n        self.mean = np.zeros(shape, 'float64')\n        self.var = np.ones(shape, 'float64')\n        self.count = epsilon\n\n    def update(self, x):\n        batch_mean = np.mean([x], axis=0)\n        batch_var = np.var([x], axis=0)\n        batch_count = 1\n        self.update_from_moments(batch_mean, batch_var, batch_count)\n\n    def update_from_moments(self, batch_mean, batch_var, batch_count):\n        self.mean, self.var, self.count = update_mean_var_count_from_moments(\n            self.mean, self.var, self.count, batch_mean, batch_var, batch_count)\n\ndef update_mean_var_count_from_moments(mean, var, count, batch_mean, batch_var, batch_count):\n    delta = batch_mean - mean\n    tot_count = count + batch_count\n\n    new_mean = mean + delta * batch_count / tot_count\n    m_a = var * count\n    m_b = batch_var * batch_count\n    M2 = m_a + m_b + np.square(delta) * count * batch_count / tot_count\n    new_var = M2 / tot_count\n    new_count = tot_count\n\n    return new_mean, new_var, new_count\nclass NormalizedEnv(gym.core.Wrapper):\n    def __init__(self, env, ob=True, ret=True, clipob=10., cliprew=10., gamma=0.99, epsilon=1e-8):\n        super(NormalizedEnv, self).__init__(env)\n        self.ob_rms = RunningMeanStd(shape=self.observation_space.shape) if ob else None\n        self.ret_rms = RunningMeanStd(shape=(1,)) if ret else None\n        self.clipob = clipob\n        self.cliprew = cliprew\n        self.ret = np.zeros(())\n        self.gamma = gamma\n        self.epsilon = epsilon\n\n    def step(self, action):\n        obs, rews, news, infos = self.env.step(action)\n        infos['real_reward'] = rews\n        # print(\"before\", self.ret)\n        self.ret = self.ret * self.gamma + rews\n        # print(\"after\", self.ret)\n        obs = self._obfilt(obs)\n        if self.ret_rms:\n            self.ret_rms.update(np.array([self.ret].copy()))\n            rews = np.clip(rews / np.sqrt(self.ret_rms.var + self.epsilon), -self.cliprew, self.cliprew)\n        self.ret = self.ret * (1-float(news))\n        return obs, rews, news, infos\n\n    def _obfilt(self, obs):\n        if self.ob_rms:\n            self.ob_rms.update(obs)\n            obs = np.clip((obs - self.ob_rms.mean) / np.sqrt(self.ob_rms.var + self.epsilon), -self.clipob, self.clipob)\n            return obs\n        else:\n            return obs\n\n    def reset(self):\n        self.ret = np.zeros(())\n        obs = self.env.reset()\n        return self._obfilt(obs)\n\nif __name__ == \"__main__\":\n    parser = argparse.ArgumentParser(description='PPO agent')\n    # Common arguments\n    parser.add_argument('--exp-name', type=str, default=os.path.basename(__file__).rstrip(\".py\"),\n                       help='the name of this experiment')\n    parser.add_argument('--gym-id', type=str, default=\"MicrortsMining10x10F9-v0\",\n                       help='the id of the gym environment')\n    parser.add_argument('--seed', type=int, default=1,\n                       help='seed of the experiment')\n    parser.add_argument('--episode-length', type=int, default=0,\n                       help='the maximum length of each episode')\n    parser.add_argument('--total-timesteps', type=int, default=100000,\n                       help='total timesteps of the experiments')\n    parser.add_argument('--no-torch-deterministic', action='store_false', dest=\"torch_deterministic\", default=True,\n                       help='if toggled, `torch.backends.cudnn.deterministic=False`')\n    parser.add_argument('--no-cuda', action='store_false', dest=\"cuda\", default=True,\n                       help='if toggled, cuda will not be enabled by default')\n    parser.add_argument('--prod-mode', action='store_true', default=False,\n                       help='run the script in production mode and use wandb to log outputs')\n    parser.add_argument('--capture-video', action='store_true', default=False,\n                       help='weather to capture videos of the agent performances (check out `videos` folder)')\n    parser.add_argument('--wandb-project-name', type=str, default=\"cleanRL\",\n                       help=\"the wandb's project name\")\n    parser.add_argument('--wandb-entity', type=str, default=None,\n                       help=\"the entity (team) of wandb's project\")\n\n    # Algorithm specific arguments\n    parser.add_argument('--batch-size', type=int, default=2048,\n                       help='the batch size of ppo')\n    parser.add_argument('--minibatch-size', type=int, default=256,\n                       help='the mini batch size of ppo')\n    parser.add_argument('--gamma', type=float, default=0.99,\n                       help='the discount factor gamma')\n    parser.add_argument('--gae-lambda', type=float, default=0.97,\n                       help='the lambda for the general advantage estimation')\n    parser.add_argument('--ent-coef', type=float, default=0.01,\n                       help=\"coefficient of the entropy\")\n    parser.add_argument('--max-grad-norm', type=float, default=0.5,\n                       help='the maximum norm for the gradient clipping')\n    parser.add_argument('--clip-coef', type=float, default=0.2,\n                       help=\"the surrogate clipping coefficient\")\n    parser.add_argument('--update-epochs', type=int, default=10,\n                        help=\"the K epochs to update the policy\")\n    parser.add_argument('--kle-stop', action='store_true', default=False,\n                        help='If toggled, the policy updates will be early stopped w.r.t target-kl')\n    parser.add_argument('--kle-rollback', action='store_true', default=False,\n                        help='If toggled, the policy updates will roll back to previous policy if KL exceeds target-kl')\n    parser.add_argument('--target-kl', type=float, default=0.015,\n                        help='the target-kl variable that is referred by --kl')\n    parser.add_argument('--gae', action='store_true', default=True,\n                        help='Use GAE for advantage computation')\n    parser.add_argument('--policy-lr', type=float, default=3e-4,\n                        help=\"the learning rate of the policy optimizer\")\n    parser.add_argument('--value-lr', type=float, default=3e-4,\n                        help=\"the learning rate of the critic optimizer\")\n    parser.add_argument('--norm-obs', action='store_true', default=True,\n                        help=\"Toggles observation normalization\")\n    parser.add_argument('--norm-returns', action='store_true', default=False,\n                        help=\"Toggles returns normalization\")\n    parser.add_argument('--norm-adv', action='store_true', default=True,\n                        help=\"Toggles advantages normalization\")\n    parser.add_argument('--obs-clip', type=float, default=10.0,\n                        help=\"Value for reward clipping, as per the paper\")\n    parser.add_argument('--rew-clip', type=float, default=10.0,\n                        help=\"Value for observation clipping, as per the paper\")\n    parser.add_argument('--anneal-lr', action='store_true', default=True,\n                        help=\"Toggle learning rate annealing for policy and value networks\")\n    parser.add_argument('--weights-init', default=\"orthogonal\", choices=[\"xavier\", 'orthogonal'],\n                        help='Selects the scheme to be used for weights initialization'),\n    parser.add_argument('--clip-vloss', action=\"store_true\", default=True,\n                        help='Toggles wheter or not to use a clipped loss for the value function, as per the paper.')\n    parser.add_argument('--pol-layer-norm', action='store_true', default=False,\n                       help='Enables layer normalization in the policy network')\n    args = parser.parse_args()\n    if not args.seed:\n        args.seed = int(time.time())\n\nargs.features_turned_on = sum([args.kle_stop, args.kle_rollback, args.gae, args.norm_obs, args.norm_returns, args.norm_adv, args.anneal_lr, args.clip_vloss, args.pol_layer_norm])\n# TRY NOT TO MODIFY: setup the environment\nexperiment_name = f\"{args.gym_id}__{args.exp_name}__{args.seed}__{int(time.time())}\"\nwriter = SummaryWriter(f\"runs/{experiment_name}\")\nwriter.add_text('hyperparameters', \"|param|value|\\n|-|-|\\n%s\" % (\n        '\\n'.join([f\"|{key}|{value}|\" for key, value in vars(args).items()])))\n\nif args.prod_mode:\n    import wandb\n    wandb.init(project=args.wandb_project_name, entity=args.wandb_entity, tensorboard=True, config=vars(args), name=experiment_name, monitor_gym=True)\n    writer = SummaryWriter(f\"/tmp/{experiment_name}\")\n    wandb.save(os.path.abspath(__file__))\n\n# TRY NOT TO MODIFY: seeding\ndevice = torch.device('cuda' if torch.cuda.is_available() and args.cuda else 'cpu')\nenv = gym.make(args.gym_id)\n# respect the default timelimit\nassert isinstance(env.action_space, MultiDiscrete), \"only MultiDiscrete action space is supported\"\nassert isinstance(env, TimeLimit) or int(args.episode_length), \"the gym env does not have a built in TimeLimit, please specify by using --episode-length\"\nif isinstance(env, TimeLimit):\n    if int(args.episode_length):\n        env._max_episode_steps = int(args.episode_length)\n    args.episode_length = env._max_episode_steps\nelse:\n    env = TimeLimit(env, int(args.episode_length))\nenv = NormalizedEnv(env.env, ob=args.norm_obs, ret=args.norm_returns, clipob=args.obs_clip, cliprew=args.rew_clip, gamma=args.gamma)\nenv = TimeLimit(env, int(args.episode_length))\nrandom.seed(args.seed)\nnp.random.seed(args.seed)\ntorch.manual_seed(args.seed)\ntorch.backends.cudnn.deterministic = args.torch_deterministic\nenv.seed(args.seed)\nenv.action_space.seed(args.seed)\nenv.observation_space.seed(args.seed)\nif args.capture_video:\n    env = Monitor(env, f'videos/{experiment_name}')\n\n# ALGO LOGIC: initialize agent here:\nclass CategoricalMasked(Categorical):\n    def __init__(self, probs=None, logits=None, validate_args=None, masks=[]):\n        self.masks = masks\n        if len(self.masks) == 0:\n            super(CategoricalMasked, self).__init__(probs, logits, validate_args)\n        else:\n            self.masks = masks.type(torch.BoolTensor).to(device)\n            logits = torch.where(self.masks, logits, torch.tensor(-1e+8).to(device))\n            super(CategoricalMasked, self).__init__(probs, logits, validate_args)\n    \n    def entropy(self):\n        if len(self.masks) == 0:\n            return super(CategoricalMasked, self).entropy()\n        p_log_p = self.logits * self.probs\n        p_log_p = torch.where(self.masks, p_log_p, torch.tensor(0.).to(device))\n        return -p_log_p.sum(-1)\n\nclass Policy(nn.Module):\n    def __init__(self):\n        super(Policy, self).__init__()\n        self.features = nn.Sequential(\n            nn.Conv2d(27, 16, kernel_size=3,),\n            nn.MaxPool2d(1),\n            nn.ReLU(),\n            nn.Conv2d(16, 32, kernel_size=3),\n            nn.MaxPool2d(1),\n            nn.ReLU())\n        self.fc = nn.Sequential(\n            nn.Linear(32*6*6, 128),\n            nn.ReLU(),\n            nn.Linear(128, env.action_space.nvec.sum())\n        )\n\n    def forward(self, x):\n        x = torch.Tensor(np.moveaxis(x, -1, 1)).to(device)\n        x = self.features(x)\n        x = x.reshape(x.size(0), -1)\n        x = self.fc(x)\n        return x\n\n    def get_action(self, x, action=None, invalid_action_masks=None):\n        logits = self.forward(x)\n        split_logits = torch.split(logits, env.action_space.nvec.tolist(), dim=1)\n        \n        if invalid_action_masks is not None:\n            split_invalid_action_masks = torch.split(invalid_action_masks, env.action_space.nvec.tolist(), dim=1)\n            multi_categoricals = [CategoricalMasked(logits=logits, masks=iam) for (logits, iam) in zip(split_logits, split_invalid_action_masks)]\n        else:\n            multi_categoricals = [Categorical(logits=logits) for logits in split_logits]\n        \n        if action is None:\n            action = torch.stack([categorical.sample() for categorical in multi_categoricals])\n        logprob = torch.stack([categorical.log_prob(a) for a, categorical in zip(action, multi_categoricals)])\n        return action, logprob, [], multi_categoricals\n\nclass Value(nn.Module):\n    def __init__(self):\n        super(Value, self).__init__()\n        self.features = nn.Sequential(\n            nn.Conv2d(27, 16, kernel_size=3,),\n            nn.MaxPool2d(1),\n            nn.ReLU(),\n            nn.Conv2d(16, 32, kernel_size=3),\n            nn.MaxPool2d(1),\n            nn.ReLU())\n        self.fc = nn.Sequential(\n            nn.Linear(32*6*6, 128),\n            nn.ReLU(),\n            nn.Linear(128, 1)\n        )\n\n    def forward(self, x):\n        x = torch.Tensor(np.moveaxis(x, -1, 1)).to(device)\n        x = self.features(x)\n        x = x.reshape(x.size(0), -1)\n        x = self.fc(x)\n        return x\n\ndef discount_cumsum(x, dones, gamma):\n    \"\"\"\n    computing discounted cumulative sums of vectors that resets with dones\n    input:\n        vector x,  vector dones,\n        [x0,       [0,\n         x1,        0,\n         x2         1,\n         x3         0, \n         x4]        0]\n    output:\n        [x0 + discount * x1 + discount^2 * x2,\n         x1 + discount * x2,\n         x2,\n         x3 + discount * x4,\n         x4]\n    \"\"\"\n    discount_cumsum = np.zeros_like(x)\n    discount_cumsum[-1] = x[-1]\n    for t in reversed(range(x.shape[0]-1)):\n        discount_cumsum[t] = x[t] + gamma * discount_cumsum[t+1] * (1-dones[t])\n    return discount_cumsum\n\npg = Policy().to(device)\nvf = Value().to(device)\n\n# MODIFIED: Separate optimizer and learning rates\npg_optimizer = optim.Adam(list(pg.parameters()), lr=args.policy_lr)\nv_optimizer = optim.Adam(list(vf.parameters()), lr=args.value_lr)\n\n# MODIFIED: Initializing learning rate anneal scheduler when need\nif args.anneal_lr:\n    anneal_fn = lambda f: max(0, 1-f / args.total_timesteps)\n    pg_lr_scheduler = optim.lr_scheduler.LambdaLR(pg_optimizer, lr_lambda=anneal_fn)\n    vf_lr_scheduler = optim.lr_scheduler.LambdaLR(v_optimizer, lr_lambda=anneal_fn)\n\nloss_fn = nn.MSELoss()\n\n# TRY NOT TO MODIFY: start the game\nglobal_step = 0\nwhile global_step < args.total_timesteps:\n    if args.capture_video:\n        env.stats_recorder.done=True\n    next_obs = np.array(env.reset())\n\n    # ALGO Logic: Storage for epoch data\n    obs = np.empty((args.batch_size,) + env.observation_space.shape)\n\n    actions = np.empty((args.batch_size,) + env.action_space.shape)\n    logprobs = torch.zeros((env.action_space.nvec.shape[0], args.batch_size,)).to(device)\n\n    rewards = np.zeros((args.batch_size,))\n    raw_rewards = np.zeros((len(env.rfs),args.batch_size,))\n    \n    real_rewards = []\n    invalid_action_stats = []\n\n    dones = np.zeros((args.batch_size,))\n    values = torch.zeros((args.batch_size,)).to(device)\n\n    invalid_action_masks = torch.zeros((args.batch_size, env.action_space.nvec.sum()))\n    # TRY NOT TO MODIFY: prepare the execution of the game.\n    for step in range(args.batch_size):\n        env.render()\n        global_step += 1\n        obs[step] = next_obs.copy()\n\n        # ALGO LOGIC: put action logic here\n        invalid_action_mask = torch.ones(env.action_space.nvec.sum())\n        invalid_action_mask[0:env.action_space.nvec[0]] = torch.tensor(env.unit_location_mask)\n        invalid_action_mask[-env.action_space.nvec[-1]:] = torch.tensor(env.target_unit_location_mask)\n        invalid_action_masks[step] = invalid_action_mask\n        with torch.no_grad():\n            values[step] = vf.forward(obs[step:step+1])\n            action, logproba, _, probs = pg.get_action(obs[step:step+1], invalid_action_masks=invalid_action_masks[step:step+1])\n            \n            # CORE LOGIC:\n            # use the action generated by CategoricalMasked, but \n            # don't adjust the logprobability accordingly. Instead, calculate the log\n            # probability using Categorical\n            action, logproba, _, probs = pg.get_action(obs[step:step+1], action=action)\n        \n        actions[step] = action[:,0].data.cpu().numpy()\n        logprobs[:,[step]] = logproba\n\n        # TRY NOT TO MODIFY: execute the game and log data.\n        next_obs, rewards[step], dones[step], info = env.step(action[:,0].data.cpu().numpy())\n        raw_rewards[:,step] = info[\"rewards\"]\n        real_rewards += [info['real_reward']]\n        invalid_action_stats += [info['invalid_action_stats']]\n        next_obs = np.array(next_obs)\n\n        # Annealing the rate if instructed to do so.\n        if args.anneal_lr:\n            pg_lr_scheduler.step()\n            vf_lr_scheduler.step()\n\n        if dones[step]:\n            # Computing the discounted returns:\n            writer.add_scalar(\"charts/episode_reward\", np.sum(real_rewards), global_step)\n            print(f\"global_step={global_step}, episode_reward={np.sum(real_rewards)}\")\n            for i in range(len(env.rfs)):\n                writer.add_scalar(f\"charts/episode_reward/{str(env.rfs[i])}\", raw_rewards.sum(1)[i], global_step)\n            real_rewards = []\n            for key, idx in zip(info['invalid_action_stats'], range(len(info['invalid_action_stats']))):\n                writer.add_scalar(f\"stats/{key}\", pd.DataFrame(invalid_action_stats).sum(0)[idx], global_step)\n            invalid_action_stats = []\n            next_obs = np.array(env.reset())\n\n    # bootstrap reward if not done. reached the batch limit\n    last_value = 0\n    if not dones[step]:\n        last_value = vf.forward(next_obs.reshape((1,)+next_obs.shape))[0].detach().cpu().numpy()[0]\n    bootstrapped_rewards = np.append(rewards, last_value)\n\n    # calculate the returns and advantages\n    if args.gae:\n        bootstrapped_values = np.append(values.detach().cpu().numpy(), last_value)\n        deltas = bootstrapped_rewards[:-1] + args.gamma * bootstrapped_values[1:] * (1-dones) - bootstrapped_values[:-1]\n        advantages = discount_cumsum(deltas, dones, args.gamma * args.gae_lambda)\n        advantages = torch.Tensor(advantages).to(device)\n        returns = advantages + values\n    else:\n        returns = discount_cumsum(bootstrapped_rewards, dones, args.gamma)[:-1]\n        advantages = returns - values.detach().cpu().numpy()\n        advantages = torch.Tensor(advantages).to(device)\n        returns = torch.Tensor(returns).to(device)\n\n    # Advantage normalization\n    if args.norm_adv:\n        EPS = 1e-10\n        advantages = (advantages - advantages.mean()) / (advantages.std() + EPS)\n\n    # Optimizaing policy network\n    entropys = []\n    target_pg = Policy().to(device)\n    inds = np.arange(args.batch_size,)\n    for i_epoch_pi in range(args.update_epochs):\n        np.random.shuffle(inds)\n        for start in range(0, args.batch_size, args.minibatch_size):\n            end = start + args.minibatch_size\n            minibatch_ind = inds[start:end]\n            target_pg.load_state_dict(pg.state_dict())\n            \n            _, newlogproba, _, _ = pg.get_action(\n                obs[minibatch_ind],\n                torch.LongTensor(actions[minibatch_ind].astype(np.int)).to(device).T,)\n            ratio = (newlogproba - logprobs[:,minibatch_ind]).exp()\n\n            # Policy loss as in OpenAI SpinUp\n            clip_adv = torch.where(advantages[minibatch_ind] > 0,\n                                    (1.+args.clip_coef) * advantages[minibatch_ind],\n                                    (1.-args.clip_coef) * advantages[minibatch_ind]).to(device)\n\n            # Entropy computation with resampled actions\n            entropy = -(newlogproba.exp() * newlogproba).mean()\n            entropys.append(entropy.item())\n\n            policy_loss = -torch.min(ratio * advantages[minibatch_ind], clip_adv) + args.ent_coef * entropy\n            policy_loss = policy_loss.mean()\n            \n            pg_optimizer.zero_grad()\n            policy_loss.backward()\n            nn.utils.clip_grad_norm_(pg.parameters(), args.max_grad_norm)\n            pg_optimizer.step()\n\n            approx_kl = (logprobs[:,minibatch_ind] - newlogproba).mean()\n            # Optimizing value network\n            new_values = vf.forward(obs[minibatch_ind]).view(-1)\n\n            # Value loss clipping\n            if args.clip_vloss:\n                v_loss_unclipped = ((new_values - returns[minibatch_ind]) ** 2)\n                v_clipped = values[minibatch_ind] + torch.clamp(new_values - values[minibatch_ind], -args.clip_coef, args.clip_coef)\n                v_loss_clipped = (v_clipped - returns[minibatch_ind])**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 = torch.mean((returns[minibatch_ind]- new_values).pow(2))\n\n            v_optimizer.zero_grad()\n            v_loss.backward()\n            nn.utils.clip_grad_norm_(vf.parameters(), args.max_grad_norm)\n            v_optimizer.step()\n\n        if args.kle_stop:\n            if approx_kl > args.target_kl:\n                break\n        if args.kle_rollback:\n            if (logprobs[:,minibatch_ind] - \n                pg.get_action(\n                    obs[minibatch_ind],\n                    torch.LongTensor(actions[minibatch_ind].astype(np.int)).to(device).T,\n                    invalid_action_masks[minibatch_ind])[1]).mean() > args.target_kl:\n                pg.load_state_dict(target_pg.state_dict())\n                break\n\n    # TRY NOT TO MODIFY: record rewards for plotting purposes\n    writer.add_scalar(\"losses/value_loss\", v_loss.item(), global_step)\n    writer.add_scalar(\"charts/policy_learning_rate\", pg_optimizer.param_groups[0]['lr'], global_step)\n    writer.add_scalar(\"charts/value_learning_rate\", v_optimizer.param_groups[0]['lr'], global_step)\n    writer.add_scalar(\"losses/policy_loss\", policy_loss.item(), global_step)\n    writer.add_scalar(\"losses/entropy\", np.mean(entropys), global_step)\n    writer.add_scalar(\"losses/approx_kl\", approx_kl.item(), global_step)\n    if args.kle_stop or args.kle_rollback:\n        writer.add_scalar(\"debug/pg_stop_iter\", i_epoch_pi, global_step)\n\nenv.close()\nwriter.close()\n"
  },
  {
    "path": "invalid_action_masking/ppo_no_adj_16x16.py",
    "content": "import torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport torch.nn.functional as F\nfrom torch.distributions.categorical import Categorical\nfrom torch.utils.tensorboard import SummaryWriter\n\nfrom cleanrl.common import preprocess_obs_space, preprocess_ac_space\nimport argparse\nimport numpy as np\nimport gym\nimport gym_microrts\nfrom gym.wrappers import TimeLimit, Monitor\nfrom gym.spaces import Discrete, Box, MultiBinary, MultiDiscrete, Space\nimport time\nimport random\nimport os\nimport pandas as pd\n\n# taken from https://github.com/openai/baselines/blob/master/baselines/common/vec_env/vec_normalize.py\nclass RunningMeanStd(object):\n    def __init__(self, epsilon=1e-4, shape=()):\n        self.mean = np.zeros(shape, 'float64')\n        self.var = np.ones(shape, 'float64')\n        self.count = epsilon\n\n    def update(self, x):\n        batch_mean = np.mean([x], axis=0)\n        batch_var = np.var([x], axis=0)\n        batch_count = 1\n        self.update_from_moments(batch_mean, batch_var, batch_count)\n\n    def update_from_moments(self, batch_mean, batch_var, batch_count):\n        self.mean, self.var, self.count = update_mean_var_count_from_moments(\n            self.mean, self.var, self.count, batch_mean, batch_var, batch_count)\n\ndef update_mean_var_count_from_moments(mean, var, count, batch_mean, batch_var, batch_count):\n    delta = batch_mean - mean\n    tot_count = count + batch_count\n\n    new_mean = mean + delta * batch_count / tot_count\n    m_a = var * count\n    m_b = batch_var * batch_count\n    M2 = m_a + m_b + np.square(delta) * count * batch_count / tot_count\n    new_var = M2 / tot_count\n    new_count = tot_count\n\n    return new_mean, new_var, new_count\nclass NormalizedEnv(gym.core.Wrapper):\n    def __init__(self, env, ob=True, ret=True, clipob=10., cliprew=10., gamma=0.99, epsilon=1e-8):\n        super(NormalizedEnv, self).__init__(env)\n        self.ob_rms = RunningMeanStd(shape=self.observation_space.shape) if ob else None\n        self.ret_rms = RunningMeanStd(shape=(1,)) if ret else None\n        self.clipob = clipob\n        self.cliprew = cliprew\n        self.ret = np.zeros(())\n        self.gamma = gamma\n        self.epsilon = epsilon\n\n    def step(self, action):\n        obs, rews, news, infos = self.env.step(action)\n        infos['real_reward'] = rews\n        # print(\"before\", self.ret)\n        self.ret = self.ret * self.gamma + rews\n        # print(\"after\", self.ret)\n        obs = self._obfilt(obs)\n        if self.ret_rms:\n            self.ret_rms.update(np.array([self.ret].copy()))\n            rews = np.clip(rews / np.sqrt(self.ret_rms.var + self.epsilon), -self.cliprew, self.cliprew)\n        self.ret = self.ret * (1-float(news))\n        return obs, rews, news, infos\n\n    def _obfilt(self, obs):\n        if self.ob_rms:\n            self.ob_rms.update(obs)\n            obs = np.clip((obs - self.ob_rms.mean) / np.sqrt(self.ob_rms.var + self.epsilon), -self.clipob, self.clipob)\n            return obs\n        else:\n            return obs\n\n    def reset(self):\n        self.ret = np.zeros(())\n        obs = self.env.reset()\n        return self._obfilt(obs)\n\nif __name__ == \"__main__\":\n    parser = argparse.ArgumentParser(description='PPO agent')\n    # Common arguments\n    parser.add_argument('--exp-name', type=str, default=os.path.basename(__file__).rstrip(\".py\"),\n                       help='the name of this experiment')\n    parser.add_argument('--gym-id', type=str, default=\"MicrortsMining16x16F9-v0\",\n                       help='the id of the gym environment')\n    parser.add_argument('--seed', type=int, default=1,\n                       help='seed of the experiment')\n    parser.add_argument('--episode-length', type=int, default=0,\n                       help='the maximum length of each episode')\n    parser.add_argument('--total-timesteps', type=int, default=100000,\n                       help='total timesteps of the experiments')\n    parser.add_argument('--no-torch-deterministic', action='store_false', dest=\"torch_deterministic\", default=True,\n                       help='if toggled, `torch.backends.cudnn.deterministic=False`')\n    parser.add_argument('--no-cuda', action='store_false', dest=\"cuda\", default=True,\n                       help='if toggled, cuda will not be enabled by default')\n    parser.add_argument('--prod-mode', action='store_true', default=False,\n                       help='run the script in production mode and use wandb to log outputs')\n    parser.add_argument('--capture-video', action='store_true', default=False,\n                       help='weather to capture videos of the agent performances (check out `videos` folder)')\n    parser.add_argument('--wandb-project-name', type=str, default=\"cleanRL\",\n                       help=\"the wandb's project name\")\n    parser.add_argument('--wandb-entity', type=str, default=None,\n                       help=\"the entity (team) of wandb's project\")\n\n    # Algorithm specific arguments\n    parser.add_argument('--batch-size', type=int, default=2048,\n                       help='the batch size of ppo')\n    parser.add_argument('--minibatch-size', type=int, default=256,\n                       help='the mini batch size of ppo')\n    parser.add_argument('--gamma', type=float, default=0.99,\n                       help='the discount factor gamma')\n    parser.add_argument('--gae-lambda', type=float, default=0.97,\n                       help='the lambda for the general advantage estimation')\n    parser.add_argument('--ent-coef', type=float, default=0.01,\n                       help=\"coefficient of the entropy\")\n    parser.add_argument('--max-grad-norm', type=float, default=0.5,\n                       help='the maximum norm for the gradient clipping')\n    parser.add_argument('--clip-coef', type=float, default=0.2,\n                       help=\"the surrogate clipping coefficient\")\n    parser.add_argument('--update-epochs', type=int, default=10,\n                        help=\"the K epochs to update the policy\")\n    parser.add_argument('--kle-stop', action='store_true', default=False,\n                        help='If toggled, the policy updates will be early stopped w.r.t target-kl')\n    parser.add_argument('--kle-rollback', action='store_true', default=False,\n                        help='If toggled, the policy updates will roll back to previous policy if KL exceeds target-kl')\n    parser.add_argument('--target-kl', type=float, default=0.015,\n                        help='the target-kl variable that is referred by --kl')\n    parser.add_argument('--gae', action='store_true', default=True,\n                        help='Use GAE for advantage computation')\n    parser.add_argument('--policy-lr', type=float, default=3e-4,\n                        help=\"the learning rate of the policy optimizer\")\n    parser.add_argument('--value-lr', type=float, default=3e-4,\n                        help=\"the learning rate of the critic optimizer\")\n    parser.add_argument('--norm-obs', action='store_true', default=True,\n                        help=\"Toggles observation normalization\")\n    parser.add_argument('--norm-returns', action='store_true', default=False,\n                        help=\"Toggles returns normalization\")\n    parser.add_argument('--norm-adv', action='store_true', default=True,\n                        help=\"Toggles advantages normalization\")\n    parser.add_argument('--obs-clip', type=float, default=10.0,\n                        help=\"Value for reward clipping, as per the paper\")\n    parser.add_argument('--rew-clip', type=float, default=10.0,\n                        help=\"Value for observation clipping, as per the paper\")\n    parser.add_argument('--anneal-lr', action='store_true', default=True,\n                        help=\"Toggle learning rate annealing for policy and value networks\")\n    parser.add_argument('--weights-init', default=\"orthogonal\", choices=[\"xavier\", 'orthogonal'],\n                        help='Selects the scheme to be used for weights initialization'),\n    parser.add_argument('--clip-vloss', action=\"store_true\", default=True,\n                        help='Toggles wheter or not to use a clipped loss for the value function, as per the paper.')\n    parser.add_argument('--pol-layer-norm', action='store_true', default=False,\n                       help='Enables layer normalization in the policy network')\n    args = parser.parse_args()\n    if not args.seed:\n        args.seed = int(time.time())\n\nargs.features_turned_on = sum([args.kle_stop, args.kle_rollback, args.gae, args.norm_obs, args.norm_returns, args.norm_adv, args.anneal_lr, args.clip_vloss, args.pol_layer_norm])\n# TRY NOT TO MODIFY: setup the environment\nexperiment_name = f\"{args.gym_id}__{args.exp_name}__{args.seed}__{int(time.time())}\"\nwriter = SummaryWriter(f\"runs/{experiment_name}\")\nwriter.add_text('hyperparameters', \"|param|value|\\n|-|-|\\n%s\" % (\n        '\\n'.join([f\"|{key}|{value}|\" for key, value in vars(args).items()])))\n\nif args.prod_mode:\n    import wandb\n    wandb.init(project=args.wandb_project_name, entity=args.wandb_entity, tensorboard=True, config=vars(args), name=experiment_name, monitor_gym=True)\n    writer = SummaryWriter(f\"/tmp/{experiment_name}\")\n    wandb.save(os.path.abspath(__file__))\n\n# TRY NOT TO MODIFY: seeding\ndevice = torch.device('cuda' if torch.cuda.is_available() and args.cuda else 'cpu')\nenv = gym.make(args.gym_id)\n# respect the default timelimit\nassert isinstance(env.action_space, MultiDiscrete), \"only MultiDiscrete action space is supported\"\nassert isinstance(env, TimeLimit) or int(args.episode_length), \"the gym env does not have a built in TimeLimit, please specify by using --episode-length\"\nif isinstance(env, TimeLimit):\n    if int(args.episode_length):\n        env._max_episode_steps = int(args.episode_length)\n    args.episode_length = env._max_episode_steps\nelse:\n    env = TimeLimit(env, int(args.episode_length))\nenv = NormalizedEnv(env.env, ob=args.norm_obs, ret=args.norm_returns, clipob=args.obs_clip, cliprew=args.rew_clip, gamma=args.gamma)\nenv = TimeLimit(env, int(args.episode_length))\nrandom.seed(args.seed)\nnp.random.seed(args.seed)\ntorch.manual_seed(args.seed)\ntorch.backends.cudnn.deterministic = args.torch_deterministic\nenv.seed(args.seed)\nenv.action_space.seed(args.seed)\nenv.observation_space.seed(args.seed)\nif args.capture_video:\n    env = Monitor(env, f'videos/{experiment_name}')\n\n# ALGO LOGIC: initialize agent here:\nclass CategoricalMasked(Categorical):\n    def __init__(self, probs=None, logits=None, validate_args=None, masks=[]):\n        self.masks = masks\n        if len(self.masks) == 0:\n            super(CategoricalMasked, self).__init__(probs, logits, validate_args)\n        else:\n            self.masks = masks.type(torch.BoolTensor).to(device)\n            logits = torch.where(self.masks, logits, torch.tensor(-1e+8).to(device))\n            super(CategoricalMasked, self).__init__(probs, logits, validate_args)\n    \n    def entropy(self):\n        if len(self.masks) == 0:\n            return super(CategoricalMasked, self).entropy()\n        p_log_p = self.logits * self.probs\n        p_log_p = torch.where(self.masks, p_log_p, torch.tensor(0.).to(device))\n        return -p_log_p.sum(-1)\n\nclass Policy(nn.Module):\n    def __init__(self):\n        super(Policy, self).__init__()\n        self.features = nn.Sequential(\n            nn.Conv2d(27, 16, kernel_size=3),\n            nn.MaxPool2d(2),\n            nn.ReLU(),\n            nn.Conv2d(16, 32, kernel_size=3),\n            nn.MaxPool2d(1),\n            nn.ReLU())\n        self.fc = nn.Sequential(\n            nn.Linear(32*5*5, 128),\n            nn.ReLU(),\n            nn.Linear(128, env.action_space.nvec.sum())\n        )\n\n    def forward(self, x):\n        x = torch.Tensor(np.moveaxis(x, -1, 1)).to(device)\n        x = self.features(x)\n        x = x.reshape(x.size(0), -1)\n        x = self.fc(x)\n        return x\n\n    def get_action(self, x, action=None, invalid_action_masks=None):\n        logits = self.forward(x)\n        split_logits = torch.split(logits, env.action_space.nvec.tolist(), dim=1)\n        \n        if invalid_action_masks is not None:\n            split_invalid_action_masks = torch.split(invalid_action_masks, env.action_space.nvec.tolist(), dim=1)\n            multi_categoricals = [CategoricalMasked(logits=logits, masks=iam) for (logits, iam) in zip(split_logits, split_invalid_action_masks)]\n        else:\n            multi_categoricals = [Categorical(logits=logits) for logits in split_logits]\n        \n        if action is None:\n            action = torch.stack([categorical.sample() for categorical in multi_categoricals])\n        logprob = torch.stack([categorical.log_prob(a) for a, categorical in zip(action, multi_categoricals)])\n        return action, logprob, [], multi_categoricals\n\nclass Value(nn.Module):\n    def __init__(self):\n        super(Value, self).__init__()\n        self.features = nn.Sequential(\n            nn.Conv2d(27, 16, kernel_size=3),\n            nn.MaxPool2d(2),\n            nn.ReLU(),\n            nn.Conv2d(16, 32, kernel_size=3),\n            nn.MaxPool2d(1),\n            nn.ReLU())\n        self.fc = nn.Sequential(\n            nn.Linear(32*5*5, 128),\n            nn.ReLU(),\n            nn.Linear(128, 1)\n        )\n\n    def forward(self, x):\n        x = torch.Tensor(np.moveaxis(x, -1, 1)).to(device)\n        x = self.features(x)\n        x = x.reshape(x.size(0), -1)\n        x = self.fc(x)\n        return x\n\ndef discount_cumsum(x, dones, gamma):\n    \"\"\"\n    computing discounted cumulative sums of vectors that resets with dones\n    input:\n        vector x,  vector dones,\n        [x0,       [0,\n         x1,        0,\n         x2         1,\n         x3         0, \n         x4]        0]\n    output:\n        [x0 + discount * x1 + discount^2 * x2,\n         x1 + discount * x2,\n         x2,\n         x3 + discount * x4,\n         x4]\n    \"\"\"\n    discount_cumsum = np.zeros_like(x)\n    discount_cumsum[-1] = x[-1]\n    for t in reversed(range(x.shape[0]-1)):\n        discount_cumsum[t] = x[t] + gamma * discount_cumsum[t+1] * (1-dones[t])\n    return discount_cumsum\n\npg = Policy().to(device)\nvf = Value().to(device)\n\n# MODIFIED: Separate optimizer and learning rates\npg_optimizer = optim.Adam(list(pg.parameters()), lr=args.policy_lr)\nv_optimizer = optim.Adam(list(vf.parameters()), lr=args.value_lr)\n\n# MODIFIED: Initializing learning rate anneal scheduler when need\nif args.anneal_lr:\n    anneal_fn = lambda f: max(0, 1-f / args.total_timesteps)\n    pg_lr_scheduler = optim.lr_scheduler.LambdaLR(pg_optimizer, lr_lambda=anneal_fn)\n    vf_lr_scheduler = optim.lr_scheduler.LambdaLR(v_optimizer, lr_lambda=anneal_fn)\n\nloss_fn = nn.MSELoss()\n\n# TRY NOT TO MODIFY: start the game\nglobal_step = 0\nwhile global_step < args.total_timesteps:\n    if args.capture_video:\n        env.stats_recorder.done=True\n    next_obs = np.array(env.reset())\n\n    # ALGO Logic: Storage for epoch data\n    obs = np.empty((args.batch_size,) + env.observation_space.shape)\n\n    actions = np.empty((args.batch_size,) + env.action_space.shape)\n    logprobs = torch.zeros((env.action_space.nvec.shape[0], args.batch_size,)).to(device)\n\n    rewards = np.zeros((args.batch_size,))\n    raw_rewards = np.zeros((len(env.rfs),args.batch_size,))\n    \n    real_rewards = []\n    invalid_action_stats = []\n\n    dones = np.zeros((args.batch_size,))\n    values = torch.zeros((args.batch_size,)).to(device)\n\n    invalid_action_masks = torch.zeros((args.batch_size, env.action_space.nvec.sum()))\n    # TRY NOT TO MODIFY: prepare the execution of the game.\n    for step in range(args.batch_size):\n        env.render()\n        global_step += 1\n        obs[step] = next_obs.copy()\n\n        # ALGO LOGIC: put action logic here\n        invalid_action_mask = torch.ones(env.action_space.nvec.sum())\n        invalid_action_mask[0:env.action_space.nvec[0]] = torch.tensor(env.unit_location_mask)\n        invalid_action_mask[-env.action_space.nvec[-1]:] = torch.tensor(env.target_unit_location_mask)\n        invalid_action_masks[step] = invalid_action_mask\n        with torch.no_grad():\n            values[step] = vf.forward(obs[step:step+1])\n            action, logproba, _, probs = pg.get_action(obs[step:step+1], invalid_action_masks=invalid_action_masks[step:step+1])\n            \n            # CORE LOGIC:\n            # use the action generated by CategoricalMasked, but \n            # don't adjust the logprobability accordingly. Instead, calculate the log\n            # probability using Categorical\n            action, logproba, _, probs = pg.get_action(obs[step:step+1], action=action)\n        \n        actions[step] = action[:,0].data.cpu().numpy()\n        logprobs[:,[step]] = logproba\n\n        # TRY NOT TO MODIFY: execute the game and log data.\n        next_obs, rewards[step], dones[step], info = env.step(action[:,0].data.cpu().numpy())\n        raw_rewards[:,step] = info[\"rewards\"]\n        real_rewards += [info['real_reward']]\n        invalid_action_stats += [info['invalid_action_stats']]\n        next_obs = np.array(next_obs)\n\n        # Annealing the rate if instructed to do so.\n        if args.anneal_lr:\n            pg_lr_scheduler.step()\n            vf_lr_scheduler.step()\n\n        if dones[step]:\n            # Computing the discounted returns:\n            writer.add_scalar(\"charts/episode_reward\", np.sum(real_rewards), global_step)\n            print(f\"global_step={global_step}, episode_reward={np.sum(real_rewards)}\")\n            for i in range(len(env.rfs)):\n                writer.add_scalar(f\"charts/episode_reward/{str(env.rfs[i])}\", raw_rewards.sum(1)[i], global_step)\n            real_rewards = []\n            for key, idx in zip(info['invalid_action_stats'], range(len(info['invalid_action_stats']))):\n                writer.add_scalar(f\"stats/{key}\", pd.DataFrame(invalid_action_stats).sum(0)[idx], global_step)\n            invalid_action_stats = []\n            next_obs = np.array(env.reset())\n\n    # bootstrap reward if not done. reached the batch limit\n    last_value = 0\n    if not dones[step]:\n        last_value = vf.forward(next_obs.reshape((1,)+next_obs.shape))[0].detach().cpu().numpy()[0]\n    bootstrapped_rewards = np.append(rewards, last_value)\n\n    # calculate the returns and advantages\n    if args.gae:\n        bootstrapped_values = np.append(values.detach().cpu().numpy(), last_value)\n        deltas = bootstrapped_rewards[:-1] + args.gamma * bootstrapped_values[1:] * (1-dones) - bootstrapped_values[:-1]\n        advantages = discount_cumsum(deltas, dones, args.gamma * args.gae_lambda)\n        advantages = torch.Tensor(advantages).to(device)\n        returns = advantages + values\n    else:\n        returns = discount_cumsum(bootstrapped_rewards, dones, args.gamma)[:-1]\n        advantages = returns - values.detach().cpu().numpy()\n        advantages = torch.Tensor(advantages).to(device)\n        returns = torch.Tensor(returns).to(device)\n\n    # Advantage normalization\n    if args.norm_adv:\n        EPS = 1e-10\n        advantages = (advantages - advantages.mean()) / (advantages.std() + EPS)\n\n    # Optimizaing policy network\n    entropys = []\n    target_pg = Policy().to(device)\n    inds = np.arange(args.batch_size,)\n    for i_epoch_pi in range(args.update_epochs):\n        np.random.shuffle(inds)\n        for start in range(0, args.batch_size, args.minibatch_size):\n            end = start + args.minibatch_size\n            minibatch_ind = inds[start:end]\n            target_pg.load_state_dict(pg.state_dict())\n            \n            _, newlogproba, _, _ = pg.get_action(\n                obs[minibatch_ind],\n                torch.LongTensor(actions[minibatch_ind].astype(np.int)).to(device).T,)\n            ratio = (newlogproba - logprobs[:,minibatch_ind]).exp()\n\n            # Policy loss as in OpenAI SpinUp\n            clip_adv = torch.where(advantages[minibatch_ind] > 0,\n                                    (1.+args.clip_coef) * advantages[minibatch_ind],\n                                    (1.-args.clip_coef) * advantages[minibatch_ind]).to(device)\n\n            # Entropy computation with resampled actions\n            entropy = -(newlogproba.exp() * newlogproba).mean()\n            entropys.append(entropy.item())\n\n            policy_loss = -torch.min(ratio * advantages[minibatch_ind], clip_adv) + args.ent_coef * entropy\n            policy_loss = policy_loss.mean()\n            \n            pg_optimizer.zero_grad()\n            policy_loss.backward()\n            nn.utils.clip_grad_norm_(pg.parameters(), args.max_grad_norm)\n            pg_optimizer.step()\n\n            approx_kl = (logprobs[:,minibatch_ind] - newlogproba).mean()\n            # Optimizing value network\n            new_values = vf.forward(obs[minibatch_ind]).view(-1)\n\n            # Value loss clipping\n            if args.clip_vloss:\n                v_loss_unclipped = ((new_values - returns[minibatch_ind]) ** 2)\n                v_clipped = values[minibatch_ind] + torch.clamp(new_values - values[minibatch_ind], -args.clip_coef, args.clip_coef)\n                v_loss_clipped = (v_clipped - returns[minibatch_ind])**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 = torch.mean((returns[minibatch_ind]- new_values).pow(2))\n\n            v_optimizer.zero_grad()\n            v_loss.backward()\n            nn.utils.clip_grad_norm_(vf.parameters(), args.max_grad_norm)\n            v_optimizer.step()\n\n        if args.kle_stop:\n            if approx_kl > args.target_kl:\n                break\n        if args.kle_rollback:\n            if (logprobs[:,minibatch_ind] - \n                pg.get_action(\n                    obs[minibatch_ind],\n                    torch.LongTensor(actions[minibatch_ind].astype(np.int)).to(device).T,\n                    invalid_action_masks[minibatch_ind])[1]).mean() > args.target_kl:\n                pg.load_state_dict(target_pg.state_dict())\n                break\n\n    # TRY NOT TO MODIFY: record rewards for plotting purposes\n    writer.add_scalar(\"losses/value_loss\", v_loss.item(), global_step)\n    writer.add_scalar(\"charts/policy_learning_rate\", pg_optimizer.param_groups[0]['lr'], global_step)\n    writer.add_scalar(\"charts/value_learning_rate\", v_optimizer.param_groups[0]['lr'], global_step)\n    writer.add_scalar(\"losses/policy_loss\", policy_loss.item(), global_step)\n    writer.add_scalar(\"losses/entropy\", np.mean(entropys), global_step)\n    writer.add_scalar(\"losses/approx_kl\", approx_kl.item(), global_step)\n    if args.kle_stop or args.kle_rollback:\n        writer.add_scalar(\"debug/pg_stop_iter\", i_epoch_pi, global_step)\n\nenv.close()\nwriter.close()\n"
  },
  {
    "path": "invalid_action_masking/ppo_no_adj_24x24.py",
    "content": "import torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport torch.nn.functional as F\nfrom torch.distributions.categorical import Categorical\nfrom torch.utils.tensorboard import SummaryWriter\n\nfrom cleanrl.common import preprocess_obs_space, preprocess_ac_space\nimport argparse\nimport numpy as np\nimport gym\nimport gym_microrts\nfrom gym.wrappers import TimeLimit, Monitor\nfrom gym.spaces import Discrete, Box, MultiBinary, MultiDiscrete, Space\nimport time\nimport random\nimport os\nimport pandas as pd\n\n# taken from https://github.com/openai/baselines/blob/master/baselines/common/vec_env/vec_normalize.py\nclass RunningMeanStd(object):\n    def __init__(self, epsilon=1e-4, shape=()):\n        self.mean = np.zeros(shape, 'float64')\n        self.var = np.ones(shape, 'float64')\n        self.count = epsilon\n\n    def update(self, x):\n        batch_mean = np.mean([x], axis=0)\n        batch_var = np.var([x], axis=0)\n        batch_count = 1\n        self.update_from_moments(batch_mean, batch_var, batch_count)\n\n    def update_from_moments(self, batch_mean, batch_var, batch_count):\n        self.mean, self.var, self.count = update_mean_var_count_from_moments(\n            self.mean, self.var, self.count, batch_mean, batch_var, batch_count)\n\ndef update_mean_var_count_from_moments(mean, var, count, batch_mean, batch_var, batch_count):\n    delta = batch_mean - mean\n    tot_count = count + batch_count\n\n    new_mean = mean + delta * batch_count / tot_count\n    m_a = var * count\n    m_b = batch_var * batch_count\n    M2 = m_a + m_b + np.square(delta) * count * batch_count / tot_count\n    new_var = M2 / tot_count\n    new_count = tot_count\n\n    return new_mean, new_var, new_count\nclass NormalizedEnv(gym.core.Wrapper):\n    def __init__(self, env, ob=True, ret=True, clipob=10., cliprew=10., gamma=0.99, epsilon=1e-8):\n        super(NormalizedEnv, self).__init__(env)\n        self.ob_rms = RunningMeanStd(shape=self.observation_space.shape) if ob else None\n        self.ret_rms = RunningMeanStd(shape=(1,)) if ret else None\n        self.clipob = clipob\n        self.cliprew = cliprew\n        self.ret = np.zeros(())\n        self.gamma = gamma\n        self.epsilon = epsilon\n\n    def step(self, action):\n        obs, rews, news, infos = self.env.step(action)\n        infos['real_reward'] = rews\n        # print(\"before\", self.ret)\n        self.ret = self.ret * self.gamma + rews\n        # print(\"after\", self.ret)\n        obs = self._obfilt(obs)\n        if self.ret_rms:\n            self.ret_rms.update(np.array([self.ret].copy()))\n            rews = np.clip(rews / np.sqrt(self.ret_rms.var + self.epsilon), -self.cliprew, self.cliprew)\n        self.ret = self.ret * (1-float(news))\n        return obs, rews, news, infos\n\n    def _obfilt(self, obs):\n        if self.ob_rms:\n            self.ob_rms.update(obs)\n            obs = np.clip((obs - self.ob_rms.mean) / np.sqrt(self.ob_rms.var + self.epsilon), -self.clipob, self.clipob)\n            return obs\n        else:\n            return obs\n\n    def reset(self):\n        self.ret = np.zeros(())\n        obs = self.env.reset()\n        return self._obfilt(obs)\n\nif __name__ == \"__main__\":\n    parser = argparse.ArgumentParser(description='PPO agent')\n    # Common arguments\n    parser.add_argument('--exp-name', type=str, default=os.path.basename(__file__).rstrip(\".py\"),\n                       help='the name of this experiment')\n    parser.add_argument('--gym-id', type=str, default=\"MicrortsMining24x24F9-v0\",\n                       help='the id of the gym environment')\n    parser.add_argument('--seed', type=int, default=1,\n                       help='seed of the experiment')\n    parser.add_argument('--episode-length', type=int, default=0,\n                       help='the maximum length of each episode')\n    parser.add_argument('--total-timesteps', type=int, default=100000,\n                       help='total timesteps of the experiments')\n    parser.add_argument('--no-torch-deterministic', action='store_false', dest=\"torch_deterministic\", default=True,\n                       help='if toggled, `torch.backends.cudnn.deterministic=False`')\n    parser.add_argument('--no-cuda', action='store_false', dest=\"cuda\", default=True,\n                       help='if toggled, cuda will not be enabled by default')\n    parser.add_argument('--prod-mode', action='store_true', default=False,\n                       help='run the script in production mode and use wandb to log outputs')\n    parser.add_argument('--capture-video', action='store_true', default=False,\n                       help='weather to capture videos of the agent performances (check out `videos` folder)')\n    parser.add_argument('--wandb-project-name', type=str, default=\"cleanRL\",\n                       help=\"the wandb's project name\")\n    parser.add_argument('--wandb-entity', type=str, default=None,\n                       help=\"the entity (team) of wandb's project\")\n\n    # Algorithm specific arguments\n    parser.add_argument('--batch-size', type=int, default=2048,\n                       help='the batch size of ppo')\n    parser.add_argument('--minibatch-size', type=int, default=256,\n                       help='the mini batch size of ppo')\n    parser.add_argument('--gamma', type=float, default=0.99,\n                       help='the discount factor gamma')\n    parser.add_argument('--gae-lambda', type=float, default=0.97,\n                       help='the lambda for the general advantage estimation')\n    parser.add_argument('--ent-coef', type=float, default=0.01,\n                       help=\"coefficient of the entropy\")\n    parser.add_argument('--max-grad-norm', type=float, default=0.5,\n                       help='the maximum norm for the gradient clipping')\n    parser.add_argument('--clip-coef', type=float, default=0.2,\n                       help=\"the surrogate clipping coefficient\")\n    parser.add_argument('--update-epochs', type=int, default=10,\n                        help=\"the K epochs to update the policy\")\n    parser.add_argument('--kle-stop', action='store_true', default=False,\n                        help='If toggled, the policy updates will be early stopped w.r.t target-kl')\n    parser.add_argument('--kle-rollback', action='store_true', default=False,\n                        help='If toggled, the policy updates will roll back to previous policy if KL exceeds target-kl')\n    parser.add_argument('--target-kl', type=float, default=0.015,\n                        help='the target-kl variable that is referred by --kl')\n    parser.add_argument('--gae', action='store_true', default=True,\n                        help='Use GAE for advantage computation')\n    parser.add_argument('--policy-lr', type=float, default=3e-4,\n                        help=\"the learning rate of the policy optimizer\")\n    parser.add_argument('--value-lr', type=float, default=3e-4,\n                        help=\"the learning rate of the critic optimizer\")\n    parser.add_argument('--norm-obs', action='store_true', default=True,\n                        help=\"Toggles observation normalization\")\n    parser.add_argument('--norm-returns', action='store_true', default=False,\n                        help=\"Toggles returns normalization\")\n    parser.add_argument('--norm-adv', action='store_true', default=True,\n                        help=\"Toggles advantages normalization\")\n    parser.add_argument('--obs-clip', type=float, default=10.0,\n                        help=\"Value for reward clipping, as per the paper\")\n    parser.add_argument('--rew-clip', type=float, default=10.0,\n                        help=\"Value for observation clipping, as per the paper\")\n    parser.add_argument('--anneal-lr', action='store_true', default=True,\n                        help=\"Toggle learning rate annealing for policy and value networks\")\n    parser.add_argument('--weights-init', default=\"orthogonal\", choices=[\"xavier\", 'orthogonal'],\n                        help='Selects the scheme to be used for weights initialization'),\n    parser.add_argument('--clip-vloss', action=\"store_true\", default=True,\n                        help='Toggles wheter or not to use a clipped loss for the value function, as per the paper.')\n    parser.add_argument('--pol-layer-norm', action='store_true', default=False,\n                       help='Enables layer normalization in the policy network')\n    args = parser.parse_args()\n    if not args.seed:\n        args.seed = int(time.time())\n\nargs.features_turned_on = sum([args.kle_stop, args.kle_rollback, args.gae, args.norm_obs, args.norm_returns, args.norm_adv, args.anneal_lr, args.clip_vloss, args.pol_layer_norm])\n# TRY NOT TO MODIFY: setup the environment\nexperiment_name = f\"{args.gym_id}__{args.exp_name}__{args.seed}__{int(time.time())}\"\nwriter = SummaryWriter(f\"runs/{experiment_name}\")\nwriter.add_text('hyperparameters', \"|param|value|\\n|-|-|\\n%s\" % (\n        '\\n'.join([f\"|{key}|{value}|\" for key, value in vars(args).items()])))\n\nif args.prod_mode:\n    import wandb\n    wandb.init(project=args.wandb_project_name, entity=args.wandb_entity, tensorboard=True, config=vars(args), name=experiment_name, monitor_gym=True)\n    writer = SummaryWriter(f\"/tmp/{experiment_name}\")\n    wandb.save(os.path.abspath(__file__))\n\n# TRY NOT TO MODIFY: seeding\ndevice = torch.device('cuda' if torch.cuda.is_available() and args.cuda else 'cpu')\nenv = gym.make(args.gym_id)\n# respect the default timelimit\nassert isinstance(env.action_space, MultiDiscrete), \"only MultiDiscrete action space is supported\"\nassert isinstance(env, TimeLimit) or int(args.episode_length), \"the gym env does not have a built in TimeLimit, please specify by using --episode-length\"\nif isinstance(env, TimeLimit):\n    if int(args.episode_length):\n        env._max_episode_steps = int(args.episode_length)\n    args.episode_length = env._max_episode_steps\nelse:\n    env = TimeLimit(env, int(args.episode_length))\nenv = NormalizedEnv(env.env, ob=args.norm_obs, ret=args.norm_returns, clipob=args.obs_clip, cliprew=args.rew_clip, gamma=args.gamma)\nenv = TimeLimit(env, int(args.episode_length))\nrandom.seed(args.seed)\nnp.random.seed(args.seed)\ntorch.manual_seed(args.seed)\ntorch.backends.cudnn.deterministic = args.torch_deterministic\nenv.seed(args.seed)\nenv.action_space.seed(args.seed)\nenv.observation_space.seed(args.seed)\nif args.capture_video:\n    env = Monitor(env, f'videos/{experiment_name}')\n\n# ALGO LOGIC: initialize agent here:\nclass CategoricalMasked(Categorical):\n    def __init__(self, probs=None, logits=None, validate_args=None, masks=[]):\n        self.masks = masks\n        if len(self.masks) == 0:\n            super(CategoricalMasked, self).__init__(probs, logits, validate_args)\n        else:\n            self.masks = masks.type(torch.BoolTensor).to(device)\n            logits = torch.where(self.masks, logits, torch.tensor(-1e+8).to(device))\n            super(CategoricalMasked, self).__init__(probs, logits, validate_args)\n    \n    def entropy(self):\n        if len(self.masks) == 0:\n            return super(CategoricalMasked, self).entropy()\n        p_log_p = self.logits * self.probs\n        p_log_p = torch.where(self.masks, p_log_p, torch.tensor(0.).to(device))\n        return -p_log_p.sum(-1)\n\nclass Policy(nn.Module):\n    def __init__(self):\n        super(Policy, self).__init__()\n        self.features = nn.Sequential(\n            nn.Conv2d(27, 16, kernel_size=3, stride=1),\n            nn.MaxPool2d(2),\n            nn.ReLU(),\n            nn.Conv2d(16, 32, kernel_size=2, stride=1),\n            nn.MaxPool2d(2),\n            nn.ReLU())\n        self.fc = nn.Sequential(\n            nn.Linear(32*5*5, 128),\n            nn.ReLU(),\n            nn.Linear(128, env.action_space.nvec.sum())\n        )\n\n    def forward(self, x):\n        x = torch.Tensor(np.moveaxis(x, -1, 1)).to(device)\n        x = self.features(x)\n        x = x.reshape(x.size(0), -1)\n        x = self.fc(x)\n        return x\n\n    def get_action(self, x, action=None, invalid_action_masks=None):\n        logits = self.forward(x)\n        split_logits = torch.split(logits, env.action_space.nvec.tolist(), dim=1)\n        \n        if invalid_action_masks is not None:\n            split_invalid_action_masks = torch.split(invalid_action_masks, env.action_space.nvec.tolist(), dim=1)\n            multi_categoricals = [CategoricalMasked(logits=logits, masks=iam) for (logits, iam) in zip(split_logits, split_invalid_action_masks)]\n        else:\n            multi_categoricals = [Categorical(logits=logits) for logits in split_logits]\n        \n        if action is None:\n            action = torch.stack([categorical.sample() for categorical in multi_categoricals])\n        logprob = torch.stack([categorical.log_prob(a) for a, categorical in zip(action, multi_categoricals)])\n        return action, logprob, [], multi_categoricals\n\nclass Value(nn.Module):\n    def __init__(self):\n        super(Value, self).__init__()\n        self.features = nn.Sequential(\n            nn.Conv2d(27, 16, kernel_size=3, stride=1),\n            nn.MaxPool2d(2),\n            nn.ReLU(),\n            nn.Conv2d(16, 32, kernel_size=2, stride=1),\n            nn.MaxPool2d(2),\n            nn.ReLU())\n        self.fc = nn.Sequential(\n            nn.Linear(32*5*5, 128),\n            nn.ReLU(),\n            nn.Linear(128, 1)\n        )\n\n    def forward(self, x):\n        x = torch.Tensor(np.moveaxis(x, -1, 1)).to(device)\n        x = self.features(x)\n        x = x.reshape(x.size(0), -1)\n        x = self.fc(x)\n        return x\n\ndef discount_cumsum(x, dones, gamma):\n    \"\"\"\n    computing discounted cumulative sums of vectors that resets with dones\n    input:\n        vector x,  vector dones,\n        [x0,       [0,\n         x1,        0,\n         x2         1,\n         x3         0, \n         x4]        0]\n    output:\n        [x0 + discount * x1 + discount^2 * x2,\n         x1 + discount * x2,\n         x2,\n         x3 + discount * x4,\n         x4]\n    \"\"\"\n    discount_cumsum = np.zeros_like(x)\n    discount_cumsum[-1] = x[-1]\n    for t in reversed(range(x.shape[0]-1)):\n        discount_cumsum[t] = x[t] + gamma * discount_cumsum[t+1] * (1-dones[t])\n    return discount_cumsum\n\npg = Policy().to(device)\nvf = Value().to(device)\n\n# MODIFIED: Separate optimizer and learning rates\npg_optimizer = optim.Adam(list(pg.parameters()), lr=args.policy_lr)\nv_optimizer = optim.Adam(list(vf.parameters()), lr=args.value_lr)\n\n# MODIFIED: Initializing learning rate anneal scheduler when need\nif args.anneal_lr:\n    anneal_fn = lambda f: max(0, 1-f / args.total_timesteps)\n    pg_lr_scheduler = optim.lr_scheduler.LambdaLR(pg_optimizer, lr_lambda=anneal_fn)\n    vf_lr_scheduler = optim.lr_scheduler.LambdaLR(v_optimizer, lr_lambda=anneal_fn)\n\nloss_fn = nn.MSELoss()\n\n# TRY NOT TO MODIFY: start the game\nglobal_step = 0\nwhile global_step < args.total_timesteps:\n    if args.capture_video:\n        env.stats_recorder.done=True\n    next_obs = np.array(env.reset())\n\n    # ALGO Logic: Storage for epoch data\n    obs = np.empty((args.batch_size,) + env.observation_space.shape)\n\n    actions = np.empty((args.batch_size,) + env.action_space.shape)\n    logprobs = torch.zeros((env.action_space.nvec.shape[0], args.batch_size,)).to(device)\n\n    rewards = np.zeros((args.batch_size,))\n    raw_rewards = np.zeros((len(env.rfs),args.batch_size,))\n    \n    real_rewards = []\n    invalid_action_stats = []\n\n    dones = np.zeros((args.batch_size,))\n    values = torch.zeros((args.batch_size,)).to(device)\n\n    invalid_action_masks = torch.zeros((args.batch_size, env.action_space.nvec.sum()))\n    # TRY NOT TO MODIFY: prepare the execution of the game.\n    for step in range(args.batch_size):\n        env.render()\n        global_step += 1\n        obs[step] = next_obs.copy()\n\n        # ALGO LOGIC: put action logic here\n        invalid_action_mask = torch.ones(env.action_space.nvec.sum())\n        invalid_action_mask[0:env.action_space.nvec[0]] = torch.tensor(env.unit_location_mask)\n        invalid_action_mask[-env.action_space.nvec[-1]:] = torch.tensor(env.target_unit_location_mask)\n        invalid_action_masks[step] = invalid_action_mask\n        with torch.no_grad():\n            values[step] = vf.forward(obs[step:step+1])\n            action, logproba, _, probs = pg.get_action(obs[step:step+1], invalid_action_masks=invalid_action_masks[step:step+1])\n            \n            # CORE LOGIC:\n            # use the action generated by CategoricalMasked, but \n            # don't adjust the logprobability accordingly. Instead, calculate the log\n            # probability using Categorical\n            action, logproba, _, probs = pg.get_action(obs[step:step+1], action=action)\n        \n        actions[step] = action[:,0].data.cpu().numpy()\n        logprobs[:,[step]] = logproba\n\n        # TRY NOT TO MODIFY: execute the game and log data.\n        next_obs, rewards[step], dones[step], info = env.step(action[:,0].data.cpu().numpy())\n        raw_rewards[:,step] = info[\"rewards\"]\n        real_rewards += [info['real_reward']]\n        invalid_action_stats += [info['invalid_action_stats']]\n        next_obs = np.array(next_obs)\n\n        # Annealing the rate if instructed to do so.\n        if args.anneal_lr:\n            pg_lr_scheduler.step()\n            vf_lr_scheduler.step()\n\n        if dones[step]:\n            # Computing the discounted returns:\n            writer.add_scalar(\"charts/episode_reward\", np.sum(real_rewards), global_step)\n            print(f\"global_step={global_step}, episode_reward={np.sum(real_rewards)}\")\n            for i in range(len(env.rfs)):\n                writer.add_scalar(f\"charts/episode_reward/{str(env.rfs[i])}\", raw_rewards.sum(1)[i], global_step)\n            real_rewards = []\n            for key, idx in zip(info['invalid_action_stats'], range(len(info['invalid_action_stats']))):\n                writer.add_scalar(f\"stats/{key}\", pd.DataFrame(invalid_action_stats).sum(0)[idx], global_step)\n            invalid_action_stats = []\n            next_obs = np.array(env.reset())\n\n    # bootstrap reward if not done. reached the batch limit\n    last_value = 0\n    if not dones[step]:\n        last_value = vf.forward(next_obs.reshape((1,)+next_obs.shape))[0].detach().cpu().numpy()[0]\n    bootstrapped_rewards = np.append(rewards, last_value)\n\n    # calculate the returns and advantages\n    if args.gae:\n        bootstrapped_values = np.append(values.detach().cpu().numpy(), last_value)\n        deltas = bootstrapped_rewards[:-1] + args.gamma * bootstrapped_values[1:] * (1-dones) - bootstrapped_values[:-1]\n        advantages = discount_cumsum(deltas, dones, args.gamma * args.gae_lambda)\n        advantages = torch.Tensor(advantages).to(device)\n        returns = advantages + values\n    else:\n        returns = discount_cumsum(bootstrapped_rewards, dones, args.gamma)[:-1]\n        advantages = returns - values.detach().cpu().numpy()\n        advantages = torch.Tensor(advantages).to(device)\n        returns = torch.Tensor(returns).to(device)\n\n    # Advantage normalization\n    if args.norm_adv:\n        EPS = 1e-10\n        advantages = (advantages - advantages.mean()) / (advantages.std() + EPS)\n\n    # Optimizaing policy network\n    entropys = []\n    target_pg = Policy().to(device)\n    inds = np.arange(args.batch_size,)\n    for i_epoch_pi in range(args.update_epochs):\n        np.random.shuffle(inds)\n        for start in range(0, args.batch_size, args.minibatch_size):\n            end = start + args.minibatch_size\n            minibatch_ind = inds[start:end]\n            target_pg.load_state_dict(pg.state_dict())\n            \n            _, newlogproba, _, _ = pg.get_action(\n                obs[minibatch_ind],\n                torch.LongTensor(actions[minibatch_ind].astype(np.int)).to(device).T,)\n            ratio = (newlogproba - logprobs[:,minibatch_ind]).exp()\n\n            # Policy loss as in OpenAI SpinUp\n            clip_adv = torch.where(advantages[minibatch_ind] > 0,\n                                    (1.+args.clip_coef) * advantages[minibatch_ind],\n                                    (1.-args.clip_coef) * advantages[minibatch_ind]).to(device)\n\n            # Entropy computation with resampled actions\n            entropy = -(newlogproba.exp() * newlogproba).mean()\n            entropys.append(entropy.item())\n\n            policy_loss = -torch.min(ratio * advantages[minibatch_ind], clip_adv) + args.ent_coef * entropy\n            policy_loss = policy_loss.mean()\n            \n            pg_optimizer.zero_grad()\n            policy_loss.backward()\n            nn.utils.clip_grad_norm_(pg.parameters(), args.max_grad_norm)\n            pg_optimizer.step()\n\n            approx_kl = (logprobs[:,minibatch_ind] - newlogproba).mean()\n            # Optimizing value network\n            new_values = vf.forward(obs[minibatch_ind]).view(-1)\n\n            # Value loss clipping\n            if args.clip_vloss:\n                v_loss_unclipped = ((new_values - returns[minibatch_ind]) ** 2)\n                v_clipped = values[minibatch_ind] + torch.clamp(new_values - values[minibatch_ind], -args.clip_coef, args.clip_coef)\n                v_loss_clipped = (v_clipped - returns[minibatch_ind])**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 = torch.mean((returns[minibatch_ind]- new_values).pow(2))\n\n            v_optimizer.zero_grad()\n            v_loss.backward()\n            nn.utils.clip_grad_norm_(vf.parameters(), args.max_grad_norm)\n            v_optimizer.step()\n\n        if args.kle_stop:\n            if approx_kl > args.target_kl:\n                break\n        if args.kle_rollback:\n            if (logprobs[:,minibatch_ind] - \n                pg.get_action(\n                    obs[minibatch_ind],\n                    torch.LongTensor(actions[minibatch_ind].astype(np.int)).to(device).T,\n                    invalid_action_masks[minibatch_ind])[1]).mean() > args.target_kl:\n                pg.load_state_dict(target_pg.state_dict())\n                break\n\n    # TRY NOT TO MODIFY: record rewards for plotting purposes\n    writer.add_scalar(\"losses/value_loss\", v_loss.item(), global_step)\n    writer.add_scalar(\"charts/policy_learning_rate\", pg_optimizer.param_groups[0]['lr'], global_step)\n    writer.add_scalar(\"charts/value_learning_rate\", v_optimizer.param_groups[0]['lr'], global_step)\n    writer.add_scalar(\"losses/policy_loss\", policy_loss.item(), global_step)\n    writer.add_scalar(\"losses/entropy\", np.mean(entropys), global_step)\n    writer.add_scalar(\"losses/approx_kl\", approx_kl.item(), global_step)\n    if args.kle_stop or args.kle_rollback:\n        writer.add_scalar(\"debug/pg_stop_iter\", i_epoch_pi, global_step)\n\nenv.close()\nwriter.close()\n"
  },
  {
    "path": "invalid_action_masking/ppo_no_adj_4x4.py",
    "content": "import torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport torch.nn.functional as F\nfrom torch.distributions.categorical import Categorical\nfrom torch.utils.tensorboard import SummaryWriter\n\nfrom cleanrl.common import preprocess_obs_space, preprocess_ac_space\nimport argparse\nimport numpy as np\nimport gym\nimport gym_microrts\nfrom gym.wrappers import TimeLimit, Monitor\nfrom gym.spaces import Discrete, Box, MultiBinary, MultiDiscrete, Space\nimport time\nimport random\nimport os\nimport pandas as pd\n\n# taken from https://github.com/openai/baselines/blob/master/baselines/common/vec_env/vec_normalize.py\nclass RunningMeanStd(object):\n    def __init__(self, epsilon=1e-4, shape=()):\n        self.mean = np.zeros(shape, 'float64')\n        self.var = np.ones(shape, 'float64')\n        self.count = epsilon\n\n    def update(self, x):\n        batch_mean = np.mean([x], axis=0)\n        batch_var = np.var([x], axis=0)\n        batch_count = 1\n        self.update_from_moments(batch_mean, batch_var, batch_count)\n\n    def update_from_moments(self, batch_mean, batch_var, batch_count):\n        self.mean, self.var, self.count = update_mean_var_count_from_moments(\n            self.mean, self.var, self.count, batch_mean, batch_var, batch_count)\n\ndef update_mean_var_count_from_moments(mean, var, count, batch_mean, batch_var, batch_count):\n    delta = batch_mean - mean\n    tot_count = count + batch_count\n\n    new_mean = mean + delta * batch_count / tot_count\n    m_a = var * count\n    m_b = batch_var * batch_count\n    M2 = m_a + m_b + np.square(delta) * count * batch_count / tot_count\n    new_var = M2 / tot_count\n    new_count = tot_count\n\n    return new_mean, new_var, new_count\nclass NormalizedEnv(gym.core.Wrapper):\n    def __init__(self, env, ob=True, ret=True, clipob=10., cliprew=10., gamma=0.99, epsilon=1e-8):\n        super(NormalizedEnv, self).__init__(env)\n        self.ob_rms = RunningMeanStd(shape=self.observation_space.shape) if ob else None\n        self.ret_rms = RunningMeanStd(shape=(1,)) if ret else None\n        self.clipob = clipob\n        self.cliprew = cliprew\n        self.ret = np.zeros(())\n        self.gamma = gamma\n        self.epsilon = epsilon\n\n    def step(self, action):\n        obs, rews, news, infos = self.env.step(action)\n        infos['real_reward'] = rews\n        # print(\"before\", self.ret)\n        self.ret = self.ret * self.gamma + rews\n        # print(\"after\", self.ret)\n        obs = self._obfilt(obs)\n        if self.ret_rms:\n            self.ret_rms.update(np.array([self.ret].copy()))\n            rews = np.clip(rews / np.sqrt(self.ret_rms.var + self.epsilon), -self.cliprew, self.cliprew)\n        self.ret = self.ret * (1-float(news))\n        return obs, rews, news, infos\n\n    def _obfilt(self, obs):\n        if self.ob_rms:\n            self.ob_rms.update(obs)\n            obs = np.clip((obs - self.ob_rms.mean) / np.sqrt(self.ob_rms.var + self.epsilon), -self.clipob, self.clipob)\n            return obs\n        else:\n            return obs\n\n    def reset(self):\n        self.ret = np.zeros(())\n        obs = self.env.reset()\n        return self._obfilt(obs)\n\nif __name__ == \"__main__\":\n    parser = argparse.ArgumentParser(description='PPO agent')\n    # Common arguments\n    parser.add_argument('--exp-name', type=str, default=os.path.basename(__file__).rstrip(\".py\"),\n                       help='the name of this experiment')\n    parser.add_argument('--gym-id', type=str, default=\"MicrortsMining4x4F9-v0\",\n                       help='the id of the gym environment')\n    parser.add_argument('--seed', type=int, default=1,\n                       help='seed of the experiment')\n    parser.add_argument('--episode-length', type=int, default=0,\n                       help='the maximum length of each episode')\n    parser.add_argument('--total-timesteps', type=int, default=100000,\n                       help='total timesteps of the experiments')\n    parser.add_argument('--no-torch-deterministic', action='store_false', dest=\"torch_deterministic\", default=True,\n                       help='if toggled, `torch.backends.cudnn.deterministic=False`')\n    parser.add_argument('--no-cuda', action='store_false', dest=\"cuda\", default=True,\n                       help='if toggled, cuda will not be enabled by default')\n    parser.add_argument('--prod-mode', action='store_true', default=False,\n                       help='run the script in production mode and use wandb to log outputs')\n    parser.add_argument('--capture-video', action='store_true', default=False,\n                       help='weather to capture videos of the agent performances (check out `videos` folder)')\n    parser.add_argument('--wandb-project-name', type=str, default=\"cleanRL\",\n                       help=\"the wandb's project name\")\n    parser.add_argument('--wandb-entity', type=str, default=None,\n                       help=\"the entity (team) of wandb's project\")\n\n    # Algorithm specific arguments\n    parser.add_argument('--batch-size', type=int, default=2048,\n                       help='the batch size of ppo')\n    parser.add_argument('--minibatch-size', type=int, default=256,\n                       help='the mini batch size of ppo')\n    parser.add_argument('--gamma', type=float, default=0.99,\n                       help='the discount factor gamma')\n    parser.add_argument('--gae-lambda', type=float, default=0.97,\n                       help='the lambda for the general advantage estimation')\n    parser.add_argument('--ent-coef', type=float, default=0.01,\n                       help=\"coefficient of the entropy\")\n    parser.add_argument('--max-grad-norm', type=float, default=0.5,\n                       help='the maximum norm for the gradient clipping')\n    parser.add_argument('--clip-coef', type=float, default=0.2,\n                       help=\"the surrogate clipping coefficient\")\n    parser.add_argument('--update-epochs', type=int, default=10,\n                        help=\"the K epochs to update the policy\")\n    parser.add_argument('--kle-stop', action='store_true', default=False,\n                        help='If toggled, the policy updates will be early stopped w.r.t target-kl')\n    parser.add_argument('--kle-rollback', action='store_true', default=False,\n                        help='If toggled, the policy updates will roll back to previous policy if KL exceeds target-kl')\n    parser.add_argument('--target-kl', type=float, default=0.015,\n                        help='the target-kl variable that is referred by --kl')\n    parser.add_argument('--gae', action='store_true', default=True,\n                        help='Use GAE for advantage computation')\n    parser.add_argument('--policy-lr', type=float, default=3e-4,\n                        help=\"the learning rate of the policy optimizer\")\n    parser.add_argument('--value-lr', type=float, default=3e-4,\n                        help=\"the learning rate of the critic optimizer\")\n    parser.add_argument('--norm-obs', action='store_true', default=True,\n                        help=\"Toggles observation normalization\")\n    parser.add_argument('--norm-returns', action='store_true', default=False,\n                        help=\"Toggles returns normalization\")\n    parser.add_argument('--norm-adv', action='store_true', default=True,\n                        help=\"Toggles advantages normalization\")\n    parser.add_argument('--obs-clip', type=float, default=10.0,\n                        help=\"Value for reward clipping, as per the paper\")\n    parser.add_argument('--rew-clip', type=float, default=10.0,\n                        help=\"Value for observation clipping, as per the paper\")\n    parser.add_argument('--anneal-lr', action='store_true', default=True,\n                        help=\"Toggle learning rate annealing for policy and value networks\")\n    parser.add_argument('--weights-init', default=\"orthogonal\", choices=[\"xavier\", 'orthogonal'],\n                        help='Selects the scheme to be used for weights initialization'),\n    parser.add_argument('--clip-vloss', action=\"store_true\", default=True,\n                        help='Toggles wheter or not to use a clipped loss for the value function, as per the paper.')\n    parser.add_argument('--pol-layer-norm', action='store_true', default=False,\n                       help='Enables layer normalization in the policy network')\n    args = parser.parse_args()\n    if not args.seed:\n        args.seed = int(time.time())\n\nargs.features_turned_on = sum([args.kle_stop, args.kle_rollback, args.gae, args.norm_obs, args.norm_returns, args.norm_adv, args.anneal_lr, args.clip_vloss, args.pol_layer_norm])\n# TRY NOT TO MODIFY: setup the environment\nexperiment_name = f\"{args.gym_id}__{args.exp_name}__{args.seed}__{int(time.time())}\"\nwriter = SummaryWriter(f\"runs/{experiment_name}\")\nwriter.add_text('hyperparameters', \"|param|value|\\n|-|-|\\n%s\" % (\n        '\\n'.join([f\"|{key}|{value}|\" for key, value in vars(args).items()])))\n\nif args.prod_mode:\n    import wandb\n    wandb.init(project=args.wandb_project_name, entity=args.wandb_entity, tensorboard=True, config=vars(args), name=experiment_name, monitor_gym=True)\n    writer = SummaryWriter(f\"/tmp/{experiment_name}\")\n    wandb.save(os.path.abspath(__file__))\n\n# TRY NOT TO MODIFY: seeding\ndevice = torch.device('cuda' if torch.cuda.is_available() and args.cuda else 'cpu')\nenv = gym.make(args.gym_id)\n# respect the default timelimit\nassert isinstance(env.action_space, MultiDiscrete), \"only MultiDiscrete action space is supported\"\nassert isinstance(env, TimeLimit) or int(args.episode_length), \"the gym env does not have a built in TimeLimit, please specify by using --episode-length\"\nif isinstance(env, TimeLimit):\n    if int(args.episode_length):\n        env._max_episode_steps = int(args.episode_length)\n    args.episode_length = env._max_episode_steps\nelse:\n    env = TimeLimit(env, int(args.episode_length))\nenv = NormalizedEnv(env.env, ob=args.norm_obs, ret=args.norm_returns, clipob=args.obs_clip, cliprew=args.rew_clip, gamma=args.gamma)\nenv = TimeLimit(env, int(args.episode_length))\nrandom.seed(args.seed)\nnp.random.seed(args.seed)\ntorch.manual_seed(args.seed)\ntorch.backends.cudnn.deterministic = args.torch_deterministic\nenv.seed(args.seed)\nenv.action_space.seed(args.seed)\nenv.observation_space.seed(args.seed)\nif args.capture_video:\n    env = Monitor(env, f'videos/{experiment_name}')\n\n# ALGO LOGIC: initialize agent here:\nclass CategoricalMasked(Categorical):\n    def __init__(self, probs=None, logits=None, validate_args=None, masks=[]):\n        self.masks = masks\n        if len(self.masks) == 0:\n            super(CategoricalMasked, self).__init__(probs, logits, validate_args)\n        else:\n            self.masks = masks.type(torch.BoolTensor).to(device)\n            logits = torch.where(self.masks, logits, torch.tensor(-1e+8).to(device))\n            super(CategoricalMasked, self).__init__(probs, logits, validate_args)\n    \n    def entropy(self):\n        if len(self.masks) == 0:\n            return super(CategoricalMasked, self).entropy()\n        p_log_p = self.logits * self.probs\n        p_log_p = torch.where(self.masks, p_log_p, torch.tensor(0.).to(device))\n        return -p_log_p.sum(-1)\n\nclass Policy(nn.Module):\n    def __init__(self):\n        super(Policy, self).__init__()\n        self.features = nn.Sequential(\n            nn.Conv2d(27, 16, kernel_size=2,),\n            nn.MaxPool2d(1),\n            nn.ReLU())\n        self.fc = nn.Sequential(\n            nn.Linear(16*3*3, 128),\n            nn.ReLU(),\n            nn.Linear(128, env.action_space.nvec.sum())\n        )\n\n    def forward(self, x):\n        x = torch.Tensor(np.moveaxis(x, -1, 1)).to(device)\n        x = self.features(x)\n        x = x.reshape(x.size(0), -1)\n        x = self.fc(x)\n        return x\n\n    def get_action(self, x, action=None, invalid_action_masks=None):\n        logits = self.forward(x)\n        split_logits = torch.split(logits, env.action_space.nvec.tolist(), dim=1)\n        \n        if invalid_action_masks is not None:\n            split_invalid_action_masks = torch.split(invalid_action_masks, env.action_space.nvec.tolist(), dim=1)\n            multi_categoricals = [CategoricalMasked(logits=logits, masks=iam) for (logits, iam) in zip(split_logits, split_invalid_action_masks)]\n        else:\n            multi_categoricals = [Categorical(logits=logits) for logits in split_logits]\n        \n        if action is None:\n            action = torch.stack([categorical.sample() for categorical in multi_categoricals])\n        logprob = torch.stack([categorical.log_prob(a) for a, categorical in zip(action, multi_categoricals)])\n        return action, logprob, [], multi_categoricals\n\nclass Value(nn.Module):\n    def __init__(self):\n        super(Value, self).__init__()\n        self.features = nn.Sequential(\n            nn.Conv2d(27, 16, kernel_size=2,),\n            nn.MaxPool2d(1),\n            nn.ReLU())\n        self.fc = nn.Sequential(\n            nn.Linear(16*3*3, 128),\n            nn.ReLU(),\n            nn.Linear(128, 1)\n        )\n\n    def forward(self, x):\n        x = torch.Tensor(np.moveaxis(x, -1, 1)).to(device)\n        x = self.features(x)\n        x = x.reshape(x.size(0), -1)\n        x = self.fc(x)\n        return x\n\ndef discount_cumsum(x, dones, gamma):\n    \"\"\"\n    computing discounted cumulative sums of vectors that resets with dones\n    input:\n        vector x,  vector dones,\n        [x0,       [0,\n         x1,        0,\n         x2         1,\n         x3         0, \n         x4]        0]\n    output:\n        [x0 + discount * x1 + discount^2 * x2,\n         x1 + discount * x2,\n         x2,\n         x3 + discount * x4,\n         x4]\n    \"\"\"\n    discount_cumsum = np.zeros_like(x)\n    discount_cumsum[-1] = x[-1]\n    for t in reversed(range(x.shape[0]-1)):\n        discount_cumsum[t] = x[t] + gamma * discount_cumsum[t+1] * (1-dones[t])\n    return discount_cumsum\n\npg = Policy().to(device)\nvf = Value().to(device)\n\n# MODIFIED: Separate optimizer and learning rates\npg_optimizer = optim.Adam(list(pg.parameters()), lr=args.policy_lr)\nv_optimizer = optim.Adam(list(vf.parameters()), lr=args.value_lr)\n\n# MODIFIED: Initializing learning rate anneal scheduler when need\nif args.anneal_lr:\n    anneal_fn = lambda f: max(0, 1-f / args.total_timesteps)\n    pg_lr_scheduler = optim.lr_scheduler.LambdaLR(pg_optimizer, lr_lambda=anneal_fn)\n    vf_lr_scheduler = optim.lr_scheduler.LambdaLR(v_optimizer, lr_lambda=anneal_fn)\n\nloss_fn = nn.MSELoss()\n\n# TRY NOT TO MODIFY: start the game\nglobal_step = 0\nwhile global_step < args.total_timesteps:\n    if args.capture_video:\n        env.stats_recorder.done=True\n    next_obs = np.array(env.reset())\n\n    # ALGO Logic: Storage for epoch data\n    obs = np.empty((args.batch_size,) + env.observation_space.shape)\n\n    actions = np.empty((args.batch_size,) + env.action_space.shape)\n    logprobs = torch.zeros((env.action_space.nvec.shape[0], args.batch_size,)).to(device)\n\n    rewards = np.zeros((args.batch_size,))\n    raw_rewards = np.zeros((len(env.rfs),args.batch_size,))\n    \n    real_rewards = []\n    invalid_action_stats = []\n\n    dones = np.zeros((args.batch_size,))\n    values = torch.zeros((args.batch_size,)).to(device)\n\n    invalid_action_masks = torch.zeros((args.batch_size, env.action_space.nvec.sum()))\n    # TRY NOT TO MODIFY: prepare the execution of the game.\n    for step in range(args.batch_size):\n        env.render()\n        global_step += 1\n        obs[step] = next_obs.copy()\n\n        # ALGO LOGIC: put action logic here\n        invalid_action_mask = torch.ones(env.action_space.nvec.sum())\n        invalid_action_mask[0:env.action_space.nvec[0]] = torch.tensor(env.unit_location_mask)\n        invalid_action_mask[-env.action_space.nvec[-1]:] = torch.tensor(env.target_unit_location_mask)\n        invalid_action_masks[step] = invalid_action_mask\n        with torch.no_grad():\n            values[step] = vf.forward(obs[step:step+1])\n            action, logproba, _, probs = pg.get_action(obs[step:step+1], invalid_action_masks=invalid_action_masks[step:step+1])\n            \n            # CORE LOGIC:\n            # use the action generated by CategoricalMasked, but \n            # don't adjust the logprobability accordingly. Instead, calculate the log\n            # probability using Categorical\n            action, logproba, _, probs = pg.get_action(obs[step:step+1], action=action)\n        \n        actions[step] = action[:,0].data.cpu().numpy()\n        logprobs[:,[step]] = logproba\n\n        # TRY NOT TO MODIFY: execute the game and log data.\n        next_obs, rewards[step], dones[step], info = env.step(action[:,0].data.cpu().numpy())\n        raw_rewards[:,step] = info[\"rewards\"]\n        real_rewards += [info['real_reward']]\n        invalid_action_stats += [info['invalid_action_stats']]\n        next_obs = np.array(next_obs)\n\n        # Annealing the rate if instructed to do so.\n        if args.anneal_lr:\n            pg_lr_scheduler.step()\n            vf_lr_scheduler.step()\n\n        if dones[step]:\n            # Computing the discounted returns:\n            writer.add_scalar(\"charts/episode_reward\", np.sum(real_rewards), global_step)\n            print(f\"global_step={global_step}, episode_reward={np.sum(real_rewards)}\")\n            for i in range(len(env.rfs)):\n                writer.add_scalar(f\"charts/episode_reward/{str(env.rfs[i])}\", raw_rewards.sum(1)[i], global_step)\n            real_rewards = []\n            for key, idx in zip(info['invalid_action_stats'], range(len(info['invalid_action_stats']))):\n                writer.add_scalar(f\"stats/{key}\", pd.DataFrame(invalid_action_stats).sum(0)[idx], global_step)\n            invalid_action_stats = []\n            next_obs = np.array(env.reset())\n\n    # bootstrap reward if not done. reached the batch limit\n    last_value = 0\n    if not dones[step]:\n        last_value = vf.forward(next_obs.reshape((1,)+next_obs.shape))[0].detach().cpu().numpy()[0]\n    bootstrapped_rewards = np.append(rewards, last_value)\n\n    # calculate the returns and advantages\n    if args.gae:\n        bootstrapped_values = np.append(values.detach().cpu().numpy(), last_value)\n        deltas = bootstrapped_rewards[:-1] + args.gamma * bootstrapped_values[1:] * (1-dones) - bootstrapped_values[:-1]\n        advantages = discount_cumsum(deltas, dones, args.gamma * args.gae_lambda)\n        advantages = torch.Tensor(advantages).to(device)\n        returns = advantages + values\n    else:\n        returns = discount_cumsum(bootstrapped_rewards, dones, args.gamma)[:-1]\n        advantages = returns - values.detach().cpu().numpy()\n        advantages = torch.Tensor(advantages).to(device)\n        returns = torch.Tensor(returns).to(device)\n\n    # Advantage normalization\n    if args.norm_adv:\n        EPS = 1e-10\n        advantages = (advantages - advantages.mean()) / (advantages.std() + EPS)\n\n    # Optimizaing policy network\n    entropys = []\n    target_pg = Policy().to(device)\n    inds = np.arange(args.batch_size,)\n    for i_epoch_pi in range(args.update_epochs):\n        np.random.shuffle(inds)\n        for start in range(0, args.batch_size, args.minibatch_size):\n            end = start + args.minibatch_size\n            minibatch_ind = inds[start:end]\n            target_pg.load_state_dict(pg.state_dict())\n            \n            _, newlogproba, _, _ = pg.get_action(\n                obs[minibatch_ind],\n                torch.LongTensor(actions[minibatch_ind].astype(np.int)).to(device).T,)\n            ratio = (newlogproba - logprobs[:,minibatch_ind]).exp()\n\n            # Policy loss as in OpenAI SpinUp\n            clip_adv = torch.where(advantages[minibatch_ind] > 0,\n                                    (1.+args.clip_coef) * advantages[minibatch_ind],\n                                    (1.-args.clip_coef) * advantages[minibatch_ind]).to(device)\n\n            # Entropy computation with resampled actions\n            entropy = -(newlogproba.exp() * newlogproba).mean()\n            entropys.append(entropy.item())\n\n            policy_loss = -torch.min(ratio * advantages[minibatch_ind], clip_adv) + args.ent_coef * entropy\n            policy_loss = policy_loss.mean()\n            \n            pg_optimizer.zero_grad()\n            policy_loss.backward()\n            nn.utils.clip_grad_norm_(pg.parameters(), args.max_grad_norm)\n            pg_optimizer.step()\n\n            approx_kl = (logprobs[:,minibatch_ind] - newlogproba).mean()\n            # Optimizing value network\n            new_values = vf.forward(obs[minibatch_ind]).view(-1)\n\n            # Value loss clipping\n            if args.clip_vloss:\n                v_loss_unclipped = ((new_values - returns[minibatch_ind]) ** 2)\n                v_clipped = values[minibatch_ind] + torch.clamp(new_values - values[minibatch_ind], -args.clip_coef, args.clip_coef)\n                v_loss_clipped = (v_clipped - returns[minibatch_ind])**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 = torch.mean((returns[minibatch_ind]- new_values).pow(2))\n\n            v_optimizer.zero_grad()\n            v_loss.backward()\n            nn.utils.clip_grad_norm_(vf.parameters(), args.max_grad_norm)\n            v_optimizer.step()\n\n        if args.kle_stop:\n            if approx_kl > args.target_kl:\n                break\n        if args.kle_rollback:\n            if (logprobs[:,minibatch_ind] - \n                pg.get_action(\n                    obs[minibatch_ind],\n                    torch.LongTensor(actions[minibatch_ind].astype(np.int)).to(device).T,\n                    invalid_action_masks[minibatch_ind])[1]).mean() > args.target_kl:\n                pg.load_state_dict(target_pg.state_dict())\n                break\n\n    # TRY NOT TO MODIFY: record rewards for plotting purposes\n    writer.add_scalar(\"losses/value_loss\", v_loss.item(), global_step)\n    writer.add_scalar(\"charts/policy_learning_rate\", pg_optimizer.param_groups[0]['lr'], global_step)\n    writer.add_scalar(\"charts/value_learning_rate\", v_optimizer.param_groups[0]['lr'], global_step)\n    writer.add_scalar(\"losses/policy_loss\", policy_loss.item(), global_step)\n    writer.add_scalar(\"losses/entropy\", np.mean(entropys), global_step)\n    writer.add_scalar(\"losses/approx_kl\", approx_kl.item(), global_step)\n    if args.kle_stop or args.kle_rollback:\n        writer.add_scalar(\"debug/pg_stop_iter\", i_epoch_pi, global_step)\n\nenv.close()\nwriter.close()\n"
  },
  {
    "path": "invalid_action_masking/ppo_no_mask_10x10.py",
    "content": "import torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport torch.nn.functional as F\nfrom torch.distributions.categorical import Categorical\nfrom torch.utils.tensorboard import SummaryWriter\n\nfrom cleanrl.common import preprocess_obs_space, preprocess_ac_space\nimport argparse\nimport numpy as np\nimport gym\nimport gym_microrts\nfrom gym.wrappers import TimeLimit, Monitor\nfrom gym.spaces import Discrete, Box, MultiBinary, MultiDiscrete, Space\nimport time\nimport random\nimport os\nimport pandas as pd\n\n# taken from https://github.com/openai/baselines/blob/master/baselines/common/vec_env/vec_normalize.py\nclass RunningMeanStd(object):\n    def __init__(self, epsilon=1e-4, shape=()):\n        self.mean = np.zeros(shape, 'float64')\n        self.var = np.ones(shape, 'float64')\n        self.count = epsilon\n\n    def update(self, x):\n        batch_mean = np.mean([x], axis=0)\n        batch_var = np.var([x], axis=0)\n        batch_count = 1\n        self.update_from_moments(batch_mean, batch_var, batch_count)\n\n    def update_from_moments(self, batch_mean, batch_var, batch_count):\n        self.mean, self.var, self.count = update_mean_var_count_from_moments(\n            self.mean, self.var, self.count, batch_mean, batch_var, batch_count)\n\ndef update_mean_var_count_from_moments(mean, var, count, batch_mean, batch_var, batch_count):\n    delta = batch_mean - mean\n    tot_count = count + batch_count\n\n    new_mean = mean + delta * batch_count / tot_count\n    m_a = var * count\n    m_b = batch_var * batch_count\n    M2 = m_a + m_b + np.square(delta) * count * batch_count / tot_count\n    new_var = M2 / tot_count\n    new_count = tot_count\n\n    return new_mean, new_var, new_count\nclass NormalizedEnv(gym.core.Wrapper):\n    def __init__(self, env, ob=True, ret=True, clipob=10., cliprew=10., gamma=0.99, epsilon=1e-8):\n        super(NormalizedEnv, self).__init__(env)\n        self.ob_rms = RunningMeanStd(shape=self.observation_space.shape) if ob else None\n        self.ret_rms = RunningMeanStd(shape=(1,)) if ret else None\n        self.clipob = clipob\n        self.cliprew = cliprew\n        self.ret = np.zeros(())\n        self.gamma = gamma\n        self.epsilon = epsilon\n\n    def step(self, action):\n        obs, rews, news, infos = self.env.step(action)\n        infos['real_reward'] = rews\n        # print(\"before\", self.ret)\n        self.ret = self.ret * self.gamma + rews\n        # print(\"after\", self.ret)\n        obs = self._obfilt(obs)\n        if self.ret_rms:\n            self.ret_rms.update(np.array([self.ret].copy()))\n            rews = np.clip(rews / np.sqrt(self.ret_rms.var + self.epsilon), -self.cliprew, self.cliprew)\n        self.ret = self.ret * (1-float(news))\n        return obs, rews, news, infos\n\n    def _obfilt(self, obs):\n        if self.ob_rms:\n            self.ob_rms.update(obs)\n            obs = np.clip((obs - self.ob_rms.mean) / np.sqrt(self.ob_rms.var + self.epsilon), -self.clipob, self.clipob)\n            return obs\n        else:\n            return obs\n\n    def reset(self):\n        self.ret = np.zeros(())\n        obs = self.env.reset()\n        return self._obfilt(obs)\n\nif __name__ == \"__main__\":\n    parser = argparse.ArgumentParser(description='PPO agent')\n    # Common arguments\n    parser.add_argument('--exp-name', type=str, default=os.path.basename(__file__).rstrip(\".py\"),\n                       help='the name of this experiment')\n    parser.add_argument('--gym-id', type=str, default=\"MicrortsMining10x10F9-v0\",\n                       help='the id of the gym environment')\n    parser.add_argument('--seed', type=int, default=1,\n                       help='seed of the experiment')\n    parser.add_argument('--episode-length', type=int, default=0,\n                       help='the maximum length of each episode')\n    parser.add_argument('--total-timesteps', type=int, default=100000,\n                       help='total timesteps of the experiments')\n    parser.add_argument('--no-torch-deterministic', action='store_false', dest=\"torch_deterministic\", default=True,\n                       help='if toggled, `torch.backends.cudnn.deterministic=False`')\n    parser.add_argument('--no-cuda', action='store_false', dest=\"cuda\", default=True,\n                       help='if toggled, cuda will not be enabled by default')\n    parser.add_argument('--prod-mode', action='store_true', default=False,\n                       help='run the script in production mode and use wandb to log outputs')\n    parser.add_argument('--capture-video', action='store_true', default=False,\n                       help='weather to capture videos of the agent performances (check out `videos` folder)')\n    parser.add_argument('--wandb-project-name', type=str, default=\"cleanRL\",\n                       help=\"the wandb's project name\")\n    parser.add_argument('--wandb-entity', type=str, default=None,\n                       help=\"the entity (team) of wandb's project\")\n\n    # Algorithm specific arguments\n    parser.add_argument('--batch-size', type=int, default=2048,\n                       help='the batch size of ppo')\n    parser.add_argument('--minibatch-size', type=int, default=256,\n                       help='the mini batch size of ppo')\n    parser.add_argument('--gamma', type=float, default=0.99,\n                       help='the discount factor gamma')\n    parser.add_argument('--gae-lambda', type=float, default=0.97,\n                       help='the lambda for the general advantage estimation')\n    parser.add_argument('--ent-coef', type=float, default=0.01,\n                       help=\"coefficient of the entropy\")\n    parser.add_argument('--max-grad-norm', type=float, default=0.5,\n                       help='the maximum norm for the gradient clipping')\n    parser.add_argument('--clip-coef', type=float, default=0.2,\n                       help=\"the surrogate clipping coefficient\")\n    parser.add_argument('--update-epochs', type=int, default=10,\n                        help=\"the K epochs to update the policy\")\n    parser.add_argument('--kle-stop', action='store_true', default=False,\n                        help='If toggled, the policy updates will be early stopped w.r.t target-kl')\n    parser.add_argument('--kle-rollback', action='store_true', default=False,\n                        help='If toggled, the policy updates will roll back to previous policy if KL exceeds target-kl')\n    parser.add_argument('--target-kl', type=float, default=0.015,\n                        help='the target-kl variable that is referred by --kl')\n    parser.add_argument('--gae', action='store_true', default=True,\n                        help='Use GAE for advantage computation')\n    parser.add_argument('--policy-lr', type=float, default=3e-4,\n                        help=\"the learning rate of the policy optimizer\")\n    parser.add_argument('--value-lr', type=float, default=3e-4,\n                        help=\"the learning rate of the critic optimizer\")\n    parser.add_argument('--norm-obs', action='store_true', default=True,\n                        help=\"Toggles observation normalization\")\n    parser.add_argument('--norm-returns', action='store_true', default=False,\n                        help=\"Toggles returns normalization\")\n    parser.add_argument('--norm-adv', action='store_true', default=True,\n                        help=\"Toggles advantages normalization\")\n    parser.add_argument('--obs-clip', type=float, default=10.0,\n                        help=\"Value for reward clipping, as per the paper\")\n    parser.add_argument('--rew-clip', type=float, default=10.0,\n                        help=\"Value for observation clipping, as per the paper\")\n    parser.add_argument('--anneal-lr', action='store_true', default=True,\n                        help=\"Toggle learning rate annealing for policy and value networks\")\n    parser.add_argument('--weights-init', default=\"orthogonal\", choices=[\"xavier\", 'orthogonal'],\n                        help='Selects the scheme to be used for weights initialization'),\n    parser.add_argument('--clip-vloss', action=\"store_true\", default=True,\n                        help='Toggles wheter or not to use a clipped loss for the value function, as per the paper.')\n    parser.add_argument('--pol-layer-norm', action='store_true', default=False,\n                       help='Enables layer normalization in the policy network')\n    parser.add_argument('--invalid-action-penalty', type=float, default=0.00,\n                       help='the negative reward penalty for invalid actions')\n    args = parser.parse_args()\n    if not args.seed:\n        args.seed = int(time.time())\n\nargs.features_turned_on = sum([args.kle_stop, args.kle_rollback, args.gae, args.norm_obs, args.norm_returns, args.norm_adv, args.anneal_lr, args.clip_vloss, args.pol_layer_norm])\n# TRY NOT TO MODIFY: setup the environment\nexperiment_name = f\"{args.gym_id}__{args.exp_name}__{args.seed}__{int(time.time())}\"\nwriter = SummaryWriter(f\"runs/{experiment_name}\")\nwriter.add_text('hyperparameters', \"|param|value|\\n|-|-|\\n%s\" % (\n        '\\n'.join([f\"|{key}|{value}|\" for key, value in vars(args).items()])))\n\nif args.prod_mode:\n    import wandb\n    wandb.init(project=args.wandb_project_name, entity=args.wandb_entity, tensorboard=True, config=vars(args), name=experiment_name, monitor_gym=True)\n    writer = SummaryWriter(f\"/tmp/{experiment_name}\")\n    wandb.save(os.path.abspath(__file__))\n\n# TRY NOT TO MODIFY: seeding\ndevice = torch.device('cuda' if torch.cuda.is_available() and args.cuda else 'cpu')\nenv = gym.make(args.gym_id)\n# respect the default timelimit\nassert isinstance(env.action_space, MultiDiscrete), \"only MultiDiscrete action space is supported\"\nassert isinstance(env, TimeLimit) or int(args.episode_length), \"the gym env does not have a built in TimeLimit, please specify by using --episode-length\"\nif isinstance(env, TimeLimit):\n    if int(args.episode_length):\n        env._max_episode_steps = int(args.episode_length)\n    args.episode_length = env._max_episode_steps\nelse:\n    env = TimeLimit(env, int(args.episode_length))\nenv = NormalizedEnv(env.env, ob=args.norm_obs, ret=args.norm_returns, clipob=args.obs_clip, cliprew=args.rew_clip, gamma=args.gamma)\nenv = TimeLimit(env, int(args.episode_length))\nrandom.seed(args.seed)\nnp.random.seed(args.seed)\ntorch.manual_seed(args.seed)\ntorch.backends.cudnn.deterministic = args.torch_deterministic\nenv.seed(args.seed)\nenv.action_space.seed(args.seed)\nenv.observation_space.seed(args.seed)\nif args.capture_video:\n    env = Monitor(env, f'videos/{experiment_name}')\n\n# ALGO LOGIC: initialize agent here:\nclass CategoricalMasked(Categorical):\n    def __init__(self, probs=None, logits=None, validate_args=None, masks=[]):\n        self.masks = masks\n        if len(self.masks) == 0:\n            super(CategoricalMasked, self).__init__(probs, logits, validate_args)\n        else:\n            self.masks = masks.type(torch.BoolTensor).to(device)\n            logits = torch.where(self.masks, logits, torch.tensor(-1e+8).to(device))\n            super(CategoricalMasked, self).__init__(probs, logits, validate_args)\n    \n    def entropy(self):\n        if len(self.masks) == 0:\n            return super(CategoricalMasked, self).entropy()\n        p_log_p = self.logits * self.probs\n        p_log_p = torch.where(self.masks, p_log_p, torch.tensor(0.).to(device))\n        return -p_log_p.sum(-1)\n\nclass Policy(nn.Module):\n    def __init__(self):\n        super(Policy, self).__init__()\n        self.features = nn.Sequential(\n            nn.Conv2d(27, 16, kernel_size=3,),\n            nn.MaxPool2d(1),\n            nn.ReLU(),\n            nn.Conv2d(16, 32, kernel_size=3),\n            nn.MaxPool2d(1),\n            nn.ReLU())\n        self.fc = nn.Sequential(\n            nn.Linear(32*6*6, 128),\n            nn.ReLU(),\n            nn.Linear(128, env.action_space.nvec.sum())\n        )\n\n    def forward(self, x):\n        x = torch.Tensor(np.moveaxis(x, -1, 1)).to(device)\n        x = self.features(x)\n        x = x.reshape(x.size(0), -1)\n        x = self.fc(x)\n        return x\n\n    def get_action(self, x, action=None):\n        logits = self.forward(x)\n        split_logits = torch.split(logits, env.action_space.nvec.tolist(), dim=1)\n        multi_categoricals = [Categorical(logits=logits) for logits in split_logits]\n        \n        if action is None:\n            action = torch.stack([categorical.sample() for categorical in multi_categoricals])\n        logprob = torch.stack([categorical.log_prob(a) for a, categorical in zip(action, multi_categoricals)])\n        # entropy = torch.stack([categorical.entropy() for categorical in multi_categoricals])\n\n        return action, logprob, [], multi_categoricals\n\nclass Value(nn.Module):\n    def __init__(self):\n        super(Value, self).__init__()\n        self.features = nn.Sequential(\n            nn.Conv2d(27, 16, kernel_size=3,),\n            nn.MaxPool2d(1),\n            nn.ReLU(),\n            nn.Conv2d(16, 32, kernel_size=3),\n            nn.MaxPool2d(1),\n            nn.ReLU())\n        self.fc = nn.Sequential(\n            nn.Linear(32*6*6, 128),\n            nn.ReLU(),\n            nn.Linear(128, 1)\n        )\n\n    def forward(self, x):\n        x = torch.Tensor(np.moveaxis(x, -1, 1)).to(device)\n        x = self.features(x)\n        x = x.reshape(x.size(0), -1)\n        x = self.fc(x)\n        return x\n\ndef discount_cumsum(x, dones, gamma):\n    \"\"\"\n    computing discounted cumulative sums of vectors that resets with dones\n    input:\n        vector x,  vector dones,\n        [x0,       [0,\n         x1,        0,\n         x2         1,\n         x3         0, \n         x4]        0]\n    output:\n        [x0 + discount * x1 + discount^2 * x2,\n         x1 + discount * x2,\n         x2,\n         x3 + discount * x4,\n         x4]\n    \"\"\"\n    discount_cumsum = np.zeros_like(x)\n    discount_cumsum[-1] = x[-1]\n    for t in reversed(range(x.shape[0]-1)):\n        discount_cumsum[t] = x[t] + gamma * discount_cumsum[t+1] * (1-dones[t])\n    return discount_cumsum\n\npg = Policy().to(device)\nvf = Value().to(device)\n\n# MODIFIED: Separate optimizer and learning rates\npg_optimizer = optim.Adam(list(pg.parameters()), lr=args.policy_lr)\nv_optimizer = optim.Adam(list(vf.parameters()), lr=args.value_lr)\n\n# MODIFIED: Initializing learning rate anneal scheduler when need\nif args.anneal_lr:\n    anneal_fn = lambda f: max(0, 1-f / args.total_timesteps)\n    pg_lr_scheduler = optim.lr_scheduler.LambdaLR(pg_optimizer, lr_lambda=anneal_fn)\n    vf_lr_scheduler = optim.lr_scheduler.LambdaLR(v_optimizer, lr_lambda=anneal_fn)\n\nloss_fn = nn.MSELoss()\n\n# TRY NOT TO MODIFY: start the game\nglobal_step = 0\nwhile global_step < args.total_timesteps:\n    if args.capture_video:\n        env.stats_recorder.done=True\n    next_obs = np.array(env.reset())\n\n    # ALGO Logic: Storage for epoch data\n    obs = np.empty((args.batch_size,) + env.observation_space.shape)\n\n    actions = np.empty((args.batch_size,) + env.action_space.shape)\n    logprobs = torch.zeros((env.action_space.nvec.shape[0], args.batch_size,)).to(device)\n\n    rewards = np.zeros((args.batch_size,))\n    real_rewards = []\n    invalid_action_stats = []\n\n    dones = np.zeros((args.batch_size,))\n    values = torch.zeros((args.batch_size,)).to(device)\n\n    # TRY NOT TO MODIFY: prepare the execution of the game.\n    for step in range(args.batch_size):\n        env.render()\n        global_step += 1\n        obs[step] = next_obs.copy()\n\n        # ALGO LOGIC: put action logic here\n        with torch.no_grad():\n            values[step] = vf.forward(obs[step:step+1])\n            action, logproba, _, probs = pg.get_action(obs[step:step+1])\n        \n        actions[step] = action[:,0].data.cpu().numpy()\n        logprobs[:,[step]] = logproba\n\n        # TRY NOT TO MODIFY: execute the game and log data.\n        try:\n            next_obs, rewards[step], dones[step], info = env.step(action[:,0].data.cpu().numpy())\n            rewards[step] += pd.DataFrame([info['invalid_action_stats']]).sum().sum() * args.invalid_action_penalty\n        except Exception as e:\n            print(e)\n            print(e.stacktrace())\n        real_rewards += [info['real_reward']]\n        invalid_action_stats += [info['invalid_action_stats']]\n        next_obs = np.array(next_obs)\n\n        # Annealing the rate if instructed to do so.\n        if args.anneal_lr:\n            pg_lr_scheduler.step()\n            vf_lr_scheduler.step()\n\n        if dones[step]:\n            # Computing the discounted returns:\n            writer.add_scalar(\"charts/episode_reward\", np.sum(real_rewards), global_step)\n            print(f\"global_step={global_step}, episode_reward={np.sum(real_rewards)}\")\n            real_rewards = []\n            for key, idx in zip(info['invalid_action_stats'], range(len(info['invalid_action_stats']))):\n                writer.add_scalar(f\"stats/{key}\", pd.DataFrame(invalid_action_stats).sum(0)[idx], global_step)\n            invalid_action_stats = []\n            next_obs = np.array(env.reset())\n\n    # bootstrap reward if not done. reached the batch limit\n    last_value = 0\n    if not dones[step]:\n        last_value = vf.forward(next_obs.reshape((1,)+next_obs.shape))[0].detach().cpu().numpy()[0]\n    bootstrapped_rewards = np.append(rewards, last_value)\n\n    # calculate the returns and advantages\n    if args.gae:\n        bootstrapped_values = np.append(values.detach().cpu().numpy(), last_value)\n        deltas = bootstrapped_rewards[:-1] + args.gamma * bootstrapped_values[1:] * (1-dones) - bootstrapped_values[:-1]\n        advantages = discount_cumsum(deltas, dones, args.gamma * args.gae_lambda)\n        advantages = torch.Tensor(advantages).to(device)\n        returns = advantages + values\n    else:\n        returns = discount_cumsum(bootstrapped_rewards, dones, args.gamma)[:-1]\n        advantages = returns - values.detach().cpu().numpy()\n        advantages = torch.Tensor(advantages).to(device)\n        returns = torch.Tensor(returns).to(device)\n\n    # Advantage normalization\n    if args.norm_adv:\n        EPS = 1e-10\n        advantages = (advantages - advantages.mean()) / (advantages.std() + EPS)\n\n    # Optimizaing policy network\n    entropys = []\n    target_pg = Policy().to(device)\n    inds = np.arange(args.batch_size,)\n    for i_epoch_pi in range(args.update_epochs):\n        np.random.shuffle(inds)\n        for start in range(0, args.batch_size, args.minibatch_size):\n            end = start + args.minibatch_size\n            minibatch_ind = inds[start:end]\n            target_pg.load_state_dict(pg.state_dict())\n            \n            _, newlogproba, _, _ = pg.get_action(\n                obs[minibatch_ind],\n                torch.LongTensor(actions[minibatch_ind].astype(np.int)).to(device).T)\n            ratio = (newlogproba - logprobs[:,minibatch_ind]).exp()\n\n            # Policy loss as in OpenAI SpinUp\n            clip_adv = torch.where(advantages[minibatch_ind] > 0,\n                                    (1.+args.clip_coef) * advantages[minibatch_ind],\n                                    (1.-args.clip_coef) * advantages[minibatch_ind]).to(device)\n\n            # Entropy computation with resampled actions\n            entropy = -(newlogproba.exp() * newlogproba).mean()\n            entropys.append(entropy.item())\n\n            policy_loss = -torch.min(ratio * advantages[minibatch_ind], clip_adv) + args.ent_coef * entropy\n            policy_loss = policy_loss.mean()\n            \n            pg_optimizer.zero_grad()\n            policy_loss.backward()\n            nn.utils.clip_grad_norm_(pg.parameters(), args.max_grad_norm)\n            pg_optimizer.step()\n\n            approx_kl = (logprobs[:,minibatch_ind] - newlogproba).mean()\n            # Resample values\n            new_values = vf.forward(obs[minibatch_ind]).view(-1)\n\n            # Value loss clipping\n            if args.clip_vloss:\n                v_loss_unclipped = ((new_values - returns[minibatch_ind]) ** 2)\n                v_clipped = values[minibatch_ind] + torch.clamp(new_values - values[minibatch_ind], -args.clip_coef, args.clip_coef)\n                v_loss_clipped = (v_clipped - returns[minibatch_ind])**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 = torch.mean((returns[minibatch_ind]- new_values).pow(2))\n\n            v_optimizer.zero_grad()\n            v_loss.backward()\n            nn.utils.clip_grad_norm_(vf.parameters(), args.max_grad_norm)\n            v_optimizer.step()\n\n        if args.kle_stop:\n            if approx_kl > args.target_kl:\n                break\n        if args.kle_rollback:\n            if (logprobs[:,minibatch_ind] - \n                pg.get_action(\n                    obs[minibatch_ind],\n                    torch.LongTensor(actions[minibatch_ind].astype(np.int)).to(device).T,\n                    )[1]).mean() > args.target_kl:\n                pg.load_state_dict(target_pg.state_dict())\n                break\n\n    # TRY NOT TO MODIFY: record rewards for plotting purposes\n    writer.add_scalar(\"losses/value_loss\", v_loss.item(), global_step)\n    writer.add_scalar(\"charts/policy_learning_rate\", pg_optimizer.param_groups[0]['lr'], global_step)\n    writer.add_scalar(\"charts/value_learning_rate\", v_optimizer.param_groups[0]['lr'], global_step)\n    writer.add_scalar(\"losses/policy_loss\", policy_loss.item(), global_step)\n    writer.add_scalar(\"losses/entropy\", np.mean(entropys), global_step)\n    writer.add_scalar(\"losses/approx_kl\", approx_kl.item(), global_step)\n    if args.kle_stop or args.kle_rollback:\n        writer.add_scalar(\"debug/pg_stop_iter\", i_epoch_pi, global_step)\n\nenv.close()\nwriter.close()\n"
  },
  {
    "path": "invalid_action_masking/ppo_no_mask_16x16.py",
    "content": "import torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport torch.nn.functional as F\nfrom torch.distributions.categorical import Categorical\nfrom torch.utils.tensorboard import SummaryWriter\n\nfrom cleanrl.common import preprocess_obs_space, preprocess_ac_space\nimport argparse\nimport numpy as np\nimport gym\nimport gym_microrts\nfrom gym.wrappers import TimeLimit, Monitor\nfrom gym.spaces import Discrete, Box, MultiBinary, MultiDiscrete, Space\nimport time\nimport random\nimport os\nimport pandas as pd\n\n# taken from https://github.com/openai/baselines/blob/master/baselines/common/vec_env/vec_normalize.py\nclass RunningMeanStd(object):\n    def __init__(self, epsilon=1e-4, shape=()):\n        self.mean = np.zeros(shape, 'float64')\n        self.var = np.ones(shape, 'float64')\n        self.count = epsilon\n\n    def update(self, x):\n        batch_mean = np.mean([x], axis=0)\n        batch_var = np.var([x], axis=0)\n        batch_count = 1\n        self.update_from_moments(batch_mean, batch_var, batch_count)\n\n    def update_from_moments(self, batch_mean, batch_var, batch_count):\n        self.mean, self.var, self.count = update_mean_var_count_from_moments(\n            self.mean, self.var, self.count, batch_mean, batch_var, batch_count)\n\ndef update_mean_var_count_from_moments(mean, var, count, batch_mean, batch_var, batch_count):\n    delta = batch_mean - mean\n    tot_count = count + batch_count\n\n    new_mean = mean + delta * batch_count / tot_count\n    m_a = var * count\n    m_b = batch_var * batch_count\n    M2 = m_a + m_b + np.square(delta) * count * batch_count / tot_count\n    new_var = M2 / tot_count\n    new_count = tot_count\n\n    return new_mean, new_var, new_count\nclass NormalizedEnv(gym.core.Wrapper):\n    def __init__(self, env, ob=True, ret=True, clipob=10., cliprew=10., gamma=0.99, epsilon=1e-8):\n        super(NormalizedEnv, self).__init__(env)\n        self.ob_rms = RunningMeanStd(shape=self.observation_space.shape) if ob else None\n        self.ret_rms = RunningMeanStd(shape=(1,)) if ret else None\n        self.clipob = clipob\n        self.cliprew = cliprew\n        self.ret = np.zeros(())\n        self.gamma = gamma\n        self.epsilon = epsilon\n\n    def step(self, action):\n        obs, rews, news, infos = self.env.step(action)\n        infos['real_reward'] = rews\n        # print(\"before\", self.ret)\n        self.ret = self.ret * self.gamma + rews\n        # print(\"after\", self.ret)\n        obs = self._obfilt(obs)\n        if self.ret_rms:\n            self.ret_rms.update(np.array([self.ret].copy()))\n            rews = np.clip(rews / np.sqrt(self.ret_rms.var + self.epsilon), -self.cliprew, self.cliprew)\n        self.ret = self.ret * (1-float(news))\n        return obs, rews, news, infos\n\n    def _obfilt(self, obs):\n        if self.ob_rms:\n            self.ob_rms.update(obs)\n            obs = np.clip((obs - self.ob_rms.mean) / np.sqrt(self.ob_rms.var + self.epsilon), -self.clipob, self.clipob)\n            return obs\n        else:\n            return obs\n\n    def reset(self):\n        self.ret = np.zeros(())\n        obs = self.env.reset()\n        return self._obfilt(obs)\n\nif __name__ == \"__main__\":\n    parser = argparse.ArgumentParser(description='PPO agent')\n    # Common arguments\n    parser.add_argument('--exp-name', type=str, default=os.path.basename(__file__).rstrip(\".py\"),\n                       help='the name of this experiment')\n    parser.add_argument('--gym-id', type=str, default=\"MicrortsMining16x16F9-v0\",\n                       help='the id of the gym environment')\n    parser.add_argument('--seed', type=int, default=1,\n                       help='seed of the experiment')\n    parser.add_argument('--episode-length', type=int, default=0,\n                       help='the maximum length of each episode')\n    parser.add_argument('--total-timesteps', type=int, default=100000,\n                       help='total timesteps of the experiments')\n    parser.add_argument('--no-torch-deterministic', action='store_false', dest=\"torch_deterministic\", default=True,\n                       help='if toggled, `torch.backends.cudnn.deterministic=False`')\n    parser.add_argument('--no-cuda', action='store_false', dest=\"cuda\", default=True,\n                       help='if toggled, cuda will not be enabled by default')\n    parser.add_argument('--prod-mode', action='store_true', default=False,\n                       help='run the script in production mode and use wandb to log outputs')\n    parser.add_argument('--capture-video', action='store_true', default=False,\n                       help='weather to capture videos of the agent performances (check out `videos` folder)')\n    parser.add_argument('--wandb-project-name', type=str, default=\"cleanRL\",\n                       help=\"the wandb's project name\")\n    parser.add_argument('--wandb-entity', type=str, default=None,\n                       help=\"the entity (team) of wandb's project\")\n\n    # Algorithm specific arguments\n    parser.add_argument('--batch-size', type=int, default=2048,\n                       help='the batch size of ppo')\n    parser.add_argument('--minibatch-size', type=int, default=256,\n                       help='the mini batch size of ppo')\n    parser.add_argument('--gamma', type=float, default=0.99,\n                       help='the discount factor gamma')\n    parser.add_argument('--gae-lambda', type=float, default=0.97,\n                       help='the lambda for the general advantage estimation')\n    parser.add_argument('--ent-coef', type=float, default=0.01,\n                       help=\"coefficient of the entropy\")\n    parser.add_argument('--max-grad-norm', type=float, default=0.5,\n                       help='the maximum norm for the gradient clipping')\n    parser.add_argument('--clip-coef', type=float, default=0.2,\n                       help=\"the surrogate clipping coefficient\")\n    parser.add_argument('--update-epochs', type=int, default=10,\n                        help=\"the K epochs to update the policy\")\n    parser.add_argument('--kle-stop', action='store_true', default=False,\n                        help='If toggled, the policy updates will be early stopped w.r.t target-kl')\n    parser.add_argument('--kle-rollback', action='store_true', default=False,\n                        help='If toggled, the policy updates will roll back to previous policy if KL exceeds target-kl')\n    parser.add_argument('--target-kl', type=float, default=0.015,\n                        help='the target-kl variable that is referred by --kl')\n    parser.add_argument('--gae', action='store_true', default=True,\n                        help='Use GAE for advantage computation')\n    parser.add_argument('--policy-lr', type=float, default=3e-4,\n                        help=\"the learning rate of the policy optimizer\")\n    parser.add_argument('--value-lr', type=float, default=3e-4,\n                        help=\"the learning rate of the critic optimizer\")\n    parser.add_argument('--norm-obs', action='store_true', default=True,\n                        help=\"Toggles observation normalization\")\n    parser.add_argument('--norm-returns', action='store_true', default=False,\n                        help=\"Toggles returns normalization\")\n    parser.add_argument('--norm-adv', action='store_true', default=True,\n                        help=\"Toggles advantages normalization\")\n    parser.add_argument('--obs-clip', type=float, default=10.0,\n                        help=\"Value for reward clipping, as per the paper\")\n    parser.add_argument('--rew-clip', type=float, default=10.0,\n                        help=\"Value for observation clipping, as per the paper\")\n    parser.add_argument('--anneal-lr', action='store_true', default=True,\n                        help=\"Toggle learning rate annealing for policy and value networks\")\n    parser.add_argument('--weights-init', default=\"orthogonal\", choices=[\"xavier\", 'orthogonal'],\n                        help='Selects the scheme to be used for weights initialization'),\n    parser.add_argument('--clip-vloss', action=\"store_true\", default=True,\n                        help='Toggles wheter or not to use a clipped loss for the value function, as per the paper.')\n    parser.add_argument('--pol-layer-norm', action='store_true', default=False,\n                       help='Enables layer normalization in the policy network')\n    parser.add_argument('--invalid-action-penalty', type=float, default=0.00,\n                       help='the negative reward penalty for invalid actions')\n    args = parser.parse_args()\n    if not args.seed:\n        args.seed = int(time.time())\n\nargs.features_turned_on = sum([args.kle_stop, args.kle_rollback, args.gae, args.norm_obs, args.norm_returns, args.norm_adv, args.anneal_lr, args.clip_vloss, args.pol_layer_norm])\n# TRY NOT TO MODIFY: setup the environment\nexperiment_name = f\"{args.gym_id}__{args.exp_name}__{args.seed}__{int(time.time())}\"\nwriter = SummaryWriter(f\"runs/{experiment_name}\")\nwriter.add_text('hyperparameters', \"|param|value|\\n|-|-|\\n%s\" % (\n        '\\n'.join([f\"|{key}|{value}|\" for key, value in vars(args).items()])))\n\nif args.prod_mode:\n    import wandb\n    wandb.init(project=args.wandb_project_name, entity=args.wandb_entity, tensorboard=True, config=vars(args), name=experiment_name, monitor_gym=True)\n    writer = SummaryWriter(f\"/tmp/{experiment_name}\")\n    wandb.save(os.path.abspath(__file__))\n\n# TRY NOT TO MODIFY: seeding\ndevice = torch.device('cuda' if torch.cuda.is_available() and args.cuda else 'cpu')\nenv = gym.make(args.gym_id)\n# respect the default timelimit\nassert isinstance(env.action_space, MultiDiscrete), \"only MultiDiscrete action space is supported\"\nassert isinstance(env, TimeLimit) or int(args.episode_length), \"the gym env does not have a built in TimeLimit, please specify by using --episode-length\"\nif isinstance(env, TimeLimit):\n    if int(args.episode_length):\n        env._max_episode_steps = int(args.episode_length)\n    args.episode_length = env._max_episode_steps\nelse:\n    env = TimeLimit(env, int(args.episode_length))\nenv = NormalizedEnv(env.env, ob=args.norm_obs, ret=args.norm_returns, clipob=args.obs_clip, cliprew=args.rew_clip, gamma=args.gamma)\nenv = TimeLimit(env, int(args.episode_length))\nrandom.seed(args.seed)\nnp.random.seed(args.seed)\ntorch.manual_seed(args.seed)\ntorch.backends.cudnn.deterministic = args.torch_deterministic\nenv.seed(args.seed)\nenv.action_space.seed(args.seed)\nenv.observation_space.seed(args.seed)\nif args.capture_video:\n    env = Monitor(env, f'videos/{experiment_name}')\n\n# ALGO LOGIC: initialize agent here:\nclass CategoricalMasked(Categorical):\n    def __init__(self, probs=None, logits=None, validate_args=None, masks=[]):\n        self.masks = masks\n        if len(self.masks) == 0:\n            super(CategoricalMasked, self).__init__(probs, logits, validate_args)\n        else:\n            self.masks = masks.type(torch.BoolTensor).to(device)\n            logits = torch.where(self.masks, logits, torch.tensor(-1e+8).to(device))\n            super(CategoricalMasked, self).__init__(probs, logits, validate_args)\n    \n    def entropy(self):\n        if len(self.masks) == 0:\n            return super(CategoricalMasked, self).entropy()\n        p_log_p = self.logits * self.probs\n        p_log_p = torch.where(self.masks, p_log_p, torch.tensor(0.).to(device))\n        return -p_log_p.sum(-1)\n\nclass Policy(nn.Module):\n    def __init__(self):\n        super(Policy, self).__init__()\n        self.features = nn.Sequential(\n            nn.Conv2d(27, 16, kernel_size=3),\n            nn.MaxPool2d(2),\n            nn.ReLU(),\n            nn.Conv2d(16, 32, kernel_size=3),\n            nn.MaxPool2d(1),\n            nn.ReLU())\n        self.fc = nn.Sequential(\n            nn.Linear(32*5*5, 128),\n            nn.ReLU(),\n            nn.Linear(128, env.action_space.nvec.sum())\n        )\n\n    def forward(self, x):\n        x = torch.Tensor(np.moveaxis(x, -1, 1)).to(device)\n        x = self.features(x)\n        x = x.reshape(x.size(0), -1)\n        x = self.fc(x)\n        return x\n\n    def get_action(self, x, action=None):\n        logits = self.forward(x)\n        split_logits = torch.split(logits, env.action_space.nvec.tolist(), dim=1)\n        multi_categoricals = [Categorical(logits=logits) for logits in split_logits]\n        \n        if action is None:\n            action = torch.stack([categorical.sample() for categorical in multi_categoricals])\n        logprob = torch.stack([categorical.log_prob(a) for a, categorical in zip(action, multi_categoricals)])\n        # entropy = torch.stack([categorical.entropy() for categorical in multi_categoricals])\n\n        return action, logprob, [], multi_categoricals\n\nclass Value(nn.Module):\n    def __init__(self):\n        super(Value, self).__init__()\n        self.features = nn.Sequential(\n            nn.Conv2d(27, 16, kernel_size=3),\n            nn.MaxPool2d(2),\n            nn.ReLU(),\n            nn.Conv2d(16, 32, kernel_size=3),\n            nn.MaxPool2d(1),\n            nn.ReLU())\n        self.fc = nn.Sequential(\n            nn.Linear(32*5*5, 128),\n            nn.ReLU(),\n            nn.Linear(128, 1)\n        )\n\n    def forward(self, x):\n        x = torch.Tensor(np.moveaxis(x, -1, 1)).to(device)\n        x = self.features(x)\n        x = x.reshape(x.size(0), -1)\n        x = self.fc(x)\n        return x\n\ndef discount_cumsum(x, dones, gamma):\n    \"\"\"\n    computing discounted cumulative sums of vectors that resets with dones\n    input:\n        vector x,  vector dones,\n        [x0,       [0,\n         x1,        0,\n         x2         1,\n         x3         0, \n         x4]        0]\n    output:\n        [x0 + discount * x1 + discount^2 * x2,\n         x1 + discount * x2,\n         x2,\n         x3 + discount * x4,\n         x4]\n    \"\"\"\n    discount_cumsum = np.zeros_like(x)\n    discount_cumsum[-1] = x[-1]\n    for t in reversed(range(x.shape[0]-1)):\n        discount_cumsum[t] = x[t] + gamma * discount_cumsum[t+1] * (1-dones[t])\n    return discount_cumsum\n\npg = Policy().to(device)\nvf = Value().to(device)\n\n# MODIFIED: Separate optimizer and learning rates\npg_optimizer = optim.Adam(list(pg.parameters()), lr=args.policy_lr)\nv_optimizer = optim.Adam(list(vf.parameters()), lr=args.value_lr)\n\n# MODIFIED: Initializing learning rate anneal scheduler when need\nif args.anneal_lr:\n    anneal_fn = lambda f: max(0, 1-f / args.total_timesteps)\n    pg_lr_scheduler = optim.lr_scheduler.LambdaLR(pg_optimizer, lr_lambda=anneal_fn)\n    vf_lr_scheduler = optim.lr_scheduler.LambdaLR(v_optimizer, lr_lambda=anneal_fn)\n\nloss_fn = nn.MSELoss()\n\n# TRY NOT TO MODIFY: start the game\nglobal_step = 0\nwhile global_step < args.total_timesteps:\n    if args.capture_video:\n        env.stats_recorder.done=True\n    next_obs = np.array(env.reset())\n\n    # ALGO Logic: Storage for epoch data\n    obs = np.empty((args.batch_size,) + env.observation_space.shape)\n\n    actions = np.empty((args.batch_size,) + env.action_space.shape)\n    logprobs = torch.zeros((env.action_space.nvec.shape[0], args.batch_size,)).to(device)\n\n    rewards = np.zeros((args.batch_size,))\n    real_rewards = []\n    invalid_action_stats = []\n\n    dones = np.zeros((args.batch_size,))\n    values = torch.zeros((args.batch_size,)).to(device)\n\n    # TRY NOT TO MODIFY: prepare the execution of the game.\n    for step in range(args.batch_size):\n        env.render()\n        global_step += 1\n        obs[step] = next_obs.copy()\n\n        # ALGO LOGIC: put action logic here\n        with torch.no_grad():\n            values[step] = vf.forward(obs[step:step+1])\n            action, logproba, _, probs = pg.get_action(obs[step:step+1])\n        \n        actions[step] = action[:,0].data.cpu().numpy()\n        logprobs[:,[step]] = logproba\n\n        # TRY NOT TO MODIFY: execute the game and log data.\n        try:\n            next_obs, rewards[step], dones[step], info = env.step(action[:,0].data.cpu().numpy())\n            rewards[step] += pd.DataFrame([info['invalid_action_stats']]).sum().sum() * args.invalid_action_penalty\n        except Exception as e:\n            print(e)\n            print(e.stacktrace())\n        real_rewards += [info['real_reward']]\n        invalid_action_stats += [info['invalid_action_stats']]\n        next_obs = np.array(next_obs)\n\n        # Annealing the rate if instructed to do so.\n        if args.anneal_lr:\n            pg_lr_scheduler.step()\n            vf_lr_scheduler.step()\n\n        if dones[step]:\n            # Computing the discounted returns:\n            writer.add_scalar(\"charts/episode_reward\", np.sum(real_rewards), global_step)\n            print(f\"global_step={global_step}, episode_reward={np.sum(real_rewards)}\")\n            real_rewards = []\n            for key, idx in zip(info['invalid_action_stats'], range(len(info['invalid_action_stats']))):\n                writer.add_scalar(f\"stats/{key}\", pd.DataFrame(invalid_action_stats).sum(0)[idx], global_step)\n            invalid_action_stats = []\n            next_obs = np.array(env.reset())\n\n    # bootstrap reward if not done. reached the batch limit\n    last_value = 0\n    if not dones[step]:\n        last_value = vf.forward(next_obs.reshape((1,)+next_obs.shape))[0].detach().cpu().numpy()[0]\n    bootstrapped_rewards = np.append(rewards, last_value)\n\n    # calculate the returns and advantages\n    if args.gae:\n        bootstrapped_values = np.append(values.detach().cpu().numpy(), last_value)\n        deltas = bootstrapped_rewards[:-1] + args.gamma * bootstrapped_values[1:] * (1-dones) - bootstrapped_values[:-1]\n        advantages = discount_cumsum(deltas, dones, args.gamma * args.gae_lambda)\n        advantages = torch.Tensor(advantages).to(device)\n        returns = advantages + values\n    else:\n        returns = discount_cumsum(bootstrapped_rewards, dones, args.gamma)[:-1]\n        advantages = returns - values.detach().cpu().numpy()\n        advantages = torch.Tensor(advantages).to(device)\n        returns = torch.Tensor(returns).to(device)\n\n    # Advantage normalization\n    if args.norm_adv:\n        EPS = 1e-10\n        advantages = (advantages - advantages.mean()) / (advantages.std() + EPS)\n\n    # Optimizaing policy network\n    entropys = []\n    target_pg = Policy().to(device)\n    inds = np.arange(args.batch_size,)\n    for i_epoch_pi in range(args.update_epochs):\n        np.random.shuffle(inds)\n        for start in range(0, args.batch_size, args.minibatch_size):\n            end = start + args.minibatch_size\n            minibatch_ind = inds[start:end]\n            target_pg.load_state_dict(pg.state_dict())\n            \n            _, newlogproba, _, _ = pg.get_action(\n                obs[minibatch_ind],\n                torch.LongTensor(actions[minibatch_ind].astype(np.int)).to(device).T)\n            ratio = (newlogproba - logprobs[:,minibatch_ind]).exp()\n\n            # Policy loss as in OpenAI SpinUp\n            clip_adv = torch.where(advantages[minibatch_ind] > 0,\n                                    (1.+args.clip_coef) * advantages[minibatch_ind],\n                                    (1.-args.clip_coef) * advantages[minibatch_ind]).to(device)\n\n            # Entropy computation with resampled actions\n            entropy = -(newlogproba.exp() * newlogproba).mean()\n            entropys.append(entropy.item())\n\n            policy_loss = -torch.min(ratio * advantages[minibatch_ind], clip_adv) + args.ent_coef * entropy\n            policy_loss = policy_loss.mean()\n            \n            pg_optimizer.zero_grad()\n            policy_loss.backward()\n            nn.utils.clip_grad_norm_(pg.parameters(), args.max_grad_norm)\n            pg_optimizer.step()\n\n            approx_kl = (logprobs[:,minibatch_ind] - newlogproba).mean()\n            # Resample values\n            new_values = vf.forward(obs[minibatch_ind]).view(-1)\n\n            # Value loss clipping\n            if args.clip_vloss:\n                v_loss_unclipped = ((new_values - returns[minibatch_ind]) ** 2)\n                v_clipped = values[minibatch_ind] + torch.clamp(new_values - values[minibatch_ind], -args.clip_coef, args.clip_coef)\n                v_loss_clipped = (v_clipped - returns[minibatch_ind])**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 = torch.mean((returns[minibatch_ind]- new_values).pow(2))\n\n            v_optimizer.zero_grad()\n            v_loss.backward()\n            nn.utils.clip_grad_norm_(vf.parameters(), args.max_grad_norm)\n            v_optimizer.step()\n\n        if args.kle_stop:\n            if approx_kl > args.target_kl:\n                break\n        if args.kle_rollback:\n            if (logprobs[:,minibatch_ind] - \n                pg.get_action(\n                    obs[minibatch_ind],\n                    torch.LongTensor(actions[minibatch_ind].astype(np.int)).to(device).T,\n                    )[1]).mean() > args.target_kl:\n                pg.load_state_dict(target_pg.state_dict())\n                break\n\n    # TRY NOT TO MODIFY: record rewards for plotting purposes\n    writer.add_scalar(\"losses/value_loss\", v_loss.item(), global_step)\n    writer.add_scalar(\"charts/policy_learning_rate\", pg_optimizer.param_groups[0]['lr'], global_step)\n    writer.add_scalar(\"charts/value_learning_rate\", v_optimizer.param_groups[0]['lr'], global_step)\n    writer.add_scalar(\"losses/policy_loss\", policy_loss.item(), global_step)\n    writer.add_scalar(\"losses/entropy\", np.mean(entropys), global_step)\n    writer.add_scalar(\"losses/approx_kl\", approx_kl.item(), global_step)\n    if args.kle_stop or args.kle_rollback:\n        writer.add_scalar(\"debug/pg_stop_iter\", i_epoch_pi, global_step)\n\nenv.close()\nwriter.close()\n"
  },
  {
    "path": "invalid_action_masking/ppo_no_mask_24x24.py",
    "content": "import torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport torch.nn.functional as F\nfrom torch.distributions.categorical import Categorical\nfrom torch.utils.tensorboard import SummaryWriter\n\nfrom cleanrl.common import preprocess_obs_space, preprocess_ac_space\nimport argparse\nimport numpy as np\nimport gym\nimport gym_microrts\nfrom gym.wrappers import TimeLimit, Monitor\nfrom gym.spaces import Discrete, Box, MultiBinary, MultiDiscrete, Space\nimport time\nimport random\nimport os\nimport pandas as pd\n\n# taken from https://github.com/openai/baselines/blob/master/baselines/common/vec_env/vec_normalize.py\nclass RunningMeanStd(object):\n    def __init__(self, epsilon=1e-4, shape=()):\n        self.mean = np.zeros(shape, 'float64')\n        self.var = np.ones(shape, 'float64')\n        self.count = epsilon\n\n    def update(self, x):\n        batch_mean = np.mean([x], axis=0)\n        batch_var = np.var([x], axis=0)\n        batch_count = 1\n        self.update_from_moments(batch_mean, batch_var, batch_count)\n\n    def update_from_moments(self, batch_mean, batch_var, batch_count):\n        self.mean, self.var, self.count = update_mean_var_count_from_moments(\n            self.mean, self.var, self.count, batch_mean, batch_var, batch_count)\n\ndef update_mean_var_count_from_moments(mean, var, count, batch_mean, batch_var, batch_count):\n    delta = batch_mean - mean\n    tot_count = count + batch_count\n\n    new_mean = mean + delta * batch_count / tot_count\n    m_a = var * count\n    m_b = batch_var * batch_count\n    M2 = m_a + m_b + np.square(delta) * count * batch_count / tot_count\n    new_var = M2 / tot_count\n    new_count = tot_count\n\n    return new_mean, new_var, new_count\nclass NormalizedEnv(gym.core.Wrapper):\n    def __init__(self, env, ob=True, ret=True, clipob=10., cliprew=10., gamma=0.99, epsilon=1e-8):\n        super(NormalizedEnv, self).__init__(env)\n        self.ob_rms = RunningMeanStd(shape=self.observation_space.shape) if ob else None\n        self.ret_rms = RunningMeanStd(shape=(1,)) if ret else None\n        self.clipob = clipob\n        self.cliprew = cliprew\n        self.ret = np.zeros(())\n        self.gamma = gamma\n        self.epsilon = epsilon\n\n    def step(self, action):\n        obs, rews, news, infos = self.env.step(action)\n        infos['real_reward'] = rews\n        # print(\"before\", self.ret)\n        self.ret = self.ret * self.gamma + rews\n        # print(\"after\", self.ret)\n        obs = self._obfilt(obs)\n        if self.ret_rms:\n            self.ret_rms.update(np.array([self.ret].copy()))\n            rews = np.clip(rews / np.sqrt(self.ret_rms.var + self.epsilon), -self.cliprew, self.cliprew)\n        self.ret = self.ret * (1-float(news))\n        return obs, rews, news, infos\n\n    def _obfilt(self, obs):\n        if self.ob_rms:\n            self.ob_rms.update(obs)\n            obs = np.clip((obs - self.ob_rms.mean) / np.sqrt(self.ob_rms.var + self.epsilon), -self.clipob, self.clipob)\n            return obs\n        else:\n            return obs\n\n    def reset(self):\n        self.ret = np.zeros(())\n        obs = self.env.reset()\n        return self._obfilt(obs)\n\nif __name__ == \"__main__\":\n    parser = argparse.ArgumentParser(description='PPO agent')\n    # Common arguments\n    parser.add_argument('--exp-name', type=str, default=os.path.basename(__file__).rstrip(\".py\"),\n                       help='the name of this experiment')\n    parser.add_argument('--gym-id', type=str, default=\"MicrortsMining24x24F9-v0\",\n                       help='the id of the gym environment')\n    parser.add_argument('--seed', type=int, default=1,\n                       help='seed of the experiment')\n    parser.add_argument('--episode-length', type=int, default=0,\n                       help='the maximum length of each episode')\n    parser.add_argument('--total-timesteps', type=int, default=100000,\n                       help='total timesteps of the experiments')\n    parser.add_argument('--no-torch-deterministic', action='store_false', dest=\"torch_deterministic\", default=True,\n                       help='if toggled, `torch.backends.cudnn.deterministic=False`')\n    parser.add_argument('--no-cuda', action='store_false', dest=\"cuda\", default=True,\n                       help='if toggled, cuda will not be enabled by default')\n    parser.add_argument('--prod-mode', action='store_true', default=False,\n                       help='run the script in production mode and use wandb to log outputs')\n    parser.add_argument('--capture-video', action='store_true', default=False,\n                       help='weather to capture videos of the agent performances (check out `videos` folder)')\n    parser.add_argument('--wandb-project-name', type=str, default=\"cleanRL\",\n                       help=\"the wandb's project name\")\n    parser.add_argument('--wandb-entity', type=str, default=None,\n                       help=\"the entity (team) of wandb's project\")\n\n    # Algorithm specific arguments\n    parser.add_argument('--batch-size', type=int, default=2048,\n                       help='the batch size of ppo')\n    parser.add_argument('--minibatch-size', type=int, default=256,\n                       help='the mini batch size of ppo')\n    parser.add_argument('--gamma', type=float, default=0.99,\n                       help='the discount factor gamma')\n    parser.add_argument('--gae-lambda', type=float, default=0.97,\n                       help='the lambda for the general advantage estimation')\n    parser.add_argument('--ent-coef', type=float, default=0.01,\n                       help=\"coefficient of the entropy\")\n    parser.add_argument('--max-grad-norm', type=float, default=0.5,\n                       help='the maximum norm for the gradient clipping')\n    parser.add_argument('--clip-coef', type=float, default=0.2,\n                       help=\"the surrogate clipping coefficient\")\n    parser.add_argument('--update-epochs', type=int, default=10,\n                        help=\"the K epochs to update the policy\")\n    parser.add_argument('--kle-stop', action='store_true', default=False,\n                        help='If toggled, the policy updates will be early stopped w.r.t target-kl')\n    parser.add_argument('--kle-rollback', action='store_true', default=False,\n                        help='If toggled, the policy updates will roll back to previous policy if KL exceeds target-kl')\n    parser.add_argument('--target-kl', type=float, default=0.015,\n                        help='the target-kl variable that is referred by --kl')\n    parser.add_argument('--gae', action='store_true', default=True,\n                        help='Use GAE for advantage computation')\n    parser.add_argument('--policy-lr', type=float, default=3e-4,\n                        help=\"the learning rate of the policy optimizer\")\n    parser.add_argument('--value-lr', type=float, default=3e-4,\n                        help=\"the learning rate of the critic optimizer\")\n    parser.add_argument('--norm-obs', action='store_true', default=True,\n                        help=\"Toggles observation normalization\")\n    parser.add_argument('--norm-returns', action='store_true', default=False,\n                        help=\"Toggles returns normalization\")\n    parser.add_argument('--norm-adv', action='store_true', default=True,\n                        help=\"Toggles advantages normalization\")\n    parser.add_argument('--obs-clip', type=float, default=10.0,\n                        help=\"Value for reward clipping, as per the paper\")\n    parser.add_argument('--rew-clip', type=float, default=10.0,\n                        help=\"Value for observation clipping, as per the paper\")\n    parser.add_argument('--anneal-lr', action='store_true', default=True,\n                        help=\"Toggle learning rate annealing for policy and value networks\")\n    parser.add_argument('--weights-init', default=\"orthogonal\", choices=[\"xavier\", 'orthogonal'],\n                        help='Selects the scheme to be used for weights initialization'),\n    parser.add_argument('--clip-vloss', action=\"store_true\", default=True,\n                        help='Toggles wheter or not to use a clipped loss for the value function, as per the paper.')\n    parser.add_argument('--pol-layer-norm', action='store_true', default=False,\n                       help='Enables layer normalization in the policy network')\n    parser.add_argument('--invalid-action-penalty', type=float, default=0.00,\n                       help='the negative reward penalty for invalid actions')\n    args = parser.parse_args()\n    if not args.seed:\n        args.seed = int(time.time())\n\nargs.features_turned_on = sum([args.kle_stop, args.kle_rollback, args.gae, args.norm_obs, args.norm_returns, args.norm_adv, args.anneal_lr, args.clip_vloss, args.pol_layer_norm])\n# TRY NOT TO MODIFY: setup the environment\nexperiment_name = f\"{args.gym_id}__{args.exp_name}__{args.seed}__{int(time.time())}\"\nwriter = SummaryWriter(f\"runs/{experiment_name}\")\nwriter.add_text('hyperparameters', \"|param|value|\\n|-|-|\\n%s\" % (\n        '\\n'.join([f\"|{key}|{value}|\" for key, value in vars(args).items()])))\n\nif args.prod_mode:\n    import wandb\n    wandb.init(project=args.wandb_project_name, entity=args.wandb_entity, tensorboard=True, config=vars(args), name=experiment_name, monitor_gym=True)\n    writer = SummaryWriter(f\"/tmp/{experiment_name}\")\n    wandb.save(os.path.abspath(__file__))\n\n# TRY NOT TO MODIFY: seeding\ndevice = torch.device('cuda' if torch.cuda.is_available() and args.cuda else 'cpu')\nenv = gym.make(args.gym_id)\n# respect the default timelimit\nassert isinstance(env.action_space, MultiDiscrete), \"only MultiDiscrete action space is supported\"\nassert isinstance(env, TimeLimit) or int(args.episode_length), \"the gym env does not have a built in TimeLimit, please specify by using --episode-length\"\nif isinstance(env, TimeLimit):\n    if int(args.episode_length):\n        env._max_episode_steps = int(args.episode_length)\n    args.episode_length = env._max_episode_steps\nelse:\n    env = TimeLimit(env, int(args.episode_length))\nenv = NormalizedEnv(env.env, ob=args.norm_obs, ret=args.norm_returns, clipob=args.obs_clip, cliprew=args.rew_clip, gamma=args.gamma)\nenv = TimeLimit(env, int(args.episode_length))\nrandom.seed(args.seed)\nnp.random.seed(args.seed)\ntorch.manual_seed(args.seed)\ntorch.backends.cudnn.deterministic = args.torch_deterministic\nenv.seed(args.seed)\nenv.action_space.seed(args.seed)\nenv.observation_space.seed(args.seed)\nif args.capture_video:\n    env = Monitor(env, f'videos/{experiment_name}')\n\n# ALGO LOGIC: initialize agent here:\nclass CategoricalMasked(Categorical):\n    def __init__(self, probs=None, logits=None, validate_args=None, masks=[]):\n        self.masks = masks\n        if len(self.masks) == 0:\n            super(CategoricalMasked, self).__init__(probs, logits, validate_args)\n        else:\n            self.masks = masks.type(torch.BoolTensor).to(device)\n            logits = torch.where(self.masks, logits, torch.tensor(-1e+8).to(device))\n            super(CategoricalMasked, self).__init__(probs, logits, validate_args)\n    \n    def entropy(self):\n        if len(self.masks) == 0:\n            return super(CategoricalMasked, self).entropy()\n        p_log_p = self.logits * self.probs\n        p_log_p = torch.where(self.masks, p_log_p, torch.tensor(0.).to(device))\n        return -p_log_p.sum(-1)\n\nclass Policy(nn.Module):\n    def __init__(self):\n        super(Policy, self).__init__()\n        self.features = nn.Sequential(\n            nn.Conv2d(27, 16, kernel_size=3, stride=1),\n            nn.MaxPool2d(2),\n            nn.ReLU(),\n            nn.Conv2d(16, 32, kernel_size=2, stride=1),\n            nn.MaxPool2d(2),\n            nn.ReLU())\n        self.fc = nn.Sequential(\n            nn.Linear(32*5*5, 128),\n            nn.ReLU(),\n            nn.Linear(128, env.action_space.nvec.sum())\n        )\n\n    def forward(self, x):\n        x = torch.Tensor(np.moveaxis(x, -1, 1)).to(device)\n        x = self.features(x)\n        x = x.reshape(x.size(0), -1)\n        x = self.fc(x)\n        return x\n\n    def get_action(self, x, action=None):\n        logits = self.forward(x)\n        split_logits = torch.split(logits, env.action_space.nvec.tolist(), dim=1)\n        multi_categoricals = [Categorical(logits=logits) for logits in split_logits]\n        \n        if action is None:\n            action = torch.stack([categorical.sample() for categorical in multi_categoricals])\n        logprob = torch.stack([categorical.log_prob(a) for a, categorical in zip(action, multi_categoricals)])\n        # entropy = torch.stack([categorical.entropy() for categorical in multi_categoricals])\n\n        return action, logprob, [], multi_categoricals\n\nclass Value(nn.Module):\n    def __init__(self):\n        super(Value, self).__init__()\n        self.features = nn.Sequential(\n            nn.Conv2d(27, 16, kernel_size=3, stride=1),\n            nn.MaxPool2d(2),\n            nn.ReLU(),\n            nn.Conv2d(16, 32, kernel_size=2, stride=1),\n            nn.MaxPool2d(2),\n            nn.ReLU())\n        self.fc = nn.Sequential(\n            nn.Linear(32*5*5, 128),\n            nn.ReLU(),\n            nn.Linear(128, 1)\n        )\n\n    def forward(self, x):\n        x = torch.Tensor(np.moveaxis(x, -1, 1)).to(device)\n        x = self.features(x)\n        x = x.reshape(x.size(0), -1)\n        x = self.fc(x)\n        return x\n\ndef discount_cumsum(x, dones, gamma):\n    \"\"\"\n    computing discounted cumulative sums of vectors that resets with dones\n    input:\n        vector x,  vector dones,\n        [x0,       [0,\n         x1,        0,\n         x2         1,\n         x3         0, \n         x4]        0]\n    output:\n        [x0 + discount * x1 + discount^2 * x2,\n         x1 + discount * x2,\n         x2,\n         x3 + discount * x4,\n         x4]\n    \"\"\"\n    discount_cumsum = np.zeros_like(x)\n    discount_cumsum[-1] = x[-1]\n    for t in reversed(range(x.shape[0]-1)):\n        discount_cumsum[t] = x[t] + gamma * discount_cumsum[t+1] * (1-dones[t])\n    return discount_cumsum\n\npg = Policy().to(device)\nvf = Value().to(device)\n\n# MODIFIED: Separate optimizer and learning rates\npg_optimizer = optim.Adam(list(pg.parameters()), lr=args.policy_lr)\nv_optimizer = optim.Adam(list(vf.parameters()), lr=args.value_lr)\n\n# MODIFIED: Initializing learning rate anneal scheduler when need\nif args.anneal_lr:\n    anneal_fn = lambda f: max(0, 1-f / args.total_timesteps)\n    pg_lr_scheduler = optim.lr_scheduler.LambdaLR(pg_optimizer, lr_lambda=anneal_fn)\n    vf_lr_scheduler = optim.lr_scheduler.LambdaLR(v_optimizer, lr_lambda=anneal_fn)\n\nloss_fn = nn.MSELoss()\n\n# TRY NOT TO MODIFY: start the game\nglobal_step = 0\nwhile global_step < args.total_timesteps:\n    if args.capture_video:\n        env.stats_recorder.done=True\n    next_obs = np.array(env.reset())\n\n    # ALGO Logic: Storage for epoch data\n    obs = np.empty((args.batch_size,) + env.observation_space.shape)\n\n    actions = np.empty((args.batch_size,) + env.action_space.shape)\n    logprobs = torch.zeros((env.action_space.nvec.shape[0], args.batch_size,)).to(device)\n\n    rewards = np.zeros((args.batch_size,))\n    real_rewards = []\n    invalid_action_stats = []\n\n    dones = np.zeros((args.batch_size,))\n    values = torch.zeros((args.batch_size,)).to(device)\n\n    # TRY NOT TO MODIFY: prepare the execution of the game.\n    for step in range(args.batch_size):\n        env.render()\n        global_step += 1\n        obs[step] = next_obs.copy()\n\n        # ALGO LOGIC: put action logic here\n        with torch.no_grad():\n            values[step] = vf.forward(obs[step:step+1])\n            action, logproba, _, probs = pg.get_action(obs[step:step+1])\n        \n        actions[step] = action[:,0].data.cpu().numpy()\n        logprobs[:,[step]] = logproba\n\n        # TRY NOT TO MODIFY: execute the game and log data.\n        try:\n            next_obs, rewards[step], dones[step], info = env.step(action[:,0].data.cpu().numpy())\n            rewards[step] += pd.DataFrame([info['invalid_action_stats']]).sum().sum() * args.invalid_action_penalty\n        except Exception as e:\n            print(e)\n            print(e.stacktrace())\n        real_rewards += [info['real_reward']]\n        invalid_action_stats += [info['invalid_action_stats']]\n        next_obs = np.array(next_obs)\n\n        # Annealing the rate if instructed to do so.\n        if args.anneal_lr:\n            pg_lr_scheduler.step()\n            vf_lr_scheduler.step()\n\n        if dones[step]:\n            # Computing the discounted returns:\n            writer.add_scalar(\"charts/episode_reward\", np.sum(real_rewards), global_step)\n            print(f\"global_step={global_step}, episode_reward={np.sum(real_rewards)}\")\n            real_rewards = []\n            for key, idx in zip(info['invalid_action_stats'], range(len(info['invalid_action_stats']))):\n                writer.add_scalar(f\"stats/{key}\", pd.DataFrame(invalid_action_stats).sum(0)[idx], global_step)\n            invalid_action_stats = []\n            next_obs = np.array(env.reset())\n\n    # bootstrap reward if not done. reached the batch limit\n    last_value = 0\n    if not dones[step]:\n        last_value = vf.forward(next_obs.reshape((1,)+next_obs.shape))[0].detach().cpu().numpy()[0]\n    bootstrapped_rewards = np.append(rewards, last_value)\n\n    # calculate the returns and advantages\n    if args.gae:\n        bootstrapped_values = np.append(values.detach().cpu().numpy(), last_value)\n        deltas = bootstrapped_rewards[:-1] + args.gamma * bootstrapped_values[1:] * (1-dones) - bootstrapped_values[:-1]\n        advantages = discount_cumsum(deltas, dones, args.gamma * args.gae_lambda)\n        advantages = torch.Tensor(advantages).to(device)\n        returns = advantages + values\n    else:\n        returns = discount_cumsum(bootstrapped_rewards, dones, args.gamma)[:-1]\n        advantages = returns - values.detach().cpu().numpy()\n        advantages = torch.Tensor(advantages).to(device)\n        returns = torch.Tensor(returns).to(device)\n\n    # Advantage normalization\n    if args.norm_adv:\n        EPS = 1e-10\n        advantages = (advantages - advantages.mean()) / (advantages.std() + EPS)\n\n    # Optimizaing policy network\n    entropys = []\n    target_pg = Policy().to(device)\n    inds = np.arange(args.batch_size,)\n    for i_epoch_pi in range(args.update_epochs):\n        np.random.shuffle(inds)\n        for start in range(0, args.batch_size, args.minibatch_size):\n            end = start + args.minibatch_size\n            minibatch_ind = inds[start:end]\n            target_pg.load_state_dict(pg.state_dict())\n            \n            _, newlogproba, _, _ = pg.get_action(\n                obs[minibatch_ind],\n                torch.LongTensor(actions[minibatch_ind].astype(np.int)).to(device).T)\n            ratio = (newlogproba - logprobs[:,minibatch_ind]).exp()\n\n            # Policy loss as in OpenAI SpinUp\n            clip_adv = torch.where(advantages[minibatch_ind] > 0,\n                                    (1.+args.clip_coef) * advantages[minibatch_ind],\n                                    (1.-args.clip_coef) * advantages[minibatch_ind]).to(device)\n\n            # Entropy computation with resampled actions\n            entropy = -(newlogproba.exp() * newlogproba).mean()\n            entropys.append(entropy.item())\n\n            policy_loss = -torch.min(ratio * advantages[minibatch_ind], clip_adv) + args.ent_coef * entropy\n            policy_loss = policy_loss.mean()\n            \n            pg_optimizer.zero_grad()\n            policy_loss.backward()\n            nn.utils.clip_grad_norm_(pg.parameters(), args.max_grad_norm)\n            pg_optimizer.step()\n\n            approx_kl = (logprobs[:,minibatch_ind] - newlogproba).mean()\n            # Resample values\n            new_values = vf.forward(obs[minibatch_ind]).view(-1)\n\n            # Value loss clipping\n            if args.clip_vloss:\n                v_loss_unclipped = ((new_values - returns[minibatch_ind]) ** 2)\n                v_clipped = values[minibatch_ind] + torch.clamp(new_values - values[minibatch_ind], -args.clip_coef, args.clip_coef)\n                v_loss_clipped = (v_clipped - returns[minibatch_ind])**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 = torch.mean((returns[minibatch_ind]- new_values).pow(2))\n\n            v_optimizer.zero_grad()\n            v_loss.backward()\n            nn.utils.clip_grad_norm_(vf.parameters(), args.max_grad_norm)\n            v_optimizer.step()\n\n        if args.kle_stop:\n            if approx_kl > args.target_kl:\n                break\n        if args.kle_rollback:\n            if (logprobs[:,minibatch_ind] - \n                pg.get_action(\n                    obs[minibatch_ind],\n                    torch.LongTensor(actions[minibatch_ind].astype(np.int)).to(device).T,\n                    )[1]).mean() > args.target_kl:\n                pg.load_state_dict(target_pg.state_dict())\n                break\n\n    # TRY NOT TO MODIFY: record rewards for plotting purposes\n    writer.add_scalar(\"losses/value_loss\", v_loss.item(), global_step)\n    writer.add_scalar(\"charts/policy_learning_rate\", pg_optimizer.param_groups[0]['lr'], global_step)\n    writer.add_scalar(\"charts/value_learning_rate\", v_optimizer.param_groups[0]['lr'], global_step)\n    writer.add_scalar(\"losses/policy_loss\", policy_loss.item(), global_step)\n    writer.add_scalar(\"losses/entropy\", np.mean(entropys), global_step)\n    writer.add_scalar(\"losses/approx_kl\", approx_kl.item(), global_step)\n    if args.kle_stop or args.kle_rollback:\n        writer.add_scalar(\"debug/pg_stop_iter\", i_epoch_pi, global_step)\n\nenv.close()\nwriter.close()\n"
  },
  {
    "path": "invalid_action_masking/ppo_no_mask_4x4.py",
    "content": "import torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport torch.nn.functional as F\nfrom torch.distributions.categorical import Categorical\nfrom torch.utils.tensorboard import SummaryWriter\n\nfrom cleanrl.common import preprocess_obs_space, preprocess_ac_space\nimport argparse\nimport numpy as np\nimport gym\nimport gym_microrts\nfrom gym.wrappers import TimeLimit, Monitor\nfrom gym.spaces import Discrete, Box, MultiBinary, MultiDiscrete, Space\nimport time\nimport random\nimport os\nimport pandas as pd\n\n# taken from https://github.com/openai/baselines/blob/master/baselines/common/vec_env/vec_normalize.py\nclass RunningMeanStd(object):\n    def __init__(self, epsilon=1e-4, shape=()):\n        self.mean = np.zeros(shape, 'float64')\n        self.var = np.ones(shape, 'float64')\n        self.count = epsilon\n\n    def update(self, x):\n        batch_mean = np.mean([x], axis=0)\n        batch_var = np.var([x], axis=0)\n        batch_count = 1\n        self.update_from_moments(batch_mean, batch_var, batch_count)\n\n    def update_from_moments(self, batch_mean, batch_var, batch_count):\n        self.mean, self.var, self.count = update_mean_var_count_from_moments(\n            self.mean, self.var, self.count, batch_mean, batch_var, batch_count)\n\ndef update_mean_var_count_from_moments(mean, var, count, batch_mean, batch_var, batch_count):\n    delta = batch_mean - mean\n    tot_count = count + batch_count\n\n    new_mean = mean + delta * batch_count / tot_count\n    m_a = var * count\n    m_b = batch_var * batch_count\n    M2 = m_a + m_b + np.square(delta) * count * batch_count / tot_count\n    new_var = M2 / tot_count\n    new_count = tot_count\n\n    return new_mean, new_var, new_count\nclass NormalizedEnv(gym.core.Wrapper):\n    def __init__(self, env, ob=True, ret=True, clipob=10., cliprew=10., gamma=0.99, epsilon=1e-8):\n        super(NormalizedEnv, self).__init__(env)\n        self.ob_rms = RunningMeanStd(shape=self.observation_space.shape) if ob else None\n        self.ret_rms = RunningMeanStd(shape=(1,)) if ret else None\n        self.clipob = clipob\n        self.cliprew = cliprew\n        self.ret = np.zeros(())\n        self.gamma = gamma\n        self.epsilon = epsilon\n\n    def step(self, action):\n        obs, rews, news, infos = self.env.step(action)\n        infos['real_reward'] = rews\n        # print(\"before\", self.ret)\n        self.ret = self.ret * self.gamma + rews\n        # print(\"after\", self.ret)\n        obs = self._obfilt(obs)\n        if self.ret_rms:\n            self.ret_rms.update(np.array([self.ret].copy()))\n            rews = np.clip(rews / np.sqrt(self.ret_rms.var + self.epsilon), -self.cliprew, self.cliprew)\n        self.ret = self.ret * (1-float(news))\n        return obs, rews, news, infos\n\n    def _obfilt(self, obs):\n        if self.ob_rms:\n            self.ob_rms.update(obs)\n            obs = np.clip((obs - self.ob_rms.mean) / np.sqrt(self.ob_rms.var + self.epsilon), -self.clipob, self.clipob)\n            return obs\n        else:\n            return obs\n\n    def reset(self):\n        self.ret = np.zeros(())\n        obs = self.env.reset()\n        return self._obfilt(obs)\n\nif __name__ == \"__main__\":\n    parser = argparse.ArgumentParser(description='PPO agent')\n    # Common arguments\n    parser.add_argument('--exp-name', type=str, default=os.path.basename(__file__).rstrip(\".py\"),\n                       help='the name of this experiment')\n    parser.add_argument('--gym-id', type=str, default=\"MicrortsMining4x4F9-v0\",\n                       help='the id of the gym environment')\n    parser.add_argument('--seed', type=int, default=1,\n                       help='seed of the experiment')\n    parser.add_argument('--episode-length', type=int, default=0,\n                       help='the maximum length of each episode')\n    parser.add_argument('--total-timesteps', type=int, default=100000,\n                       help='total timesteps of the experiments')\n    parser.add_argument('--no-torch-deterministic', action='store_false', dest=\"torch_deterministic\", default=True,\n                       help='if toggled, `torch.backends.cudnn.deterministic=False`')\n    parser.add_argument('--no-cuda', action='store_false', dest=\"cuda\", default=True,\n                       help='if toggled, cuda will not be enabled by default')\n    parser.add_argument('--prod-mode', action='store_true', default=False,\n                       help='run the script in production mode and use wandb to log outputs')\n    parser.add_argument('--capture-video', action='store_true', default=False,\n                       help='weather to capture videos of the agent performances (check out `videos` folder)')\n    parser.add_argument('--wandb-project-name', type=str, default=\"cleanRL\",\n                       help=\"the wandb's project name\")\n    parser.add_argument('--wandb-entity', type=str, default=None,\n                       help=\"the entity (team) of wandb's project\")\n\n    # Algorithm specific arguments\n    parser.add_argument('--batch-size', type=int, default=2048,\n                       help='the batch size of ppo')\n    parser.add_argument('--minibatch-size', type=int, default=256,\n                       help='the mini batch size of ppo')\n    parser.add_argument('--gamma', type=float, default=0.99,\n                       help='the discount factor gamma')\n    parser.add_argument('--gae-lambda', type=float, default=0.97,\n                       help='the lambda for the general advantage estimation')\n    parser.add_argument('--ent-coef', type=float, default=0.01,\n                       help=\"coefficient of the entropy\")\n    parser.add_argument('--max-grad-norm', type=float, default=0.5,\n                       help='the maximum norm for the gradient clipping')\n    parser.add_argument('--clip-coef', type=float, default=0.2,\n                       help=\"the surrogate clipping coefficient\")\n    parser.add_argument('--update-epochs', type=int, default=10,\n                        help=\"the K epochs to update the policy\")\n    parser.add_argument('--kle-stop', action='store_true', default=False,\n                        help='If toggled, the policy updates will be early stopped w.r.t target-kl')\n    parser.add_argument('--kle-rollback', action='store_true', default=False,\n                        help='If toggled, the policy updates will roll back to previous policy if KL exceeds target-kl')\n    parser.add_argument('--target-kl', type=float, default=0.015,\n                        help='the target-kl variable that is referred by --kl')\n    parser.add_argument('--gae', action='store_true', default=True,\n                        help='Use GAE for advantage computation')\n    parser.add_argument('--policy-lr', type=float, default=3e-4,\n                        help=\"the learning rate of the policy optimizer\")\n    parser.add_argument('--value-lr', type=float, default=3e-4,\n                        help=\"the learning rate of the critic optimizer\")\n    parser.add_argument('--norm-obs', action='store_true', default=True,\n                        help=\"Toggles observation normalization\")\n    parser.add_argument('--norm-returns', action='store_true', default=False,\n                        help=\"Toggles returns normalization\")\n    parser.add_argument('--norm-adv', action='store_true', default=True,\n                        help=\"Toggles advantages normalization\")\n    parser.add_argument('--obs-clip', type=float, default=10.0,\n                        help=\"Value for reward clipping, as per the paper\")\n    parser.add_argument('--rew-clip', type=float, default=10.0,\n                        help=\"Value for observation clipping, as per the paper\")\n    parser.add_argument('--anneal-lr', action='store_true', default=True,\n                        help=\"Toggle learning rate annealing for policy and value networks\")\n    parser.add_argument('--weights-init', default=\"orthogonal\", choices=[\"xavier\", 'orthogonal'],\n                        help='Selects the scheme to be used for weights initialization'),\n    parser.add_argument('--clip-vloss', action=\"store_true\", default=True,\n                        help='Toggles wheter or not to use a clipped loss for the value function, as per the paper.')\n    parser.add_argument('--pol-layer-norm', action='store_true', default=False,\n                       help='Enables layer normalization in the policy network')\n    parser.add_argument('--invalid-action-penalty', type=float, default=0.00,\n                       help='the negative reward penalty for invalid actions')\n    args = parser.parse_args()\n    if not args.seed:\n        args.seed = int(time.time())\n\nargs.features_turned_on = sum([args.kle_stop, args.kle_rollback, args.gae, args.norm_obs, args.norm_returns, args.norm_adv, args.anneal_lr, args.clip_vloss, args.pol_layer_norm])\n# TRY NOT TO MODIFY: setup the environment\nexperiment_name = f\"{args.gym_id}__{args.exp_name}__{args.seed}__{int(time.time())}\"\nwriter = SummaryWriter(f\"runs/{experiment_name}\")\nwriter.add_text('hyperparameters', \"|param|value|\\n|-|-|\\n%s\" % (\n        '\\n'.join([f\"|{key}|{value}|\" for key, value in vars(args).items()])))\n\nif args.prod_mode:\n    import wandb\n    wandb.init(project=args.wandb_project_name, entity=args.wandb_entity, tensorboard=True, config=vars(args), name=experiment_name, monitor_gym=True)\n    writer = SummaryWriter(f\"/tmp/{experiment_name}\")\n    wandb.save(os.path.abspath(__file__))\n\n# TRY NOT TO MODIFY: seeding\ndevice = torch.device('cuda' if torch.cuda.is_available() and args.cuda else 'cpu')\nenv = gym.make(args.gym_id)\n# respect the default timelimit\nassert isinstance(env.action_space, MultiDiscrete), \"only MultiDiscrete action space is supported\"\nassert isinstance(env, TimeLimit) or int(args.episode_length), \"the gym env does not have a built in TimeLimit, please specify by using --episode-length\"\nif isinstance(env, TimeLimit):\n    if int(args.episode_length):\n        env._max_episode_steps = int(args.episode_length)\n    args.episode_length = env._max_episode_steps\nelse:\n    env = TimeLimit(env, int(args.episode_length))\nenv = NormalizedEnv(env.env, ob=args.norm_obs, ret=args.norm_returns, clipob=args.obs_clip, cliprew=args.rew_clip, gamma=args.gamma)\nenv = TimeLimit(env, int(args.episode_length))\nrandom.seed(args.seed)\nnp.random.seed(args.seed)\ntorch.manual_seed(args.seed)\ntorch.backends.cudnn.deterministic = args.torch_deterministic\nenv.seed(args.seed)\nenv.action_space.seed(args.seed)\nenv.observation_space.seed(args.seed)\nif args.capture_video:\n    env = Monitor(env, f'videos/{experiment_name}')\n\n# ALGO LOGIC: initialize agent here:\nclass CategoricalMasked(Categorical):\n    def __init__(self, probs=None, logits=None, validate_args=None, masks=[]):\n        self.masks = masks\n        if len(self.masks) == 0:\n            super(CategoricalMasked, self).__init__(probs, logits, validate_args)\n        else:\n            self.masks = masks.type(torch.BoolTensor).to(device)\n            logits = torch.where(self.masks, logits, torch.tensor(-1e+8).to(device))\n            super(CategoricalMasked, self).__init__(probs, logits, validate_args)\n    \n    def entropy(self):\n        if len(self.masks) == 0:\n            return super(CategoricalMasked, self).entropy()\n        p_log_p = self.logits * self.probs\n        p_log_p = torch.where(self.masks, p_log_p, torch.tensor(0.).to(device))\n        return -p_log_p.sum(-1)\n\nclass Policy(nn.Module):\n    def __init__(self):\n        super(Policy, self).__init__()\n        self.features = nn.Sequential(\n            nn.Conv2d(27, 16, kernel_size=2,),\n            nn.MaxPool2d(1),\n            nn.ReLU())\n        self.fc = nn.Sequential(\n            nn.Linear(16*3*3, 128),\n            nn.ReLU(),\n            nn.Linear(128, env.action_space.nvec.sum())\n        )\n\n    def forward(self, x):\n        x = torch.Tensor(np.moveaxis(x, -1, 1)).to(device)\n        x = self.features(x)\n        x = x.reshape(x.size(0), -1)\n        x = self.fc(x)\n        return x\n\n    def get_action(self, x, action=None):\n        logits = self.forward(x)\n        split_logits = torch.split(logits, env.action_space.nvec.tolist(), dim=1)\n        multi_categoricals = [Categorical(logits=logits) for logits in split_logits]\n        \n        if action is None:\n            action = torch.stack([categorical.sample() for categorical in multi_categoricals])\n        logprob = torch.stack([categorical.log_prob(a) for a, categorical in zip(action, multi_categoricals)])\n        # entropy = torch.stack([categorical.entropy() for categorical in multi_categoricals])\n\n        return action, logprob, [], multi_categoricals\n\nclass Value(nn.Module):\n    def __init__(self):\n        super(Value, self).__init__()\n        self.features = nn.Sequential(\n            nn.Conv2d(27, 16, kernel_size=2,),\n            nn.MaxPool2d(1),\n            nn.ReLU())\n        self.fc = nn.Sequential(\n            nn.Linear(16*3*3, 128),\n            nn.ReLU(),\n            nn.Linear(128, 1)\n        )\n\n    def forward(self, x):\n        x = torch.Tensor(np.moveaxis(x, -1, 1)).to(device)\n        x = self.features(x)\n        x = x.reshape(x.size(0), -1)\n        x = self.fc(x)\n        return x\n\ndef discount_cumsum(x, dones, gamma):\n    \"\"\"\n    computing discounted cumulative sums of vectors that resets with dones\n    input:\n        vector x,  vector dones,\n        [x0,       [0,\n         x1,        0,\n         x2         1,\n         x3         0, \n         x4]        0]\n    output:\n        [x0 + discount * x1 + discount^2 * x2,\n         x1 + discount * x2,\n         x2,\n         x3 + discount * x4,\n         x4]\n    \"\"\"\n    discount_cumsum = np.zeros_like(x)\n    discount_cumsum[-1] = x[-1]\n    for t in reversed(range(x.shape[0]-1)):\n        discount_cumsum[t] = x[t] + gamma * discount_cumsum[t+1] * (1-dones[t])\n    return discount_cumsum\n\npg = Policy().to(device)\nvf = Value().to(device)\n\n# MODIFIED: Separate optimizer and learning rates\npg_optimizer = optim.Adam(list(pg.parameters()), lr=args.policy_lr)\nv_optimizer = optim.Adam(list(vf.parameters()), lr=args.value_lr)\n\n# MODIFIED: Initializing learning rate anneal scheduler when need\nif args.anneal_lr:\n    anneal_fn = lambda f: max(0, 1-f / args.total_timesteps)\n    pg_lr_scheduler = optim.lr_scheduler.LambdaLR(pg_optimizer, lr_lambda=anneal_fn)\n    vf_lr_scheduler = optim.lr_scheduler.LambdaLR(v_optimizer, lr_lambda=anneal_fn)\n\nloss_fn = nn.MSELoss()\n\n# TRY NOT TO MODIFY: start the game\nglobal_step = 0\nwhile global_step < args.total_timesteps:\n    if args.capture_video:\n        env.stats_recorder.done=True\n    next_obs = np.array(env.reset())\n\n    # ALGO Logic: Storage for epoch data\n    obs = np.empty((args.batch_size,) + env.observation_space.shape)\n\n    actions = np.empty((args.batch_size,) + env.action_space.shape)\n    logprobs = torch.zeros((env.action_space.nvec.shape[0], args.batch_size,)).to(device)\n\n    rewards = np.zeros((args.batch_size,))\n    real_rewards = []\n    invalid_action_stats = []\n\n    dones = np.zeros((args.batch_size,))\n    values = torch.zeros((args.batch_size,)).to(device)\n\n    # TRY NOT TO MODIFY: prepare the execution of the game.\n    for step in range(args.batch_size):\n        env.render()\n        global_step += 1\n        obs[step] = next_obs.copy()\n\n        # ALGO LOGIC: put action logic here\n        with torch.no_grad():\n            values[step] = vf.forward(obs[step:step+1])\n            action, logproba, _, probs = pg.get_action(obs[step:step+1])\n        \n        actions[step] = action[:,0].data.cpu().numpy()\n        logprobs[:,[step]] = logproba\n\n        # TRY NOT TO MODIFY: execute the game and log data.\n        try:\n            next_obs, rewards[step], dones[step], info = env.step(action[:,0].data.cpu().numpy())\n            rewards[step] += pd.DataFrame([info['invalid_action_stats']]).sum().sum() * args.invalid_action_penalty\n        except Exception as e:\n            print(e)\n            print(e.stacktrace())\n        real_rewards += [info['real_reward']]\n        invalid_action_stats += [info['invalid_action_stats']]\n        next_obs = np.array(next_obs)\n\n        # Annealing the rate if instructed to do so.\n        if args.anneal_lr:\n            pg_lr_scheduler.step()\n            vf_lr_scheduler.step()\n\n        if dones[step]:\n            # Computing the discounted returns:\n            writer.add_scalar(\"charts/episode_reward\", np.sum(real_rewards), global_step)\n            print(f\"global_step={global_step}, episode_reward={np.sum(real_rewards)}\")\n            real_rewards = []\n            for key, idx in zip(info['invalid_action_stats'], range(len(info['invalid_action_stats']))):\n                writer.add_scalar(f\"stats/{key}\", pd.DataFrame(invalid_action_stats).sum(0)[idx], global_step)\n            invalid_action_stats = []\n            next_obs = np.array(env.reset())\n\n    # bootstrap reward if not done. reached the batch limit\n    last_value = 0\n    if not dones[step]:\n        last_value = vf.forward(next_obs.reshape((1,)+next_obs.shape))[0].detach().cpu().numpy()[0]\n    bootstrapped_rewards = np.append(rewards, last_value)\n\n    # calculate the returns and advantages\n    if args.gae:\n        bootstrapped_values = np.append(values.detach().cpu().numpy(), last_value)\n        deltas = bootstrapped_rewards[:-1] + args.gamma * bootstrapped_values[1:] * (1-dones) - bootstrapped_values[:-1]\n        advantages = discount_cumsum(deltas, dones, args.gamma * args.gae_lambda)\n        advantages = torch.Tensor(advantages).to(device)\n        returns = advantages + values\n    else:\n        returns = discount_cumsum(bootstrapped_rewards, dones, args.gamma)[:-1]\n        advantages = returns - values.detach().cpu().numpy()\n        advantages = torch.Tensor(advantages).to(device)\n        returns = torch.Tensor(returns).to(device)\n\n    # Advantage normalization\n    if args.norm_adv:\n        EPS = 1e-10\n        advantages = (advantages - advantages.mean()) / (advantages.std() + EPS)\n\n    # Optimizaing policy network\n    entropys = []\n    target_pg = Policy().to(device)\n    inds = np.arange(args.batch_size,)\n    for i_epoch_pi in range(args.update_epochs):\n        np.random.shuffle(inds)\n        for start in range(0, args.batch_size, args.minibatch_size):\n            end = start + args.minibatch_size\n            minibatch_ind = inds[start:end]\n            target_pg.load_state_dict(pg.state_dict())\n            \n            _, newlogproba, _, _ = pg.get_action(\n                obs[minibatch_ind],\n                torch.LongTensor(actions[minibatch_ind].astype(np.int)).to(device).T)\n            ratio = (newlogproba - logprobs[:,minibatch_ind]).exp()\n\n            # Policy loss as in OpenAI SpinUp\n            clip_adv = torch.where(advantages[minibatch_ind] > 0,\n                                    (1.+args.clip_coef) * advantages[minibatch_ind],\n                                    (1.-args.clip_coef) * advantages[minibatch_ind]).to(device)\n\n            # Entropy computation with resampled actions\n            entropy = -(newlogproba.exp() * newlogproba).mean()\n            entropys.append(entropy.item())\n\n            policy_loss = -torch.min(ratio * advantages[minibatch_ind], clip_adv) + args.ent_coef * entropy\n            policy_loss = policy_loss.mean()\n            \n            pg_optimizer.zero_grad()\n            policy_loss.backward()\n            nn.utils.clip_grad_norm_(pg.parameters(), args.max_grad_norm)\n            pg_optimizer.step()\n\n            approx_kl = (logprobs[:,minibatch_ind] - newlogproba).mean()\n            # Resample values\n            new_values = vf.forward(obs[minibatch_ind]).view(-1)\n\n            # Value loss clipping\n            if args.clip_vloss:\n                v_loss_unclipped = ((new_values - returns[minibatch_ind]) ** 2)\n                v_clipped = values[minibatch_ind] + torch.clamp(new_values - values[minibatch_ind], -args.clip_coef, args.clip_coef)\n                v_loss_clipped = (v_clipped - returns[minibatch_ind])**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 = torch.mean((returns[minibatch_ind]- new_values).pow(2))\n\n            v_optimizer.zero_grad()\n            v_loss.backward()\n            nn.utils.clip_grad_norm_(vf.parameters(), args.max_grad_norm)\n            v_optimizer.step()\n\n        if args.kle_stop:\n            if approx_kl > args.target_kl:\n                break\n        if args.kle_rollback:\n            if (logprobs[:,minibatch_ind] - \n                pg.get_action(\n                    obs[minibatch_ind],\n                    torch.LongTensor(actions[minibatch_ind].astype(np.int)).to(device).T,\n                    )[1]).mean() > args.target_kl:\n                pg.load_state_dict(target_pg.state_dict())\n                break\n\n    # TRY NOT TO MODIFY: record rewards for plotting purposes\n    writer.add_scalar(\"losses/value_loss\", v_loss.item(), global_step)\n    writer.add_scalar(\"charts/policy_learning_rate\", pg_optimizer.param_groups[0]['lr'], global_step)\n    writer.add_scalar(\"charts/value_learning_rate\", v_optimizer.param_groups[0]['lr'], global_step)\n    writer.add_scalar(\"losses/policy_loss\", policy_loss.item(), global_step)\n    writer.add_scalar(\"losses/entropy\", np.mean(entropys), global_step)\n    writer.add_scalar(\"losses/approx_kl\", approx_kl.item(), global_step)\n    if args.kle_stop or args.kle_rollback:\n        writer.add_scalar(\"debug/pg_stop_iter\", i_epoch_pi, global_step)\n\nenv.close()\nwriter.close()\n"
  },
  {
    "path": "plots/analysis.py",
    "content": "import wandb\nimport numpy as np\nimport pandas as pd \napi = wandb.Api()\nwandb_entity = os.environ['WANDB_ENTITY']\n\n# Project is specified by <entity/project-name>\nruns = api.runs(f\"{wandb_entity}/invalid_action_masking\")\nanalysis = True\nsummary_list = [] \nconfig_list = [] \nname_list = [] \nfor run in runs: \n    # run.summary are the output key/values like accuracy.\n    # We call ._json_dict to omit large files \n    summary_json_dict = run.summary._json_dict\n    if analysis:\n        summary_json_dict = summary_json_dict.copy()\n        history = pd.DataFrame(run.scan_history())\n        history['rollling_e'] = history['charts/episode_reward'].dropna().rolling(10).mean()\n        first_best_reward_idx = (history[\"rollling_e\"] >= 40.0).idxmax()\n        if history.iloc[first_best_reward_idx][\"rollling_e\"] >= 40.0:\n            summary_json_dict[\"first_learned_timestep\"] = history.iloc[first_best_reward_idx][\"global_step\"] / 500000\n        else:\n            summary_json_dict[\"first_learned_timestep\"] = 1\n        summary_json_dict[\"first_reward_timestep\"] = history.iloc[(history['charts/episode_reward'] > 0).idxmax()][\"global_step\"] / 500000\n\n        # mask removed logic\n        if run.config[\"exp_name\"] == \"ppo\":\n            history['evals_rollling_e'] = history['evals/charts/episode_reward'].dropna().rolling(10).mean()\n            first_best_reward_idx = (history[\"evals_rollling_e\"] >= 40.0).idxmax()\n            if history.iloc[first_best_reward_idx][\"evals_rollling_e\"] >= 40.0:\n                summary_json_dict[\"evals_first_learned_timestep\"] = history.iloc[first_best_reward_idx][\"global_step\"] / 500000\n            else:\n                summary_json_dict[\"evals_first_learned_timestep\"] = 1\n            summary_json_dict[\"evals_first_reward_timestep\"] = history.iloc[(history['charts/episode_reward'] > 0).idxmax()][\"global_step\"] / 500000\n            \n        \n        summary_json_dict[\"charts/episode_reward\"] = history[\"charts/episode_reward\"][-10:].mean()\n        summary_json_dict[\"losses/approx_kl\"] = history['losses/approx_kl'].astype(np.float64).dropna()[-10:].mean()\n        summary_json_dict['stats/num_invalid_action_null'] = history['stats/num_invalid_action_null'].dropna()[-10:].mean()\n        summary_json_dict['stats/num_invalid_action_busy_unit'] = history['stats/num_invalid_action_busy_unit'].dropna()[-10:].mean()\n        summary_json_dict['stats/num_invalid_action_ownership'] = history['stats/num_invalid_action_ownership'].dropna()[-10:].mean()\n        \n        if run.config[\"exp_name\"] == \"ppo\":\n            summary_json_dict['evals/charts/episode_reward'] = history['evals/charts/episode_reward'][-10:].mean()\n            summary_json_dict['evals/stats/num_invalid_action_null'] = history['evals/stats/num_invalid_action_null'].dropna()[-10:].mean()\n            summary_json_dict['evals/stats/num_invalid_action_busy_unit'] = history['evals/stats/num_invalid_action_busy_unit'].dropna()[-10:].mean()\n            summary_json_dict['evals/stats/num_invalid_action_ownership'] = history['evals/stats/num_invalid_action_ownership'].dropna()[-10:].mean()\n    summary_list.append(summary_json_dict)\n    # run.config is the input metrics.\n    # We remove special values that start with _.\n    config = {k:v for k,v in run.config.items() if not k.startswith('_')}\n    config_list.append(config) \n\n    # run.name is the name of the run.\n    name_list.append(run.name)       \n\nsummary_df = pd.DataFrame.from_records(summary_list) \nconfig_df = pd.DataFrame.from_records(config_list) \nname_df = pd.DataFrame({'name': name_list}) \nall_df = pd.concat([name_df, config_df,summary_df], axis=1)\n\nall_df.to_csv(\"project.csv\")\nall_df[\"losses/approx_kl\"] = all_df[\"losses/approx_kl\"].astype(np.float64)\n\n# mask removal\nmask_removed = all_df[all_df[\"exp_name\"]==\"ppo\"].copy()\nmask_removed['charts/episode_reward'] = mask_removed['evals/charts/episode_reward']\nmask_removed['stats/num_invalid_action_null'] = mask_removed['evals/stats/num_invalid_action_null']\nmask_removed['stats/num_invalid_action_busy_unit'] = mask_removed['evals/stats/num_invalid_action_busy_unit']\nmask_removed['stats/num_invalid_action_ownership'] = mask_removed['evals/stats/num_invalid_action_ownership']\nmask_removed['first_learned_timestep'] = mask_removed['evals_first_learned_timestep']\nmask_removed['first_reward_timestep'] = mask_removed['evals_first_reward_timestep']\nmask_removed[\"exp_name\"] = \"masking removed\"\nfinal_all_df = all_df.append(mask_removed, ignore_index=True)\n\n\n# change names\nfinal_all_df.loc[final_all_df[\"gym_id\"]==\"MicrortsMining4x4F9-v0\", \"gym_id\"] = '04x04'\nfinal_all_df.loc[(final_all_df[\"gym_id\"]==\"MicrortsMining10x10F9-v0\"), \"gym_id\"] = '10x10'\nfinal_all_df.loc[final_all_df[\"gym_id\"]==\"MicrortsMining16x16F9-v0\", \"gym_id\"] = '16x16'\nfinal_all_df.loc[final_all_df[\"gym_id\"]==\"MicrortsMining24x24F9-v0\", \"gym_id\"] = '24x24'\n\nfinal_all_df.loc[final_all_df[\"exp_name\"]==\"masking removed\", \"exp_name\"] = 'Masking removed'\nfinal_all_df.loc[(final_all_df[\"exp_name\"]==\"ppo\"), \"exp_name\"] = 'Invalid action masking'\nfinal_all_df.loc[final_all_df[\"exp_name\"]==\"ppo_no_adj\", \"exp_name\"] = 'Naive invalid action masking'\nfinal_all_df.loc[final_all_df[\"exp_name\"]==\"ppo_no_mask\", \"exp_name\"] = 'Invalid action penalty'\n\nresults_df = final_all_df.fillna(0).groupby(\n    ['exp_name','gym_id',\"invalid_action_penalty\"]\n).mean()[[\n    'charts/episode_reward',\n    'losses/approx_kl',\n    'stats/num_invalid_action_null',\n    'stats/num_invalid_action_busy_unit',\n    'stats/num_invalid_action_ownership',\n    \"first_learned_timestep\",\n    \"first_reward_timestep\"\n]]\nfinal_print_df = results_df.round(2)\nfinal_print_df['losses/approx_kl'] = results_df['losses/approx_kl'].round(5)\n# final_print_df['first_learned_timestep'] = results_df['first_learned_timestep'].round(4)\n# final_print_df['first_reward_timestep'] = results_df['first_reward_timestep'].round(4)\nfinal_print_df['first_learned_timestep'] = pd.Series([\"{0:.2f}%\".format(val * 100) for val in results_df['first_learned_timestep'].round(4)], index = results_df.index)\nfinal_print_df['first_reward_timestep'] = pd.Series([\"{0:.2f}%\".format(val * 100) for val in results_df['first_reward_timestep'].round(4)], index = results_df.index)\nprint(final_print_df.to_latex())\n\nprint(final_print_df.drop(columns=['losses/approx_kl']).to_latex())\n\n\n# calculate the first time the algorithm solves the environment\n\n# , 'losses/value_loss',\n#        'losses/policy_loss', 'charts/episode_reward',\n#        , ,\n#        ,\n#        'charts/episode_reward/ResourceGatherRewardFunction',\n#        'evals/charts/episode_reward', 'evals/stats/num_invalid_action_null',\n#        'evals/stats/num_invalid_action_busy_unit',\n#        'evals/stats/num_invalid_action_ownership'"
  },
  {
    "path": "plots/approx_kl.py",
    "content": "from os import path\nimport pickle\nimport wandb\nimport pandas as pd\nimport numpy as np\nimport seaborn as sns\nimport matplotlib.pyplot as plt\nimport os\nimport argparse\nfrom distutils.util import strtobool\nimport matplotlib as mpl\nmpl.rcParams['text.usetex'] = True\nmpl.rcParams['text.latex.preamble'] = [r'\\usepackage{amsmath}'] #for \\text command\n\n\nparser = argparse.ArgumentParser(description='CleanRL Plots')\n# Common arguments\nparser.add_argument('--wandb-project', type=str, default=\"costa-huang/invalid-action-masking\",\n                   help='the name of wandb project (e.g. cleanrl/cleanrl)')\nparser.add_argument('--feature-of-interest', type=str, default='losses/approx_kl',\n                   help='which feature to be plotted on the y-axis')\nparser.add_argument('--hyper-params-tuned', nargs='+', default=[],\n                    help='the hyper parameters tuned')\n# parser.add_argument('--scan-history', type=lambda x:bool(strtobool(x)), default=False, nargs='?', const=True,\n#                     help='if toggled, cuda will not be enabled by default')\nparser.add_argument('--interested-exp-names', nargs='+', default=[],\n                    help='the hyper parameters tuned')\nparser.add_argument('--samples', type=int, default=500,\n                    help='the sampled point of the run')\nparser.add_argument('--smooth-weight', type=float, default=0.90,\n                    help='the weight parameter of the exponential moving average')\nparser.add_argument('--last-n-episodes', type=int, default=50,\n                   help='for analysis only; the last n episodes from which the mean of the feature of interest is calculated')\nparser.add_argument('--num-points-x-axis', type=int, default=500,\n                   help='the number of points in the x-axis')\nparser.add_argument('--font-size', type=int, default=18,\n                   help='the font size of the plots')\nparser.add_argument('--x-label', type=str, default=\"Time Steps\",\n                   help='the label of x-axis')\nparser.add_argument('--y-label', type=str, default=\"KL Divergence\",\n                   help='the label of y-axis')\nparser.add_argument('--y-lim-bottom', type=float, default=0.0,\n                   help='the bottom limit for the y-axis')\nparser.add_argument('--output-format', type=str, default=\"pdf\",\n                   help='either `pdf`, `png`, or `svg`')\nargs = parser.parse_args()\napi = wandb.Api()\n\n# hacks\nenv_dict = {\n    # 'MicrortsAttackShapedReward-v1': 'MicrortsAttackHRL-v1',\n    # 'MicrortsProduceCombatUnitsShapedReward-v1':  'MicrortsProduceCombatUnitHRL-v1',\n    # 'MicrortsRandomEnemyShapedReward3-v1': 'MicrortsRandomEnemyHRL3-v1',\n}\nexp_convert_dict = {\n    'ppo': 'Invalid action masking',\n    'ppo_no_mask-0': 'Invalid action penalty, $r_{\\\\text{invalid}}=0$',\n    'ppo_no_mask--0.1': 'Invalid action penalty, $r_{\\\\text{invalid}}=-0.1$',\n    'ppo_no_mask--0.01': 'Invalid action penalty, $r_{\\\\text{invalid}}=-0.01$',\n    'ppo_no_mask--1': 'Invalid action penalty, $r_{\\\\text{invalid}}=-1$',\n    'ppo-maskrm': 'Masking removed',\n    'ppo_no_adj': 'Naive invalid action masking',\n}\n\n# args.feature_of_interest = 'charts/episode_reward'\nfeature_name = args.feature_of_interest.replace(\"/\", \"_\")\nif not os.path.exists(feature_name):\n    os.makedirs(feature_name)\n\nif not path.exists(f\"{feature_name}/all_df_cache.pkl\"):\n    # Change oreilly-class/cifar to <entity/project-name>\n    runs = api.runs(args.wandb_project)\n    summary_list = [] \n    config_list = [] \n    name_list = []\n    envs = {}\n    data = []\n    exp_names = []\n    \n    for idx, run in enumerate(runs):\n        if args.feature_of_interest in run.summary:\n            metrics_dataframe = run.history(keys=[args.feature_of_interest, 'global_step'], samples=args.samples)\n            exp_name = run.config['exp_name']\n            for param in args.hyper_params_tuned:\n                if param in run.config:\n                    exp_name += \"-\" + param + \"-\" + str(run.config[param]) + \"-\"\n            \n            # hacks\n            if \"invalid_action_penalty\" in run.config:\n                exp_name = run.config['exp_name']+\"-\"+str(run.config['invalid_action_penalty'])\n                \n            # hacks\n            if run.config[\"gym_id\"] in env_dict:\n                exp_name += \"shaped\"\n                run.config[\"gym_id\"] = env_dict[run.config[\"gym_id\"]]\n\n            metrics_dataframe.insert(len(metrics_dataframe.columns), \"algo\", exp_name)\n            exp_names += [exp_name]\n            metrics_dataframe.insert(len(metrics_dataframe.columns), \"seed\", run.config['seed'])\n            \n            data += [metrics_dataframe]\n            if run.config[\"gym_id\"] not in envs:\n                envs[run.config[\"gym_id\"]] = [metrics_dataframe]\n                envs[run.config[\"gym_id\"]+\"total_timesteps\"] = run.config[\"total_timesteps\"]\n            else:\n                envs[run.config[\"gym_id\"]] += [metrics_dataframe]\n            \n            # run.summary are the output key/values like accuracy.  We call ._json_dict to omit large files \n            summary_list.append(run.summary._json_dict) \n        \n            # run.config is the input metrics.  We remove special values that start with _.\n            config_list.append({k:v for k,v in run.config.items() if not k.startswith('_')}) \n        \n            # run.name is the name of the run.\n            name_list.append(run.name)       \n    \n    \n    summary_df = pd.DataFrame.from_records(summary_list) \n    config_df = pd.DataFrame.from_records(config_list) \n    name_df = pd.DataFrame({'name': name_list}) \n    all_df = pd.concat([name_df, config_df,summary_df], axis=1)\n    data = pd.concat(data, ignore_index=True)\n    \n    with open(f'{feature_name}/all_df_cache.pkl', 'wb') as handle:\n        pickle.dump(all_df, handle, protocol=pickle.HIGHEST_PROTOCOL)\n    with open(f'{feature_name}/envs_cache.pkl', 'wb') as handle:\n        pickle.dump(envs, handle, protocol=pickle.HIGHEST_PROTOCOL)\n    with open(f'{feature_name}/exp_names_cache.pkl', 'wb') as handle:\n        pickle.dump(exp_names, handle, protocol=pickle.HIGHEST_PROTOCOL)\nelse:\n    with open(f'{feature_name}/all_df_cache.pkl', 'rb') as handle:\n        all_df = pickle.load(handle)\n    with open(f'{feature_name}/envs_cache.pkl', 'rb') as handle:\n        envs = pickle.load(handle)\n    with open(f'{feature_name}/exp_names_cache.pkl', 'rb') as handle:\n        exp_names = pickle.load(handle)\nprint(\"data loaded\")\n\n# https://stackoverflow.com/questions/42281844/what-is-the-mathematics-behind-the-smoothing-parameter-in-tensorboards-scalar#_=_\ndef smooth(scalars, weight):  # Weight between 0 and 1\n    last = scalars[0]  # First value in the plot (first timestep)\n    smoothed = list()\n    for point in scalars:\n        smoothed_val = last * weight + (1 - weight) * point  # Calculate smoothed value\n        smoothed.append(smoothed_val)                        # Save it\n        last = smoothed_val                                  # Anchor the last smoothed value\n\n    return smoothed\n\n#smoothing\nfor env in envs:\n    if not env.endswith(\"total_timesteps\"):\n        for idx, metrics_dataframe in enumerate(envs[env]):\n            envs[env][idx] = metrics_dataframe.dropna(subset=[args.feature_of_interest])\n#             envs[env][idx][args.feature_of_interest] = smooth(metrics_dataframe[args.feature_of_interest], 0.85)\n\nsns.set(style=\"darkgrid\")\ndef get_df_for_env(gym_id):\n    env_total_timesteps = envs[gym_id+\"total_timesteps\"]\n    env_increment = env_total_timesteps / 500\n    envs_same_x_axis = []\n    for sampled_run in envs[gym_id]:\n        df = pd.DataFrame(columns=sampled_run.columns)\n        x_axis = [i*env_increment for i in range(500-2)]\n        current_row = 0\n        for timestep in x_axis:\n            while sampled_run.iloc[current_row][\"global_step\"] < timestep:\n                current_row += 1\n                if current_row > len(sampled_run)-2:\n                    break\n            if current_row > len(sampled_run)-2:\n                break\n            temp_row = sampled_run.iloc[current_row].copy()\n            temp_row[\"global_step\"] = timestep\n            df = df.append(temp_row)\n        \n        envs_same_x_axis += [df]\n    return pd.concat(envs_same_x_axis, ignore_index=True)\n\ndef export_legend(ax, filename=\"legend.pdf\"):\n    # import matplotlib as mpl\n    # mpl.rcParams['text.usetex'] = True\n    # mpl.rcParams['text.latex.preamble'] = [r'\\usepackage{amsmath}'] #for \\text command\n    fig2 = plt.figure()\n    ax2 = fig2.add_subplot()\n    ax2.axis('off')\n    handles, labels = ax.get_legend_handles_labels()\n\n    legend = ax2.legend(handles=handles, labels=labels, frameon=False, loc='lower center', ncol=4, fontsize=20, handlelength=1)\n    for text in legend.get_texts():\n        if text.get_text() in exp_convert_dict:\n            text.set_text(exp_convert_dict[text.get_text()])\n    for line in legend.get_lines():\n        line.set_linewidth(4.0)\n    fig  = legend.figure\n    fig.canvas.draw()\n    bbox  = legend.get_window_extent().transformed(fig.dpi_scale_trans.inverted())\n    fig.savefig(filename, dpi=\"figure\", bbox_inches=bbox)\n    fig.clf()\n\nif not os.path.exists(f\"{feature_name}/data\"):\n    os.makedirs(f\"{feature_name}/data\")\nif not os.path.exists(f\"{feature_name}/plots\"):\n    os.makedirs(f\"{feature_name}/plots\")\nif not os.path.exists(f\"{feature_name}/legends\"):\n    os.makedirs(f\"{feature_name}/legends\")\n\n\ninterested_exp_names = sorted(list(set(exp_names))) # ['ppo_continuous_action', 'ppo_atari_visual']\ncurrent_palette = sns.color_palette(n_colors=len(interested_exp_names))\ncurrent_palette_dict = dict(zip(interested_exp_names, current_palette))\nif args.interested_exp_names:\n    interested_exp_names = args.interested_exp_names\nprint(current_palette_dict)\nlegend_df = pd.DataFrame()\n\nif args.font_size:\n    plt.rc('axes', titlesize=args.font_size)     # fontsize of the axes title\n    plt.rc('axes', labelsize=args.font_size)    # fontsize of the x and y labels\n    plt.rc('xtick', labelsize=args.font_size)    # fontsize of the tick labels\n    plt.rc('ytick', labelsize=args.font_size)    # fontsize of the tick labels\n    plt.rc('legend', fontsize=args.font_size)    # legend fontsize\n\nstats = {item: [] for item in [\"gym_id\", \"exp_name\", args.feature_of_interest]}\n# uncommenet the following to generate all figures\nfor env in set(all_df[\"gym_id\"]):\n    if not path.exists(f\"{feature_name}/data/{env}.pkl\"):\n        with open(f\"{feature_name}/data/{env}.pkl\", 'wb') as handle:\n            data = get_df_for_env(env)\n            data[\"seed\"] = data[\"seed\"].astype(float)\n            data[args.feature_of_interest] = data[args.feature_of_interest].astype(float)\n            pickle.dump(data, handle, protocol=pickle.HIGHEST_PROTOCOL)\n    else:\n        with open(f\"{feature_name}/data/{env}.pkl\", 'rb') as handle:\n            data = pickle.load(handle)\n            print(f\"{env}'s data loaded\")\n    def _smooth(df):\n        df[args.feature_of_interest] = smooth(list(df[args.feature_of_interest]), args.smooth_weight)\n        return df\n\n    legend_df = legend_df.append(data)\n    ax = sns.lineplot(data=data.groupby([\"seed\", \"algo\"]).apply(_smooth).loc[data['algo'].isin(interested_exp_names)], x=\"global_step\", y=args.feature_of_interest, hue=\"algo\", ci='sd', palette=current_palette_dict,)\n    ax.ticklabel_format(style='sci', scilimits=(0,0), axis='x')\n    ax.set(xlabel=args.x_label, ylabel=args.y_label)\n    \n    # hack \n    ax.set_ylim(0, 0.07)\n    # ax.set(ylabel=\"\")\n    # ax.set_xticks([])\n    \n    ax.legend().remove()\n    if args.y_lim_bottom:\n        plt.ylim(bottom=args.y_lim_bottom)\n    # plt.title(env)\n    plt.tight_layout()\n    plt.savefig(f\"{feature_name}/plots/{env}.{args.output_format}\")\n    plt.clf()\n    \n    for algo in interested_exp_names:\n        algo_data = data.loc[data['algo'].isin([algo])]\n        last_n_episodes_global_step = sorted(algo_data[\"global_step\"].unique())[-args.last_n_episodes]\n        last_n_episodes_features = algo_data[algo_data['global_step'] > last_n_episodes_global_step].groupby(\n            ['seed']\n        ).mean()[args.feature_of_interest]\n        \n        for item in last_n_episodes_features:\n            stats[args.feature_of_interest] += [item]\n            if algo in exp_convert_dict:\n                stats['exp_name'] += [exp_convert_dict[algo]]\n            else:\n                stats['exp_name'] += [algo]\n            stats['gym_id'] += [env]\n\n# export legend\nlegend_df = legend_df.reset_index()\nax = sns.lineplot(data=legend_df, x=\"global_step\", y=args.feature_of_interest, hue=\"algo\", ci='sd', palette=current_palette_dict,)\nax.set(xlabel='Time Steps', ylabel='Average Episode Reward')\nax.legend().remove()\nexport_legend(ax, f\"{feature_name}/legend.{args.output_format}\")\nplt.clf()\n\n\n# analysis\nstats_df = pd.DataFrame(stats)\ng = stats_df.groupby(\n    ['gym_id','exp_name']\n).agg(lambda x: f\"{np.mean(x):.2f} ± {np.std(x):.2f}\")\nprint(g.reset_index().pivot('exp_name', 'gym_id' , args.feature_of_interest).to_latex().replace(\"±\", \"$\\pm$\"))"
  },
  {
    "path": "plots/episode_reward.py",
    "content": "from os import path\nimport pickle\nimport wandb\nimport pandas as pd\nimport numpy as np\nimport seaborn as sns\nimport matplotlib.pyplot as plt\nimport os\nimport argparse\nfrom distutils.util import strtobool\nimport matplotlib as mpl\nmpl.rcParams['text.usetex'] = True\nmpl.rcParams['text.latex.preamble'] = [r'\\usepackage{amsmath}'] #for \\text command\n\n\nparser = argparse.ArgumentParser(description='CleanRL Plots')\n# Common arguments\nparser.add_argument('--wandb-project', type=str, default=\"costa-huang/invalid-action-masking\",\n                   help='the name of wandb project (e.g. cleanrl/cleanrl)')\nparser.add_argument('--feature-of-interest', type=str, default='charts/episode_reward',\n                   help='which feature to be plotted on the y-axis')\nparser.add_argument('--hyper-params-tuned', nargs='+', default=[],\n                    help='the hyper parameters tuned')\n# parser.add_argument('--scan-history', type=lambda x:bool(strtobool(x)), default=False, nargs='?', const=True,\n#                     help='if toggled, cuda will not be enabled by default')\nparser.add_argument('--interested-exp-names', nargs='+', default=[],\n                    help='the hyper parameters tuned')\nparser.add_argument('--samples', type=int, default=500,\n                    help='the sampled point of the run')\nparser.add_argument('--smooth-weight', type=float, default=0.90,\n                    help='the weight parameter of the exponential moving average')\nparser.add_argument('--last-n-episodes', type=int, default=50,\n                   help='for analysis only; the last n episodes from which the mean of the feature of interest is calculated')\nparser.add_argument('--num-points-x-axis', type=int, default=500,\n                   help='the number of points in the x-axis')\nparser.add_argument('--font-size', type=int, default=18,\n                   help='the font size of the plots')\nparser.add_argument('--x-label', type=str, default=\"Time Steps\",\n                   help='the label of x-axis')\nparser.add_argument('--y-label', type=str, default=\"Episodic Return\",\n                   help='the label of y-axis')\nparser.add_argument('--y-lim-bottom', type=float, default=0.0,\n                   help='the bottom limit for the y-axis')\nparser.add_argument('--output-format', type=str, default=\"pdf\",\n                   help='either `pdf`, `png`, or `svg`')\nargs = parser.parse_args()\napi = wandb.Api()\n\n# hacks\nenv_dict = {\n    # 'MicrortsAttackShapedReward-v1': 'MicrortsAttackHRL-v1',\n    # 'MicrortsProduceCombatUnitsShapedReward-v1':  'MicrortsProduceCombatUnitHRL-v1',\n    # 'MicrortsRandomEnemyShapedReward3-v1': 'MicrortsRandomEnemyHRL3-v1',\n}\nexp_convert_dict = {\n    'ppo': 'Invalid action masking',\n    'ppo_no_mask-0': 'Invalid action penalty, $r_{\\\\text{invalid}}=0$',\n    'ppo_no_mask--0.1': 'Invalid action penalty, $r_{\\\\text{invalid}}=-0.1$',\n    'ppo_no_mask--0.01': 'Invalid action penalty, $r_{\\\\text{invalid}}=-0.01$',\n    'ppo_no_mask--1': 'Invalid action penalty, $r_{\\\\text{invalid}}=-1$',\n    'ppo-maskrm': 'Masking removed',\n    'ppo_no_adj': 'Naive invalid action masking',\n}\n\n# args.feature_of_interest = 'charts/episode_reward'\nfeature_name = args.feature_of_interest.replace(\"/\", \"_\")\nif not os.path.exists(feature_name):\n    os.makedirs(feature_name)\n\nif not path.exists(f\"{feature_name}/all_df_cache.pkl\"):\n    # Change oreilly-class/cifar to <entity/project-name>\n    runs = api.runs(args.wandb_project)\n    summary_list = [] \n    config_list = [] \n    name_list = []\n    envs = {}\n    data = []\n    exp_names = []\n    \n    for idx, run in enumerate(runs):\n        if args.feature_of_interest in run.summary:\n            metrics_dataframe = run.history(keys=[args.feature_of_interest, 'global_step'], samples=args.samples)\n            exp_name = run.config['exp_name']\n            for param in args.hyper_params_tuned:\n                if param in run.config:\n                    exp_name += \"-\" + param + \"-\" + str(run.config[param]) + \"-\"\n            \n            # hacks\n            if \"invalid_action_penalty\" in run.config:\n                exp_name = run.config['exp_name']+\"-\"+str(run.config['invalid_action_penalty'])\n                \n            # hacks\n            if run.config[\"gym_id\"] in env_dict:\n                exp_name += \"shaped\"\n                run.config[\"gym_id\"] = env_dict[run.config[\"gym_id\"]]\n\n            metrics_dataframe.insert(len(metrics_dataframe.columns), \"algo\", exp_name)\n            exp_names += [exp_name]\n            metrics_dataframe.insert(len(metrics_dataframe.columns), \"seed\", run.config['seed'])\n            \n            data += [metrics_dataframe]\n            if run.config[\"gym_id\"] not in envs:\n                envs[run.config[\"gym_id\"]] = [metrics_dataframe]\n                envs[run.config[\"gym_id\"]+\"total_timesteps\"] = run.config[\"total_timesteps\"]\n            else:\n                envs[run.config[\"gym_id\"]] += [metrics_dataframe]\n                \n            # hacks\n            if exp_name == \"ppo\":\n                metrics_dataframe = run.history(keys=[\"evals/\"+args.feature_of_interest, 'global_step'], samples=args.samples)\n                exp_name = \"ppo-maskrm\"\n                \n                metrics_dataframe.insert(len(metrics_dataframe.columns), \"algo\", exp_name)\n                exp_names += [exp_name]\n                metrics_dataframe.insert(len(metrics_dataframe.columns), \"seed\", run.config['seed'])\n                metrics_dataframe[args.feature_of_interest] = metrics_dataframe[\"evals/\"+args.feature_of_interest]\n                data += [metrics_dataframe]\n                if run.config[\"gym_id\"] not in envs:\n                    envs[run.config[\"gym_id\"]] = [metrics_dataframe]\n                    envs[run.config[\"gym_id\"]+\"total_timesteps\"] = run.config[\"total_timesteps\"]\n                else:\n                    envs[run.config[\"gym_id\"]] += [metrics_dataframe]\n            \n            # run.summary are the output key/values like accuracy.  We call ._json_dict to omit large files \n            summary_list.append(run.summary._json_dict) \n        \n            # run.config is the input metrics.  We remove special values that start with _.\n            config_list.append({k:v for k,v in run.config.items() if not k.startswith('_')}) \n        \n            # run.name is the name of the run.\n            name_list.append(run.name)       \n    \n    \n    summary_df = pd.DataFrame.from_records(summary_list) \n    config_df = pd.DataFrame.from_records(config_list) \n    name_df = pd.DataFrame({'name': name_list}) \n    all_df = pd.concat([name_df, config_df,summary_df], axis=1)\n    data = pd.concat(data, ignore_index=True)\n    \n    with open(f'{feature_name}/all_df_cache.pkl', 'wb') as handle:\n        pickle.dump(all_df, handle, protocol=pickle.HIGHEST_PROTOCOL)\n    with open(f'{feature_name}/envs_cache.pkl', 'wb') as handle:\n        pickle.dump(envs, handle, protocol=pickle.HIGHEST_PROTOCOL)\n    with open(f'{feature_name}/exp_names_cache.pkl', 'wb') as handle:\n        pickle.dump(exp_names, handle, protocol=pickle.HIGHEST_PROTOCOL)\nelse:\n    with open(f'{feature_name}/all_df_cache.pkl', 'rb') as handle:\n        all_df = pickle.load(handle)\n    with open(f'{feature_name}/envs_cache.pkl', 'rb') as handle:\n        envs = pickle.load(handle)\n    with open(f'{feature_name}/exp_names_cache.pkl', 'rb') as handle:\n        exp_names = pickle.load(handle)\nprint(\"data loaded\")\n\n# https://stackoverflow.com/questions/42281844/what-is-the-mathematics-behind-the-smoothing-parameter-in-tensorboards-scalar#_=_\ndef smooth(scalars, weight):  # Weight between 0 and 1\n    last = scalars[0]  # First value in the plot (first timestep)\n    smoothed = list()\n    for point in scalars:\n        smoothed_val = last * weight + (1 - weight) * point  # Calculate smoothed value\n        smoothed.append(smoothed_val)                        # Save it\n        last = smoothed_val                                  # Anchor the last smoothed value\n\n    return smoothed\n\n#smoothing\nfor env in envs:\n    if not env.endswith(\"total_timesteps\"):\n        for idx, metrics_dataframe in enumerate(envs[env]):\n            envs[env][idx] = metrics_dataframe.dropna(subset=[args.feature_of_interest])\n#             envs[env][idx][args.feature_of_interest] = smooth(metrics_dataframe[args.feature_of_interest], 0.85)\n\nsns.set(style=\"darkgrid\")\ndef get_df_for_env(gym_id):\n    env_total_timesteps = envs[gym_id+\"total_timesteps\"]\n    env_increment = env_total_timesteps / 500\n    envs_same_x_axis = []\n    for sampled_run in envs[gym_id]:\n        df = pd.DataFrame(columns=sampled_run.columns)\n        x_axis = [i*env_increment for i in range(500-2)]\n        current_row = 0\n        for timestep in x_axis:\n            while sampled_run.iloc[current_row][\"global_step\"] < timestep:\n                current_row += 1\n                if current_row > len(sampled_run)-2:\n                    break\n            if current_row > len(sampled_run)-2:\n                break\n            temp_row = sampled_run.iloc[current_row].copy()\n            temp_row[\"global_step\"] = timestep\n            df = df.append(temp_row)\n        \n        envs_same_x_axis += [df]\n    return pd.concat(envs_same_x_axis, ignore_index=True)\n\ndef export_legend(ax, filename=\"legend.pdf\"):\n    # import matplotlib as mpl\n    # mpl.rcParams['text.usetex'] = True\n    # mpl.rcParams['text.latex.preamble'] = [r'\\usepackage{amsmath}'] #for \\text command\n    fig2 = plt.figure()\n    ax2 = fig2.add_subplot()\n    ax2.axis('off')\n    handles, labels = ax.get_legend_handles_labels()\n\n    legend = ax2.legend(handles=handles, labels=labels, frameon=False, loc='lower center', ncol=4, fontsize=20, handlelength=1)\n    for text in legend.get_texts():\n        if text.get_text() in exp_convert_dict:\n            text.set_text(exp_convert_dict[text.get_text()])\n    for line in legend.get_lines():\n        line.set_linewidth(4.0)\n    fig  = legend.figure\n    fig.canvas.draw()\n    bbox  = legend.get_window_extent().transformed(fig.dpi_scale_trans.inverted())\n    fig.savefig(filename, dpi=\"figure\", bbox_inches=bbox)\n    fig.clf()\n\nif not os.path.exists(f\"{feature_name}/data\"):\n    os.makedirs(f\"{feature_name}/data\")\nif not os.path.exists(f\"{feature_name}/plots\"):\n    os.makedirs(f\"{feature_name}/plots\")\nif not os.path.exists(f\"{feature_name}/legends\"):\n    os.makedirs(f\"{feature_name}/legends\")\n\n\ninterested_exp_names = sorted(list(set(exp_names))) # ['ppo_continuous_action', 'ppo_atari_visual']\ncurrent_palette = sns.color_palette(n_colors=len(interested_exp_names))\ncurrent_palette_dict = dict(zip(interested_exp_names, current_palette))\nif args.interested_exp_names:\n    interested_exp_names = args.interested_exp_names\nprint(current_palette_dict)\nlegend_df = pd.DataFrame()\n\nif args.font_size:\n    plt.rc('axes', titlesize=args.font_size)     # fontsize of the axes title\n    plt.rc('axes', labelsize=args.font_size)    # fontsize of the x and y labels\n    plt.rc('xtick', labelsize=args.font_size)    # fontsize of the tick labels\n    plt.rc('ytick', labelsize=args.font_size)    # fontsize of the tick labels\n    plt.rc('legend', fontsize=args.font_size)    # legend fontsize\n\nstats = {item: [] for item in [\"gym_id\", \"exp_name\", args.feature_of_interest]}\n# uncommenet the following to generate all figures\nfor env in set(all_df[\"gym_id\"]):\n    if not path.exists(f\"{feature_name}/data/{env}.pkl\"):\n        with open(f\"{feature_name}/data/{env}.pkl\", 'wb') as handle:\n            data = get_df_for_env(env)\n            data[\"seed\"] = data[\"seed\"].astype(float)\n            data[args.feature_of_interest] = data[args.feature_of_interest].astype(float)\n            pickle.dump(data, handle, protocol=pickle.HIGHEST_PROTOCOL)\n    else:\n        with open(f\"{feature_name}/data/{env}.pkl\", 'rb') as handle:\n            data = pickle.load(handle)\n            print(f\"{env}'s data loaded\")\n    def _smooth(df):\n        df[args.feature_of_interest] = smooth(list(df[args.feature_of_interest]), args.smooth_weight)\n        return df\n\n    legend_df = legend_df.append(data)\n    ax = sns.lineplot(data=data.groupby([\"seed\", \"algo\"]).apply(_smooth).loc[data['algo'].isin(interested_exp_names)], x=\"global_step\", y=args.feature_of_interest, hue=\"algo\", ci='sd', palette=current_palette_dict,)\n    ax.ticklabel_format(style='sci', scilimits=(0,0), axis='x')\n    ax.set(xlabel=args.x_label, ylabel=args.y_label)\n    \n    # hack \n    ax.set(xlabel=\"\")\n    ax.set_xticks([])\n    \n    ax.legend().remove()\n    if args.y_lim_bottom:\n        plt.ylim(bottom=args.y_lim_bottom)\n    # plt.title(env)\n    plt.tight_layout()\n    plt.savefig(f\"{feature_name}/plots/{env}.{args.output_format}\")\n    plt.clf()\n    \n    for algo in interested_exp_names:\n        algo_data = data.loc[data['algo'].isin([algo])]\n        last_n_episodes_global_step = sorted(algo_data[\"global_step\"].unique())[-args.last_n_episodes]\n        last_n_episodes_features = algo_data[algo_data['global_step'] > last_n_episodes_global_step].groupby(\n            ['seed']\n        ).mean()[args.feature_of_interest]\n        \n        for item in last_n_episodes_features:\n            stats[args.feature_of_interest] += [item]\n            if algo in exp_convert_dict:\n                stats['exp_name'] += [exp_convert_dict[algo]]\n            else:\n                stats['exp_name'] += [algo]\n            stats['gym_id'] += [env]\n\n# export legend\nlegend_df = legend_df.reset_index()\nax = sns.lineplot(data=legend_df, x=\"global_step\", y=args.feature_of_interest, hue=\"algo\", ci='sd', palette=current_palette_dict,)\nax.set(xlabel='Time Steps', ylabel='Average Episode Reward')\nax.legend().remove()\nexport_legend(ax, f\"{feature_name}/legend.{args.output_format}\")\nplt.clf()\n\n\n# analysis\nstats_df = pd.DataFrame(stats)\ng = stats_df.groupby(\n    ['gym_id','exp_name']\n).agg(lambda x: f\"{np.mean(x):.2f} ± {np.std(x):.2f}\")\nprint(g.reset_index().pivot('exp_name', 'gym_id' , args.feature_of_interest).to_latex().replace(\"±\", \"$\\pm$\"))"
  },
  {
    "path": "ppo.py",
    "content": "import torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport torch.nn.functional as F\nfrom torch.distributions.categorical import Categorical\nfrom torch.utils.tensorboard import SummaryWriter\n\nimport argparse\nfrom distutils.util import strtobool\nimport numpy as np\nimport gym\nimport gym_microrts\nfrom gym.wrappers import TimeLimit, Monitor\n\nfrom gym.spaces import Discrete, Box, MultiBinary, MultiDiscrete, Space\nimport time\nimport random\nimport os\nfrom stable_baselines3.common.vec_env import DummyVecEnv, SubprocVecEnv, VecEnvWrapper\n\nif __name__ == \"__main__\":\n    parser = argparse.ArgumentParser(description='PPO agent')\n    # Common arguments\n    parser.add_argument('--exp-name', type=str, default=os.path.basename(__file__).rstrip(\".py\"),\n                        help='the name of this experiment')\n    parser.add_argument('--gym-id', type=str, default=\"MicrortsMining10x10F9-v0\",\n                        help='the id of the gym environment')\n    parser.add_argument('--learning-rate', type=float, default=2.5e-4,\n                        help='the learning rate of the optimizer')\n    parser.add_argument('--seed', type=int, default=1,\n                        help='seed of the experiment')\n    parser.add_argument('--total-timesteps', type=int, default=10000000,\n                        help='total timesteps of the experiments')\n    parser.add_argument('--torch-deterministic', type=lambda x:bool(strtobool(x)), default=True, nargs='?', const=True,\n                        help='if toggled, `torch.backends.cudnn.deterministic=False`')\n    parser.add_argument('--cuda', type=lambda x:bool(strtobool(x)), default=True, nargs='?', const=True,\n                        help='if toggled, cuda will not be enabled by default')\n    parser.add_argument('--prod-mode', type=lambda x:bool(strtobool(x)), default=False, nargs='?', const=True,\n                        help='run the script in production mode and use wandb to log outputs')\n    parser.add_argument('--capture-video', type=lambda x:bool(strtobool(x)), default=False, nargs='?', const=True,\n                        help='weather to capture videos of the agent performances (check out `videos` folder)')\n    parser.add_argument('--wandb-project-name', type=str, default=\"cleanRL\",\n                        help=\"the wandb's project name\")\n    parser.add_argument('--wandb-entity', type=str, default=None,\n                        help=\"the entity (team) of wandb's project\")\n\n    # Algorithm specific arguments\n    parser.add_argument('--n-minibatch', type=int, default=4,\n                        help='the number of mini batch')\n    parser.add_argument('--num-envs', type=int, default=8,\n                        help='the number of parallel game environment')\n    parser.add_argument('--num-steps', type=int, default=128,\n                        help='the number of steps per game environment')\n    parser.add_argument('--gamma', type=float, default=0.99,\n                        help='the discount factor gamma')\n    parser.add_argument('--gae-lambda', type=float, default=0.95,\n                        help='the lambda for the general advantage estimation')\n    parser.add_argument('--ent-coef', type=float, default=0.01,\n                        help=\"coefficient of the entropy\")\n    parser.add_argument('--vf-coef', type=float, default=0.5,\n                        help=\"coefficient of the value function\")\n    parser.add_argument('--max-grad-norm', type=float, default=0.5,\n                        help='the maximum norm for the gradient clipping')\n    parser.add_argument('--clip-coef', type=float, default=0.1,\n                        help=\"the surrogate clipping coefficient\")\n    parser.add_argument('--update-epochs', type=int, default=4,\n                         help=\"the K epochs to update the policy\")\n    parser.add_argument('--kle-stop', type=lambda x:bool(strtobool(x)), default=False, nargs='?', const=True,\n                         help='If toggled, the policy updates will be early stopped w.r.t target-kl')\n    parser.add_argument('--kle-rollback', type=lambda x:bool(strtobool(x)), default=False, nargs='?', const=True,\n                         help='If toggled, the policy updates will roll back to previous policy if KL exceeds target-kl')\n    parser.add_argument('--target-kl', type=float, default=0.03,\n                         help='the target-kl variable that is referred by --kl')\n    parser.add_argument('--gae', type=lambda x:bool(strtobool(x)), default=True, nargs='?', const=True,\n                         help='Use GAE for advantage computation')\n    parser.add_argument('--norm-adv', type=lambda x:bool(strtobool(x)), default=True, nargs='?', const=True,\n                          help=\"Toggles advantages normalization\")\n    parser.add_argument('--anneal-lr', type=lambda x:bool(strtobool(x)), default=True, nargs='?', const=True,\n                          help=\"Toggle learning rate annealing for policy and value networks\")\n    parser.add_argument('--clip-vloss', type=lambda x:bool(strtobool(x)), default=True, nargs='?', const=True,\n                          help='Toggles wheter or not to use a clipped loss for the value function, as per the paper.')\n\n    args = parser.parse_args()\n    if not args.seed:\n        args.seed = int(time.time())\n\nargs.batch_size = int(args.num_envs * args.num_steps)\nargs.minibatch_size = int(args.batch_size // args.n_minibatch)\n\nclass ImageToPyTorch(gym.ObservationWrapper):\n    def __init__(self, env):\n        super(ImageToPyTorch, self).__init__(env)\n        old_shape = self.observation_space.shape\n        self.observation_space = gym.spaces.Box(\n            low=0,\n            high=1,\n            shape=(old_shape[-1], old_shape[0], old_shape[1]),\n            dtype=np.int32,\n        )\n\n    def observation(self, observation):\n        return np.transpose(observation, axes=(2, 0, 1))\n\nclass VecPyTorch(VecEnvWrapper):\n    def __init__(self, venv, device):\n        super(VecPyTorch, self).__init__(venv)\n        self.device = device\n\n    def reset(self):\n        obs = self.venv.reset()\n        obs = torch.from_numpy(obs).float().to(self.device)\n        return obs\n\n    def step_async(self, actions):\n        actions = actions.cpu().numpy()\n        self.venv.step_async(actions)\n\n    def step_wait(self):\n        obs, reward, done, info = self.venv.step_wait()\n        obs = torch.from_numpy(obs).float().to(self.device)\n        reward = torch.from_numpy(reward).unsqueeze(dim=1).float()\n        return obs, reward, done, info\n\nclass MicroRTSStatsRecorder(gym.Wrapper):\n\n    def reset(self, **kwargs):\n        observation = super(MicroRTSStatsRecorder, self).reset(**kwargs)\n        self.raw_rewards = []\n        return observation\n\n    def step(self, action):\n        observation, reward, done, info = super(MicroRTSStatsRecorder, self).step(action)\n        self.raw_rewards += [info[\"raw_rewards\"]]\n        if done:\n            raw_rewards = np.array(self.raw_rewards).sum(0)\n            raw_names = [str(rf) for rf in self.rfs]\n            info['microrts_stats'] = dict(zip(raw_names, raw_rewards))\n            self.raw_rewards = []\n        return observation, reward, done, info\n\n\n# TRY NOT TO MODIFY: setup the environment\nexperiment_name = f\"{args.gym_id}__{args.exp_name}__{args.seed}__{int(time.time())}\"\nwriter = SummaryWriter(f\"runs/{experiment_name}\")\nwriter.add_text('hyperparameters', \"|param|value|\\n|-|-|\\n%s\" % (\n        '\\n'.join([f\"|{key}|{value}|\" for key, value in vars(args).items()])))\nif args.prod_mode:\n    import wandb\n    wandb.init(project=args.wandb_project_name, entity=args.wandb_entity, sync_tensorboard=True, config=vars(args), name=experiment_name, monitor_gym=True, save_code=True)\n    writer = SummaryWriter(f\"/tmp/{experiment_name}\")\n\n# TRY NOT TO MODIFY: seeding\ndevice = torch.device('cuda' if torch.cuda.is_available() and args.cuda else 'cpu')\nrandom.seed(args.seed)\nnp.random.seed(args.seed)\ntorch.manual_seed(args.seed)\ntorch.backends.cudnn.deterministic = args.torch_deterministic\ndef make_env(gym_id, seed, idx):\n    def thunk():\n        env = gym.make(gym_id)\n        env = ImageToPyTorch(env)\n        env = gym.wrappers.RecordEpisodeStatistics(env)\n        env = MicroRTSStatsRecorder(env)\n        if args.capture_video:\n            if idx == 0:\n                env = Monitor(env, f'videos/{experiment_name}')\n        env.seed(seed)\n        env.action_space.seed(seed)\n        env.observation_space.seed(seed)\n        return env\n    return thunk\nenvs = VecPyTorch(DummyVecEnv([make_env(args.gym_id, args.seed+i, i) for i in range(args.num_envs)]), device)\n# if args.prod_mode:\n#     envs = VecPyTorch(\n#         SubprocVecEnv([make_env(args.gym_id, args.seed+i, i) for i in range(args.num_envs)], \"fork\"),\n#         device\n#     )\nassert isinstance(envs.action_space, MultiDiscrete), \"only MultiDiscrete action space is supported\"\n\n# ALGO LOGIC: initialize agent here:\nclass CategoricalMasked(Categorical):\n    def __init__(self, probs=None, logits=None, validate_args=None, masks=[]):\n        self.masks = masks\n        if len(self.masks) == 0:\n            super(CategoricalMasked, self).__init__(probs, logits, validate_args)\n        else:\n            self.masks = masks.type(torch.BoolTensor).to(device)\n            logits = torch.where(self.masks, logits, torch.tensor(-1e+8).to(device))\n            super(CategoricalMasked, self).__init__(probs, logits, validate_args)\n    \n    def entropy(self):\n        if len(self.masks) == 0:\n            return super(CategoricalMasked, self).entropy()\n        p_log_p = self.logits * self.probs\n        p_log_p = torch.where(self.masks, p_log_p, torch.tensor(0.).to(device))\n        return -p_log_p.sum(-1)\n\nclass Scale(nn.Module):\n    def __init__(self, scale):\n        super().__init__()\n        self.scale = scale\n\n    def forward(self, x):\n        return x * self.scale\n\ndef 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\nclass Agent(nn.Module):\n    def __init__(self, frames=4):\n        super(Agent, self).__init__()\n        self.network = nn.Sequential(\n            layer_init(nn.Conv2d(27, 16, kernel_size=3, stride=2)),\n            nn.ReLU(),\n            layer_init(nn.Conv2d(16, 32, kernel_size=2)),\n            nn.ReLU(),\n            nn.Flatten(),\n            layer_init(nn.Linear(32*3*3, 128)),\n            nn.ReLU(),)\n        self.actor = layer_init(nn.Linear(128, envs.action_space.nvec.sum()), std=0.01)\n        self.critic = layer_init(nn.Linear(128, 1), std=1)\n\n    def forward(self, x):\n        return self.network(x)\n\n    def get_action(self, x, action=None, invalid_action_masks=None):\n        logits = self.actor(self.forward(x))\n        split_logits = torch.split(logits, envs.action_space.nvec.tolist(), dim=1)\n        \n        if invalid_action_masks is not None:\n            split_invalid_action_masks = torch.split(invalid_action_masks, envs.action_space.nvec.tolist(), dim=1)\n            multi_categoricals = [CategoricalMasked(logits=logits, masks=iam) for (logits, iam) in zip(split_logits, split_invalid_action_masks)]\n        else:\n            multi_categoricals = [Categorical(logits=logits) for logits in split_logits]\n        \n        if action is None:\n            action = torch.stack([categorical.sample() for categorical in multi_categoricals])\n        logprob = torch.stack([categorical.log_prob(a) for a, categorical in zip(action, multi_categoricals)])\n        entropy = torch.stack([categorical.entropy() for categorical in multi_categoricals])\n        return action, logprob.sum(0), entropy.sum(0)\n\n    def get_value(self, x):\n        return self.critic(self.forward(x))\n\nagent = Agent().to(device)\noptimizer = optim.Adam(agent.parameters(), lr=args.learning_rate, eps=1e-5)\nif args.anneal_lr:\n    # https://github.com/openai/baselines/blob/ea25b9e8b234e6ee1bca43083f8f3cf974143998/baselines/ppo2/defaults.py#L20\n    lr = lambda f: f * args.learning_rate\n\n# ALGO Logic: Storage for epoch data\nobs = torch.zeros((args.num_steps, args.num_envs) + envs.observation_space.shape).to(device)\nactions = torch.zeros((args.num_steps, args.num_envs) + envs.action_space.shape).to(device)\nlogprobs = torch.zeros((args.num_steps, args.num_envs)).to(device)\nrewards = torch.zeros((args.num_steps, args.num_envs)).to(device)\ndones = torch.zeros((args.num_steps, args.num_envs)).to(device)\nvalues = torch.zeros((args.num_steps, args.num_envs)).to(device)\ninvalid_action_masks = torch.zeros((args.num_steps, args.num_envs) + (envs.action_space.nvec.sum(),)).to(device)\n# TRY NOT TO MODIFY: start the game\nglobal_step = 0\n# Note how `next_obs` and `next_done` are used; their usage is equivalent to\n# https://github.com/ikostrikov/pytorch-a2c-ppo-acktr-gail/blob/84a7582477fb0d5c82ad6d850fe476829dddd2e1/a2c_ppo_acktr/storage.py#L60\nnext_obs = envs.reset()\nnext_done = torch.zeros(args.num_envs).to(device)\nnum_updates = args.total_timesteps // args.batch_size\nfor update in range(1, num_updates+1):\n    # Annealing the rate if instructed to do so.\n    if args.anneal_lr:\n        frac = 1.0 - (update - 1.0) / num_updates\n        lrnow = lr(frac)\n        optimizer.param_groups[0]['lr'] = lrnow\n\n    # TRY NOT TO MODIFY: prepare the execution of the game.\n    for step in range(0, args.num_steps):\n        envs.env_method(\"render\", indices=0)\n        global_step += 1 * args.num_envs\n        obs[step] = next_obs\n        dones[step] = next_done\n        invalid_action_masks[step] = torch.Tensor(np.array(envs.get_attr(\"action_mask\")))\n\n        # ALGO LOGIC: put action logic here\n        with torch.no_grad():\n            values[step] = agent.get_value(obs[step]).flatten()\n            action, logproba, _ = agent.get_action(obs[step], invalid_action_masks=invalid_action_masks[step])\n        \n        actions[step] = action.T\n        logprobs[step] = logproba\n\n        # TRY NOT TO MODIFY: execute the game and log data.\n        next_obs, rs, ds, infos = envs.step(action.T)\n        rewards[step], next_done = rs.view(-1), torch.Tensor(ds).to(device)\n\n        for info in infos:\n            if 'episode' in info.keys():\n                print(f\"global_step={global_step}, episode_reward={info['episode']['r']}\")\n                writer.add_scalar(\"charts/episode_reward\", info['episode']['r'], global_step)\n                for key in info['microrts_stats']:\n                    writer.add_scalar(f\"charts/episode_reward/{key}\", info['microrts_stats'][key], global_step)\n                break\n\n    # bootstrap reward if not done. reached the batch limit\n    with torch.no_grad():\n        last_value = agent.get_value(next_obs.to(device)).reshape(1, -1)\n        if args.gae:\n            advantages = torch.zeros_like(rewards).to(device)\n            lastgaelam = 0\n            for t in reversed(range(args.num_steps)):\n                if t == args.num_steps - 1:\n                    nextnonterminal = 1.0 - next_done\n                    nextvalues = last_value\n                else:\n                    nextnonterminal = 1.0 - dones[t+1]\n                    nextvalues = values[t+1]\n                delta = rewards[t] + args.gamma * nextvalues * nextnonterminal - values[t]\n                advantages[t] = lastgaelam = delta + args.gamma * args.gae_lambda * nextnonterminal * lastgaelam\n            returns = advantages + values\n        else:\n            returns = torch.zeros_like(rewards).to(device)\n            for t in reversed(range(args.num_steps)):\n                if t == args.num_steps - 1:\n                    nextnonterminal = 1.0 - next_done\n                    next_return = last_value\n                else:\n                    nextnonterminal = 1.0 - dones[t+1]\n                    next_return = returns[t+1]\n                returns[t] = rewards[t] + args.gamma * nextnonterminal * next_return\n            advantages = returns - values\n\n    # flatten the batch\n    b_obs = obs.reshape((-1,)+envs.observation_space.shape)\n    b_logprobs = logprobs.reshape(-1)\n    b_actions = actions.reshape((-1,)+envs.action_space.shape)\n    b_advantages = advantages.reshape(-1)\n    b_returns = returns.reshape(-1)\n    b_values = values.reshape(-1)\n    b_invalid_action_masks = invalid_action_masks.reshape((-1, invalid_action_masks.shape[-1]))\n\n    # Optimizaing the policy and value network\n    target_agent = Agent().to(device)\n    inds = np.arange(args.batch_size,)\n    for i_epoch_pi in range(args.update_epochs):\n        np.random.shuffle(inds)\n        target_agent.load_state_dict(agent.state_dict())\n        for start in range(0, args.batch_size, args.minibatch_size):\n            end = start + args.minibatch_size\n            minibatch_ind = inds[start:end]\n            mb_advantages = b_advantages[minibatch_ind]\n            if args.norm_adv:\n                mb_advantages = (mb_advantages - mb_advantages.mean()) / (mb_advantages.std() + 1e-8)\n\n            _, newlogproba, entropy = agent.get_action(\n                b_obs[minibatch_ind],\n                b_actions.long()[minibatch_ind].T,\n                b_invalid_action_masks[minibatch_ind])\n            ratio = (newlogproba - b_logprobs[minibatch_ind]).exp()\n\n            # Stats\n            approx_kl = (b_logprobs[minibatch_ind] - newlogproba).mean()\n\n            # Policy loss\n            pg_loss1 = -mb_advantages * ratio\n            pg_loss2 = -mb_advantages * torch.clamp(ratio, 1-args.clip_coef, 1+args.clip_coef)\n            pg_loss = torch.max(pg_loss1, pg_loss2).mean()\n            entropy_loss = entropy.mean()\n\n            # Value loss\n            new_values = agent.get_value(b_obs[minibatch_ind]).view(-1)\n            if args.clip_vloss:\n                v_loss_unclipped = ((new_values - b_returns[minibatch_ind]) ** 2)\n                v_clipped = b_values[minibatch_ind] + torch.clamp(new_values - b_values[minibatch_ind], -args.clip_coef, args.clip_coef)\n                v_loss_clipped = (v_clipped - b_returns[minibatch_ind])**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 *((new_values - b_returns[minibatch_ind]) ** 2)\n\n            loss = pg_loss - args.ent_coef * entropy_loss + v_loss * args.vf_coef\n\n            optimizer.zero_grad()\n            loss.backward()\n            nn.utils.clip_grad_norm_(agent.parameters(), args.max_grad_norm)\n            optimizer.step()\n\n        if args.kle_stop:\n            if approx_kl > args.target_kl:\n                break\n        if args.kle_rollback:\n            if (b_logprobs[minibatch_ind] - agent.get_action(\n                    b_obs[minibatch_ind],\n                    b_actions.long()[minibatch_ind].T,\n                    b_invalid_action_masks[minibatch_ind])[1]).mean() > args.target_kl:\n                agent.load_state_dict(target_agent.state_dict())\n                break\n\n    # TRY NOT TO MODIFY: record rewards for plotting purposes\n    writer.add_scalar(\"charts/learning_rate\", optimizer.param_groups[0]['lr'], global_step)\n    writer.add_scalar(\"losses/value_loss\", v_loss.item(), global_step)\n    writer.add_scalar(\"losses/policy_loss\", pg_loss.item(), global_step)\n    writer.add_scalar(\"losses/entropy\", entropy.mean().item(), global_step)\n    writer.add_scalar(\"losses/approx_kl\", approx_kl.item(), global_step)\n    if args.kle_stop or args.kle_rollback:\n        writer.add_scalar(\"debug/pg_stop_iter\", i_epoch_pi, global_step)\n\nenvs.close()\nwriter.close()"
  },
  {
    "path": "pyproject.toml",
    "content": "[tool.poetry]\nname = \"invalid-action-masking\"\nversion = \"0.1.0\"\ndescription = \"\"\nauthors = [\"Costa Huang <costa.huang@outlook.com>\"]\n\n[tool.poetry.dependencies]\npython = \"^3.8\"\ntorch = \"1.7.1\"\nPillow = \"^8.3.1\"\ncleanrl = {git = \"https://github.com/vwxyzjn/cleanrl.git\", rev = \"V0.1\"}\ngym-microrts = {git = \"https://github.com/vwxyzjn/gym-microrts.git\", rev = \"b0cabbabc363177709b3132201d57d024b5212e9\"}\ntensorboard = \"^2.5.0\"\npandas = \"^1.3.0\"\nstable-baselines3 = \"^1.1.0\"\nwandb = \"^0.12.2\"\nseaborn = \"^0.11.2\"\nspyder = \"^5.1.5\"\nsetuptools = \"59.5.0\"\n\n[tool.poetry.dev-dependencies]\n\n[build-system]\nrequires = [\"poetry-core>=1.0.0\"]\nbuild-backend = \"poetry.core.masonry.api\"\n"
  },
  {
    "path": "requirements.txt",
    "content": "absl-py==0.13.0\ncachetools==4.2.2\ncertifi==2021.5.30\ncharset-normalizer==2.0.1; python_version >= \"3\"\ncleanrl @ git+https://github.com/vwxyzjn/cleanrl.git@V0.1\ncloudpickle==1.6.0\ndacite==1.6.0\nfuture==0.18.2\ngoogle-auth-oauthlib==0.4.4\ngoogle-auth==1.32.1\ngrpcio==1.38.1\ngym-microrts @ git+https://github.com/vwxyzjn/gym-microrts.git@b0cabbabc363177709b3132201d57d024b5212e9\ngym==0.17.3\nhilbertcurve==2.0.5\nidna==3.2; python_version >= \"3\"\njpype1==1.3.0\nmarkdown==3.3.4\nnumpy==1.21.0\noauthlib==3.1.1\npandas==1.3.0\npillow==8.3.1\nprotobuf==3.17.3\npyasn1-modules==0.2.8\npyasn1==0.4.8\npyglet==1.5.0\npython-dateutil==2.8.2\npytz==2021.1\nrequests-oauthlib==1.3.0\nrequests==2.26.0\nrsa==4.7.2; python_version >= \"3.6\"\nscipy==1.6.1\nsix==1.16.0\ntensorboard-data-server==0.6.1\ntensorboard-plugin-wit==1.8.0\ntensorboard==2.5.0\ntorch==1.7.1\ntyping-extensions==3.10.0.0\nurllib3==1.26.6\nwerkzeug==2.0.1\n"
  },
  {
    "path": "test.py",
    "content": "# suppose action 1 is invalid\n\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport torch.nn.functional as F\nfrom torch.distributions.categorical import Categorical\naction = 0\nadvantage = torch.tensor(1.)\ndevice = \"cpu\"\n\n# no invalid action masking\nprint(\"=============regular=============\")\ntarget_logits = torch.tensor([1., 1., 1., 1.,] , requires_grad=True)\ntarget_probs = Categorical(logits=target_logits)\nlog_prob = target_probs.log_prob(torch.tensor(action))\nprint(\"log_prob\", log_prob)\n(log_prob*advantage).backward()\nprint(\"gradient\", target_logits.grad)\nprint()\n\n# invalid action masking via logits\nprint(\"==================invalid action masking=============\")\ntarget_logits = torch.tensor([1., 1., 1., 1.,] , requires_grad=True)\ninvalid_action_masks = torch.tensor([1., 1., 0., 1.,])\ninvalid_action_masks = invalid_action_masks.type(torch.BoolTensor)\nadjusted_logits = torch.where(invalid_action_masks, target_logits, torch.tensor(-1e+8))\nadjusted_probs = Categorical(logits=adjusted_logits)\nadjusted_log_prob = adjusted_probs.log_prob(torch.tensor(action))\nprint(\"log_prob\", adjusted_log_prob)\n(adjusted_log_prob*advantage).backward()\nprint(\"gradient\", target_logits.grad)\nprint()\n\n# invalid action masking via importance sampling\nprint(\"==================regular importance sampling=============\")\ntarget_logits = torch.tensor([1., 1., 1., 1.,] , requires_grad=True)\ntarget_probs = Categorical(logits=target_logits)\ninvalid_action_masks = torch.tensor([1., 1., 0., 1.,])\ninvalid_action_masks = invalid_action_masks.type(torch.BoolTensor)\nadjusted_logits = torch.where(invalid_action_masks, target_logits, torch.tensor(-1e+8))\nadjusted_probs = Categorical(logits=adjusted_logits)\nlog_prob = target_probs.log_prob(torch.tensor(action))\nadjusted_log_prob = adjusted_probs.log_prob(torch.tensor(action))\n\nimportance_sampling = target_probs.probs[torch.tensor(action)] / (adjusted_probs.probs[torch.tensor(action)])\nprint(\"log_prob\", log_prob)\n(importance_sampling.detach()*log_prob*advantage).backward()\nprint(\"gradient\", target_logits.grad)\nprint()\n\n\n# invalid action masking via logits\nprint(\"==================invalid action masking=============\")\ntarget_logits = torch.tensor([1., 1., 1., 1.,] , requires_grad=True)\ninvalid_action_masks = torch.tensor([1., 1., 0., 1.,])\ninvalid_action_masks = invalid_action_masks.type(torch.BoolTensor)\nadjusted_logits = torch.where(invalid_action_masks, target_logits, torch.tensor(-2.))\nadjusted_probs = Categorical(logits=adjusted_logits)\nadjusted_log_prob = adjusted_probs.log_prob(torch.tensor(action))\nprint(\"adjusted_probs\", adjusted_probs.probs)\n(adjusted_log_prob*advantage).backward()\nprint(\"gradient\", target_logits.grad)\nprint()\n\n# no invalid action masking with different parameterization\nprint(\"=============regular but differrent parameterization=============\")\ntarget_logits = torch.tensor([1., 1., -2., 1.,] , requires_grad=True)\ntarget_probs = Categorical(logits=target_logits)\nlog_prob = target_probs.log_prob(torch.tensor(action))\nprint(\"target_probs\", target_probs.probs)\n(log_prob*advantage).backward()\nprint(\"gradient\", target_logits.grad)\nprint()\nnew_target_logits = target_logits + target_logits.grad\nnew_target_probs = Categorical(logits=new_target_logits)\nprint(\"target_probs\", new_target_probs.probs)\n\n\n"
  }
]