main_REINFORCE_LDP.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "O-o9lHV2l3wt"
      },
      "source": [
        "# **REINFORCE with Local Differential Privacy  on Cartpole**\n",
        "## In this notebook, we propose a local differentially private REINFORCE by disturbing observations i.e. the states seen by the model. We add a gaussian noise to the states before feeding the REINFORCE."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eULrCg9cDk8-"
      },
      "source": [
        "## Acknowledgement\n",
        "\n",
        "Code adapted from [this notebook](https://github.com/bentrevett/pytorch-rl/blob/master/dqn_working.ipynb).\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nqYR93FHBUhP"
      },
      "source": [
        "## Import"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6Ux9Oyv_9_w5"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "import torch.nn.functional as F\n",
        "import torch.distributions as distributions\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import gym\n",
        "import os"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Asbsh2mh-Tz1"
      },
      "source": [
        "## Initialize environments"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NlRqUoaa9_w7"
      },
      "outputs": [],
      "source": [
        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "train_env = gym.make('CartPole-v1')\n",
        "test_env = gym.make('CartPole-v1')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ufSRMcJp9_w8"
      },
      "outputs": [],
      "source": [
        "MEAN = torch.zeros(4)\n",
        "STD = torch.tensor([0.55, 0.39, 0.05, 0.32])"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Seed"
      ],
      "metadata": {
        "id": "9yCbII_eEDui"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cXYQS2wk9_w9"
      },
      "outputs": [],
      "source": [
        "SEED = 42\n",
        "\n",
        "train_env.reset(seed=SEED);\n",
        "test_env.reset(seed=SEED+1);\n",
        "np.random.seed(SEED);\n",
        "torch.manual_seed(SEED);"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4fHfwQGk9_w-"
      },
      "outputs": [],
      "source": [
        "class MLP(nn.Module):\n",
        "    def __init__(self, input_dim, hidden_dim, output_dim, dropout = 0.5):\n",
        "        super().__init__()\n",
        "\n",
        "        self.fc_1 = nn.Linear(input_dim, hidden_dim)\n",
        "        self.fc_2 = nn.Linear(hidden_dim, output_dim)\n",
        "        self.dropout = nn.Dropout(dropout)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.fc_1(x)\n",
        "        x = self.dropout(x)\n",
        "        x = F.relu(x)\n",
        "        x = self.fc_2(x)\n",
        "        return x"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AwH0vkYE9_w_"
      },
      "outputs": [],
      "source": [
        "INPUT_DIM = train_env.observation_space.shape[0]\n",
        "HIDDEN_DIM = 128\n",
        "OUTPUT_DIM = train_env.action_space.n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9Xi9Y4ZQ9_xA"
      },
      "outputs": [],
      "source": [
        "def init_weights(m):\n",
        "    if type(m) == nn.Linear:\n",
        "        torch.nn.init.xavier_normal_(m.weight)\n",
        "        m.bias.data.fill_(0)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Training"
      ],
      "metadata": {
        "id": "tZ3c_Hg_EHSg"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6MUFe0719_xB"
      },
      "outputs": [],
      "source": [
        "def train(env, policy, optimizer, discount_factor, sigma=None):\n",
        "    \n",
        "    policy.train()\n",
        "    \n",
        "    log_prob_actions = []\n",
        "    rewards = []\n",
        "    done = False\n",
        "    episode_reward = 0\n",
        "    n_samples = 0\n",
        "    state, _ = env.reset()\n",
        "    \n",
        "    while not done:\n",
        "        state = torch.FloatTensor(state).unsqueeze(0)\n",
        "\n",
        "        # Add noise on the state\n",
        "        if sigma:\n",
        "            noise = torch.normal(mean=MEAN, std=sigma*STD)\n",
        "            state += noise\n",
        "\n",
        "        action_pred = policy(state)\n",
        "        \n",
        "        action_prob = F.softmax(action_pred, dim=-1)\n",
        "                \n",
        "        dist = distributions.Categorical(action_prob)\n",
        "\n",
        "        action = dist.sample()\n",
        "        \n",
        "        log_prob_action = dist.log_prob(action)\n",
        "        \n",
        "        state, reward, done, truncated, _ = env.step(action.item())\n",
        "        done = done or truncated\n",
        "\n",
        "        log_prob_actions.append(log_prob_action)\n",
        "        rewards.append(reward)\n",
        "        n_samples += 1\n",
        "        episode_reward += reward\n",
        "\n",
        "    log_prob_actions = torch.cat(log_prob_actions)\n",
        "        \n",
        "    returns = calculate_returns(rewards, discount_factor)\n",
        "        \n",
        "    loss = update_policy(returns, log_prob_actions, optimizer)\n",
        "\n",
        "    return loss, episode_reward, n_samples"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oq6R89jH9_xC"
      },
      "outputs": [],
      "source": [
        "def calculate_returns(rewards, discount_factor, normalize = True):\n",
        "    \n",
        "    returns = []\n",
        "    R = 0\n",
        "    \n",
        "    for r in reversed(rewards):\n",
        "        R = r + R * discount_factor\n",
        "        returns.insert(0, R)\n",
        "        \n",
        "    returns = torch.tensor(returns)\n",
        "    \n",
        "    if normalize:\n",
        "        returns = (returns - returns.mean()) / returns.std()\n",
        "        \n",
        "    return returns"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ipIbel3v9_xE"
      },
      "outputs": [],
      "source": [
        "def update_policy(returns, log_prob_actions, optimizer):\n",
        "    \n",
        "    returns = returns.detach()\n",
        "    \n",
        "    loss = - (returns * log_prob_actions).sum()\n",
        "    \n",
        "    optimizer.zero_grad()\n",
        "    \n",
        "    loss.backward()\n",
        "    \n",
        "    optimizer.step()\n",
        "    \n",
        "    return loss.item()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TwQwDiSv9_xF"
      },
      "outputs": [],
      "source": [
        "def evaluate(env, policy, sigma=None):\n",
        "    \n",
        "    policy.eval()\n",
        "    \n",
        "    done = False\n",
        "    episode_reward = 0\n",
        "\n",
        "    state, _ = env.reset()\n",
        "\n",
        "    while not done:     \n",
        "        state = torch.FloatTensor(state).unsqueeze(0)\n",
        "        \n",
        "        with torch.no_grad():\n",
        "        \n",
        "            action_pred = policy(state)\n",
        "        \n",
        "            action_prob = F.softmax(action_pred, dim = -1)\n",
        "                            \n",
        "        action = torch.argmax(action_prob, dim = -1)\n",
        "       \n",
        "        state, reward, done, truncated, _ = env.step(action.item())\n",
        "        done = done or truncated\n",
        "\n",
        "        episode_reward += reward\n",
        "        \n",
        "    return episode_reward"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "scrolled": false,
        "id": "C5GEOE6c9_xG"
      },
      "outputs": [],
      "source": [
        "total_nb_episodes = 1000\n",
        "gamma = 0.99\n",
        "\n",
        "lr = 0.01\n",
        "n_trials = 25\n",
        "\n",
        "noise_levels = [0, 0.01, 0.1, 0.2, 0.3, 0.5, 2, 5] \n",
        "\n",
        "print_every = 25\n",
        "\n",
        "run_0 = 100\n",
        "nb_runs = 100\n",
        "\n",
        "os.makedirs('experiments/reinforce_LDP', exist_ok=True)\n",
        "\n",
        "for run in range(nb_runs):\n",
        "    print('Starting run {}: '.format(run))\n",
        "    train_rewards = np.zeros((len(noise_levels), total_nb_episodes))\n",
        "    test_rewards = np.zeros((len(noise_levels), total_nb_episodes))\n",
        "\n",
        "    for i, sigma in enumerate(noise_levels):\n",
        "        train_env.reset(seed=SEED);\n",
        "        test_env.reset(seed=SEED+1);\n",
        "\n",
        "        print(\"-------------- Noise: sigma = {} ---------------\\n\".format(sigma))\n",
        "        policy = MLP(INPUT_DIM, HIDDEN_DIM, OUTPUT_DIM)\n",
        "        policy.apply(init_weights)\n",
        "\n",
        "        optimizer = optim.Adam(policy.parameters(), lr=lr)\n",
        "        convergence = False\n",
        "\n",
        "        for episode in range(total_nb_episodes):\n",
        "            if not convergence:\n",
        "                loss, train_reward, _ = train(train_env, policy,\n",
        "                                                      discount_factor = gamma, optimizer=optimizer, sigma=sigma)\n",
        "                test_reward = evaluate(test_env, policy, sigma=sigma)\n",
        "\n",
        "                train_rewards[i][episode] = train_reward\n",
        "                test_rewards[i][episode] = test_reward\n",
        "\n",
        "                mean_train_rewards = np.mean(train_rewards[i][episode-n_trials:episode+1])\n",
        "                mean_test_rewards = np.mean(test_rewards[i][episode-n_trials:episode+1])\n",
        "\n",
        "                convergence = (mean_test_rewards) >= 475\n",
        "                if convergence: print(\"Achieved convergence\")\n",
        "            else:\n",
        "                train_rewards[i][episode] = train_rewards[i][episode-1]\n",
        "                test_rewards[i][episode] = test_rewards[i][episode-1]\n",
        "\n",
        "            if episode % print_every == 0:\n",
        "                print(f'| Episode: {episode:3} | Train Reward: {train_reward:5.1f} | Test Reward: {test_reward:5.1f} |')\n",
        "        \n",
        "    np.save('experiments/reinforce_LDP/rewards_run_{}.npy'.format(run_0 + run), test_rewards)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1I8aXam19_xI"
      },
      "outputs": [],
      "source": [
        "fig = plt.figure(figsize=(12,8))\n",
        "\n",
        "for i, sigma in enumerate(noise_levels):\n",
        "    plt.plot(test_rewards[i], label='sigma = {}'.format(sigma))\n",
        "\n",
        "plt.xlabel('Number of episodes trained')\n",
        "plt.ylabel('Reward')\n",
        "\n",
        "plt.legend()\n",
        "plt.legend(loc='lower right')\n",
        "plt.grid()\n",
        "plt.show()"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.9.13"
    },
    "colab": {
      "provenance": []
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}