rm_dataset_filtering.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/sae-rm/logan/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n",
      "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]/root/sae-rm/logan/lib/python3.10/site-packages/torch/_utils.py:831: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.  To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage()\n",
      "  return self.fget.__get__(instance, owner)()\n",
      "Loading checkpoint shards: 100%|██████████| 3/3 [00:17<00:00,  5.79s/it]\n"
     ]
    }
   ],
   "source": [
    "from transformers import AutoModelForSequenceClassification, AutoTokenizer\n",
    "import torch\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "model_name = \"reciprocate/dahoas-gptj-rm-static\"\n",
    "model = AutoModelForSequenceClassification.from_pretrained(model_name, torch_dtype=torch.bfloat16).to(device)\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# import sys\n",
    "# sys.path.append(\"/root/dictionary_learning/\")\n",
    "from dictionary import GatedAutoEncoder\n",
    "\n",
    "layer = 2\n",
    "activation_name = f\"transformer.h.{layer}\"\n",
    "sae_file = f\"saes/ae_layer{layer}.pt\"\n",
    "ae = GatedAutoEncoder.from_pretrained(sae_file).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "from torch.utils.data import DataLoader\n",
    "from tqdm import tqdm\n",
    "import os\n",
    "import torch \n",
    "hh = load_dataset(\"Anthropic/hh-rlhf\", split=\"train\")\n",
    "token_length_cutoff = 870 # 99% of chosen data\n",
    "\n",
    "# Remove datapoints longer than a specific token_length\n",
    "# Check if file exists\n",
    "index_file_name = \"rm_save_files/index_small_enough.pt\"\n",
    "dataset_size = hh.num_rows\n",
    "if os.path.exists(index_file_name):\n",
    "    index_small_enough = torch.load(index_file_name)\n",
    "else:\n",
    "    print(\"hey\")\n",
    "#     index_small_enough = torch.ones(dataset_size, dtype=torch.bool)\n",
    "# # \n",
    "#     for ind, text in enumerate(tqdm(hh)):\n",
    "#         chosen_text = text[\"chosen\"]\n",
    "#         rejected_text = text[\"rejected\"]\n",
    "#         #convert to tokens\n",
    "#         length_chosen = len(tokenizer(chosen_text)[\"input_ids\"])\n",
    "#         length_rejected = len(tokenizer(rejected_text)[\"input_ids\"])\n",
    "#         if length_chosen > token_length_cutoff or length_rejected > token_length_cutoff:\n",
    "#             index_small_enough[ind] = False\n",
    "#     # Save the indices\n",
    "#     torch.save(index_small_enough, \"rm_save_files/index_small_enough.pt\")\n",
    "\n",
    "hh = hh.select(index_small_enough.nonzero()[:, 0])\n",
    "batch_size = 16\n",
    "hh_dl = DataLoader(hh, batch_size=batch_size, shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([158866]), 158866)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "reward_diff = torch.load(\"/root/sae-rm/rm_save_files/rejected_chosen_reward_diff.pt\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([160800])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "index_small_enough.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAjkAAAGdCAYAAADwjmIIAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjkuMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy80BEi2AAAACXBIWXMAAA9hAAAPYQGoP6dpAAAkWUlEQVR4nO3df3AU9f3H8VcC5Ag/7kKgueNqwLRaIZViBQnnD/qDDKdGZ6j0BzVVaqO09mIbgkqYShSLDQ1FEUVSbEeYKVRgplAlFc0ES1oJEWOpgJLqFBpaehfbkDuhkoRkv3/4zQ4HUQNeuNwnz8fMznC779t7LzvMvfjc7meTLMuyBAAAYJjkeDcAAADQGwg5AADASIQcAABgJEIOAAAwEiEHAAAYiZADAACMRMgBAABGIuQAAAAjDYx3A/HU2dmpo0ePavjw4UpKSop3OwAAoAcsy9J7770nr9er5OQPH6/p1yHn6NGjyszMjHcbAADgPBw5ckQXXXTRh27v1yFn+PDhkj74S3I6nXHuBgAA9EQkElFmZqb9Pf5h+nXI6fqJyul0EnIAAEgwH3epCRceAwAAIxFyAACAkQg5AADASIQcAABgJEIOAAAwEiEHAAAY6ZxDTk1NjW6++WZ5vV4lJSVp69atUdsty1JpaalGjx6t1NRU5ebm6u23346qaW5uVn5+vpxOp9LS0lRQUKDjx49H1bzxxhu67rrrNHjwYGVmZqq8vPysXjZv3qxx48Zp8ODBmjBhgv7whz+c6+EAAABDnXPIOXHihCZOnKhVq1Z1u728vFwrV65URUWF6urqNHToUPn9fp08edKuyc/P14EDB1RVVaVt27appqZGc+fOtbdHIhHNmDFDY8eOVX19vZYtW6aHHnpIa9assWt27dqlb3/72yooKNBf/vIXzZw5UzNnztT+/fvP9ZAAAICJrE9AkrVlyxb7dWdnp+XxeKxly5bZ61paWiyHw2H99re/tSzLst58801LkrVnzx675oUXXrCSkpKsf/3rX5ZlWdZTTz1ljRgxwmptbbVrFixYYF122WX2629+85tWXl5eVD85OTnW97///R73Hw6HLUlWOBzu8XsAAEB89fT7O6bX5Bw6dEjBYFC5ubn2OpfLpZycHNXW1kqSamtrlZaWpsmTJ9s1ubm5Sk5OVl1dnV0zbdo0paSk2DV+v18NDQ06duyYXXP653TVdH0OAADo32L6WIdgMChJcrvdUevdbre9LRgMKiMjI7qJgQOVnp4eVZOVlXXWPrq2jRgxQsFg8CM/pzutra1qbW21X0cikXM5PAAAkED61d1VZWVlcrlc9sITyAEAMFdMQ47H45EkhUKhqPWhUMje5vF41NTUFLX91KlTam5ujqrpbh+nf8aH1XRt787ChQsVDoft5ciRI+d6iAAAIEHENORkZWXJ4/GourraXheJRFRXVyefzydJ8vl8amlpUX19vV2zY8cOdXZ2Kicnx66pqalRe3u7XVNVVaXLLrtMI0aMsGtO/5yumq7P6Y7D4bCfOM6TxwEAMNs5X5Nz/PhxvfPOO/brQ4cOae/evUpPT9eYMWNUVFSkJUuW6NJLL1VWVpYWLVokr9ermTNnSpLGjx+v66+/XnfddZcqKirU3t6uwsJCzZ49W16vV5J06623avHixSooKNCCBQu0f/9+Pf7443rsscfsz/3xj3+sL33pS1q+fLny8vL07LPP6rXXXou6zRzAhXVxSWXU68NL8+LUCQDo3G8hf/nlly1JZy1z5syxLOuD28gXLVpkud1uy+FwWNOnT7caGhqi9vHf//7X+va3v20NGzbMcjqd1h133GG99957UTV//etfrWuvvdZyOBzWpz/9aWvp0qVn9bJp0ybrc5/7nJWSkmJ9/vOftyorK8/pWLiFHIitsQu2RS0A0Bt6+v2dZFmWFceMFVeRSEQul0vhcJifroAYYCQHwIXQ0+/vfnV3FQAA6D8IOQAAwEiEHAAAYCRCDgAAMBIhBwAAGImQAwAAjETIAQAARiLkAAAAIxFyAACAkQg5AADASIQcAABgJEIOAAAwEiEHAAAYiZADAACMRMgBAABGIuQAAAAjEXIAAICRCDkAAMBIhBwAAGCkgfFuAIC5Li6pPGvd4aV5cegEQH/ESA4AADASIQcAABiJkAMAAIxEyAEAAEYi5AAAACMRcgAAgJEIOQAAwEiEHAAAYCRCDgAAMBIhBwAAGInHOgC4oM581AOPeQDQWxjJAQAARiLkAAAAIxFyAACAkQg5AADASIQcAABgJEIOAAAwEiEHAAAYiZADAACMRMgBAABGIuQAAAAjEXIAAICRCDkAAMBIhBwAAGAkQg4AADASIQcAABhpYLwbANC/XVxSeda6w0vz4tAJANMwkgMAAIxEyAEAAEYi5AAAACMRcgAAgJEIOQAAwEiEHAAAYCRCDgAAMBIhBwAAGInJAAGcl+4m8QOAvoSRHAAAYCRCDgAAMBIhBwAAGImQAwAAjETIAQAARop5yOno6NCiRYuUlZWl1NRUffazn9VPf/pTWZZl11iWpdLSUo0ePVqpqanKzc3V22+/HbWf5uZm5efny+l0Ki0tTQUFBTp+/HhUzRtvvKHrrrtOgwcPVmZmpsrLy2N9OAAAIEHFPOT8/Oc/1+rVq/Xkk0/qrbfe0s9//nOVl5friSeesGvKy8u1cuVKVVRUqK6uTkOHDpXf79fJkyftmvz8fB04cEBVVVXatm2bampqNHfuXHt7JBLRjBkzNHbsWNXX12vZsmV66KGHtGbNmlgfEgAASEBJ1ulDLDFw0003ye1269e//rW9btasWUpNTdVvfvMbWZYlr9er+fPn695775UkhcNhud1urV27VrNnz9Zbb72l7Oxs7dmzR5MnT5Ykbd++XTfeeKP++c9/yuv1avXq1frJT36iYDColJQUSVJJSYm2bt2qgwcP9qjXSCQil8ulcDgsp9MZy78GwHi9OU/O4aV5vbZvAImvp9/fMR/Jufrqq1VdXa2//e1vkqS//vWv+vOf/6wbbrhBknTo0CEFg0Hl5uba73G5XMrJyVFtba0kqba2VmlpaXbAkaTc3FwlJyerrq7Orpk2bZodcCTJ7/eroaFBx44d67a31tZWRSKRqAUAAJgp5jMel5SUKBKJaNy4cRowYIA6Ojr0yCOPKD8/X5IUDAYlSW63O+p9brfb3hYMBpWRkRHd6MCBSk9Pj6rJyso6ax9d20aMGHFWb2VlZVq8eHEMjhIAAPR1MR/J2bRpk9avX68NGzbo9ddf17p16/SLX/xC69ati/VHnbOFCxcqHA7by5EjR+LdEgAA6CUxH8m57777VFJSotmzZ0uSJkyYoH/84x8qKyvTnDlz5PF4JEmhUEijR4+23xcKhXTFFVdIkjwej5qamqL2e+rUKTU3N9vv93g8CoVCUTVdr7tqzuRwOORwOD75QQIAgD4v5iM5//vf/5ScHL3bAQMGqLOzU5KUlZUlj8ej6upqe3skElFdXZ18Pp8kyefzqaWlRfX19XbNjh071NnZqZycHLumpqZG7e3tdk1VVZUuu+yybn+qAgAA/UvMQ87NN9+sRx55RJWVlTp8+LC2bNmiRx99VF/72tckSUlJSSoqKtKSJUv03HPPad++fbr99tvl9Xo1c+ZMSdL48eN1/fXX66677tKrr76qV155RYWFhZo9e7a8Xq8k6dZbb1VKSooKCgp04MABbdy4UY8//riKi4tjfUgAACABxfznqieeeEKLFi3SD3/4QzU1Ncnr9er73/++SktL7Zr7779fJ06c0Ny5c9XS0qJrr71W27dv1+DBg+2a9evXq7CwUNOnT1dycrJmzZqllStX2ttdLpdeeuklBQIBTZo0SaNGjVJpaWnUXDoAAKD/ivk8OYmEeXKA88c8OQDiJW7z5AAAAPQFhBwAAGCkmF+TAwCf1Jk/hfHzFYDzwUgOAAAwEiEHAAAYiZADAACMRMgBAABGIuQAAAAjEXIAAICRCDkAAMBIhBwAAGAkQg4AADASIQcAABiJkAMAAIxEyAEAAEYi5AAAACMRcgAAgJEIOQAAwEiEHAAAYCRCDgAAMBIhBwAAGImQAwAAjDQw3g0ASAwXl1TGuwUAOCeM5AAAACMRcgAAgJH4uQpAn9fdT2WHl+bFoRMAiYSRHAAAYCRCDgAAMBIhBwAAGImQAwAAjETIAQAARiLkAAAAIxFyAACAkQg5AADASIQcAABgJEIOAAAwEiEHAAAYiZADAACMRMgBAABGIuQAAAAjEXIAAICRCDkAAMBIhBwAAGAkQg4AADASIQcAABhpYLwbAIDzcXFJZdTrw0vz4tQJgL6KkRwAAGAkQg4AADASIQcAABiJkAMAAIxEyAEAAEYi5AAAACMRcgAAgJEIOQAAwEiEHAAAYCRCDgAAMBIhBwAAGIlnVwE4y5nPhQKARMRIDgAAMBIhBwAAGKlXQs6//vUvfec739HIkSOVmpqqCRMm6LXXXrO3W5al0tJSjR49WqmpqcrNzdXbb78dtY/m5mbl5+fL6XQqLS1NBQUFOn78eFTNG2+8oeuuu06DBw9WZmamysvLe+NwAABAAop5yDl27JiuueYaDRo0SC+88ILefPNNLV++XCNGjLBrysvLtXLlSlVUVKiurk5Dhw6V3+/XyZMn7Zr8/HwdOHBAVVVV2rZtm2pqajR37lx7eyQS0YwZMzR27FjV19dr2bJleuihh7RmzZpYHxIAAEhASZZlWbHcYUlJiV555RX96U9/6na7ZVnyer2aP3++7r33XklSOByW2+3W2rVrNXv2bL311lvKzs7Wnj17NHnyZEnS9u3bdeONN+qf//ynvF6vVq9erZ/85CcKBoNKSUmxP3vr1q06ePBgj3qNRCJyuVwKh8NyOp0xOHrADIl44fHhpXnxbgHABdLT7++Yj+Q899xzmjx5sr7xjW8oIyNDX/ziF/X000/b2w8dOqRgMKjc3Fx7ncvlUk5OjmprayVJtbW1SktLswOOJOXm5io5OVl1dXV2zbRp0+yAI0l+v18NDQ06duxYt721trYqEolELQAAwEwxDzl///vftXr1al166aV68cUXdffdd+tHP/qR1q1bJ0kKBoOSJLfbHfU+t9ttbwsGg8rIyIjaPnDgQKWnp0fVdLeP0z/jTGVlZXK5XPaSmZn5CY8WAAD0VTEPOZ2dnbryyiv1s5/9TF/84hc1d+5c3XXXXaqoqIj1R52zhQsXKhwO28uRI0fi3RIAAOglMQ85o0ePVnZ2dtS68ePHq7GxUZLk8XgkSaFQKKomFArZ2zwej5qamqK2nzp1Ss3NzVE13e3j9M84k8PhkNPpjFoAAICZYh5yrrnmGjU0NESt+9vf/qaxY8dKkrKysuTxeFRdXW1vj0Qiqqurk8/nkyT5fD61tLSovr7ertmxY4c6OzuVk5Nj19TU1Ki9vd2uqaqq0mWXXRZ1JxcAAOifYh5y5s2bp927d+tnP/uZ3nnnHW3YsEFr1qxRIBCQJCUlJamoqEhLlizRc889p3379un222+X1+vVzJkzJX0w8nP99dfrrrvu0quvvqpXXnlFhYWFmj17trxeryTp1ltvVUpKigoKCnTgwAFt3LhRjz/+uIqLi2N9SAAAIAHF/NlVV111lbZs2aKFCxfq4YcfVlZWllasWKH8/Hy75v7779eJEyc0d+5ctbS06Nprr9X27ds1ePBgu2b9+vUqLCzU9OnTlZycrFmzZmnlypX2dpfLpZdeekmBQECTJk3SqFGjVFpaGjWXDgAA6L9iPk9OImGeHKB7zJMDoC+L2zw5AAAAfUHMf64CgHjobvSJ0R2gf2MkBwAAGImQAwAAjETIAQAARiLkAAAAIxFyAACAkQg5AADASIQcAABgJEIOAAAwEiEHAAAYiRmPgX4uEZ9TBQA9wUgOAAAwEiEHAAAYiZADAACMRMgBAABGIuQAAAAjEXIAAICRCDkAAMBIhBwAAGAkQg4AADASIQcAABiJkAMAAIzEs6sAGOvM53IdXpoXp04AxAMjOQAAwEiEHAAAYCRCDgAAMBIhBwAAGImQAwAAjETIAQAARiLkAAAAIxFyAACAkQg5AADASIQcAABgJEIOAAAwEiEHAAAYiZADAACMRMgBAABGIuQAAAAjEXIAAICRCDkAAMBIhBwAAGAkQg4AADDSwHg3AODCurikMt4txE13x354aV4cOgFwITCSAwAAjETIAQAARiLkAAAAIxFyAACAkQg5AADASIQcAABgJEIOAAAwEiEHAAAYiZADAACMRMgBAABGIuQAAAAjEXIAAICRCDkAAMBIhBwAAGAkQg4AADASIQcAABip10PO0qVLlZSUpKKiInvdyZMnFQgENHLkSA0bNkyzZs1SKBSKel9jY6Py8vI0ZMgQZWRk6L777tOpU6eiav74xz/qyiuvlMPh0CWXXKK1a9f29uEAAIAE0ashZ8+ePfrlL3+pL3zhC1Hr582bp+eff16bN2/Wzp07dfToUd1yyy329o6ODuXl5amtrU27du3SunXrtHbtWpWWlto1hw4dUl5enr7yla9o7969Kioq0p133qkXX3yxNw8JAAAkiF4LOcePH1d+fr6efvppjRgxwl4fDof161//Wo8++qi++tWvatKkSXrmmWe0a9cu7d69W5L00ksv6c0339RvfvMbXXHFFbrhhhv005/+VKtWrVJbW5skqaKiQllZWVq+fLnGjx+vwsJCff3rX9djjz3WW4cEAAASSK+FnEAgoLy8POXm5katr6+vV3t7e9T6cePGacyYMaqtrZUk1dbWasKECXK73XaN3+9XJBLRgQMH7Joz9+33++19dKe1tVWRSCRqAQAAZhrYGzt99tln9frrr2vPnj1nbQsGg0pJSVFaWlrUerfbrWAwaNecHnC6tndt+6iaSCSi999/X6mpqWd9dllZmRYvXnzexwXAPBeXVEa9Prw0L06dAIi1mI/kHDlyRD/+8Y+1fv16DR48ONa7/0QWLlyocDhsL0eOHIl3SwAAoJfEPOTU19erqalJV155pQYOHKiBAwdq586dWrlypQYOHCi32622tja1tLREvS8UCsnj8UiSPB7PWXdbdb3+uBqn09ntKI4kORwOOZ3OqAUAAJgp5iFn+vTp2rdvn/bu3WsvkydPVn5+vv3nQYMGqbq62n5PQ0ODGhsb5fP5JEk+n0/79u1TU1OTXVNVVSWn06ns7Gy75vR9dNV07QMAAPRvMb8mZ/jw4br88suj1g0dOlQjR4601xcUFKi4uFjp6elyOp2655575PP5NHXqVEnSjBkzlJ2drdtuu03l5eUKBoN64IEHFAgE5HA4JEk/+MEP9OSTT+r+++/X9773Pe3YsUObNm1SZWX07+tAf3bm9SYA0J/0yoXHH+exxx5TcnKyZs2apdbWVvn9fj311FP29gEDBmjbtm26++675fP5NHToUM2ZM0cPP/ywXZOVlaXKykrNmzdPjz/+uC666CL96le/kt/vj8chAQCAPibJsiwr3k3ESyQSkcvlUjgc5vocGImRnHPH3VVA39fT72+eXQUAAIxEyAEAAEYi5AAAACMRcgAAgJEIOQAAwEiEHAAAYCRCDgAAMBIhBwAAGImQAwAAjETIAQAARiLkAAAAIxFyAACAkQg5AADASAPj3QAA9CXdPbmdJ5MDiYmRHAAAYCRGcgCDdDcKAQD9FSM5AADASIQcAABgJEIOAAAwEiEHAAAYiZADAACMRMgBAABGIuQAAAAjEXIAAICRCDkAAMBIhBwAAGAkQg4AADASIQcAABiJB3QCwMc488Gnh5fmxakTAOeCkRwAAGAkQg4AADASIQcAABiJa3KABHXmdSIAgGiM5AAAACMRcgAAgJEIOQAAwEiEHAAAYCRCDgAAMBJ3VwHAOeruzjZmQQb6HkZyAACAkQg5AADASIQcAABgJEIOAAAwEiEHAAAYiZADAACMRMgBAABGIuQAAAAjEXIAAICRmPEYSBDdzbILAPhwjOQAAAAjEXIAAICR+LkKAGLgzJ8TeWAnEH+M5AAAACMRcgAAgJEIOQAAwEiEHAAAYCRCDgAAMBIhBwAAGIlbyAGgF3Q3QzW3lQMXFiM5AADASDEPOWVlZbrqqqs0fPhwZWRkaObMmWpoaIiqOXnypAKBgEaOHKlhw4Zp1qxZCoVCUTWNjY3Ky8vTkCFDlJGRofvuu0+nTp2KqvnjH/+oK6+8Ug6HQ5dcconWrl0b68MBAAAJKuYhZ+fOnQoEAtq9e7eqqqrU3t6uGTNm6MSJE3bNvHnz9Pzzz2vz5s3auXOnjh49qltuucXe3tHRoby8PLW1tWnXrl1at26d1q5dq9LSUrvm0KFDysvL01e+8hXt3btXRUVFuvPOO/Xiiy/G+pAAAEACSrIsy+rND3j33XeVkZGhnTt3atq0aQqHw/rUpz6lDRs26Otf/7ok6eDBgxo/frxqa2s1depUvfDCC7rpppt09OhRud1uSVJFRYUWLFigd999VykpKVqwYIEqKyu1f/9++7Nmz56tlpYWbd++vUe9RSIRuVwuhcNhOZ3O2B88EEM8hTzxcU0OEBs9/f7u9WtywuGwJCk9PV2SVF9fr/b2duXm5to148aN05gxY1RbWytJqq2t1YQJE+yAI0l+v1+RSEQHDhywa07fR1dN1z6609raqkgkErUAAAAz9WrI6ezsVFFRka655hpdfvnlkqRgMKiUlBSlpaVF1brdbgWDQbvm9IDTtb1r20fVRCIRvf/++932U1ZWJpfLZS+ZmZmf+BgBAEDf1KshJxAIaP/+/Xr22Wd782N6bOHChQqHw/Zy5MiReLcEAAB6Sa/Nk1NYWKht27appqZGF110kb3e4/Gora1NLS0tUaM5oVBIHo/Hrnn11Vej9td199XpNWfekRUKheR0OpWamtptTw6HQw6H4xMfG9DbuP4GAD65mI/kWJalwsJCbdmyRTt27FBWVlbU9kmTJmnQoEGqrq621zU0NKixsVE+n0+S5PP5tG/fPjU1Ndk1VVVVcjqdys7OtmtO30dXTdc+AABA/xbzkZxAIKANGzbo97//vYYPH25fQ+NyuZSamiqXy6WCggIVFxcrPT1dTqdT99xzj3w+n6ZOnSpJmjFjhrKzs3XbbbepvLxcwWBQDzzwgAKBgD0S84Mf/EBPPvmk7r//fn3ve9/Tjh07tGnTJlVW8j9gAADQCyM5q1evVjgc1pe//GWNHj3aXjZu3GjXPPbYY7rppps0a9YsTZs2TR6PR7/73e/s7QMGDNC2bds0YMAA+Xw+fec739Htt9+uhx9+2K7JyspSZWWlqqqqNHHiRC1fvly/+tWv5Pf7Y31IAAAgAfX6PDl9GfPkoK/impz+gXlzgPPTZ+bJAQAAiAdCDgAAMFKv3UIOoOf4eQoAYo+RHAAAYCRCDgAAMBIhBwAAGImQAwAAjETIAQAARiLkAAAAI3ELOQDESXdTBzALMhA7jOQAAAAjEXIAAICRCDkAAMBIXJMDXGA8wgEALgxGcgAAgJEIOQAAwEiEHAAAYCSuyQGAPuTMa7aYNwc4f4zkAAAAIxFyAACAkQg5AADASIQcAABgJEIOAAAwEndXAb2MGY4BID4YyQEAAEZiJAcA+rDuRgKZOwfoGUZyAACAkQg5AADASIQcAABgJK7JAWKIO6kAoO9gJAcAABiJkRwASDA8qRzoGUZyAACAkQg5AADASIQcAABgJK7JAYAEx6zIQPcYyQEAAEZiJAf4BJgXBwD6LkZyAACAkRjJAQADMZcOwEgOAAAwFCEHAAAYiZADAACMxDU5QA9xJxUAJBZCDgD0A0wYiP6In6sAAICRGMkBgH6K28xhOkIO8CG4BgcAEhs/VwEAACMxkgOIURtA4uJkmIeRHAAAYCRCDgAAMBIhBwAAGIlrctAvcQ0O0DPcZo5ERsgBAPRYT/6DQBBCX0HIgfEYtQGA/omQA+MQagAAEiEHABBjzLeDvoKQg4TGqA2QGLiAGfFAyEFCIdQAZmC0BxdCwoecVatWadmyZQoGg5o4caKeeOIJTZkyJd5tAQDOEXduIdYSOuRs3LhRxcXFqqioUE5OjlasWCG/36+GhgZlZGTEuz2cI0ZpAHwcfvbCuUiyLMuKdxPnKycnR1dddZWefPJJSVJnZ6cyMzN1zz33qKSk5GPfH4lE5HK5FA6H5XQ6e7vdfo0AAyBeugtChKXE1tPv74QdyWlra1N9fb0WLlxor0tOTlZubq5qa2u7fU9ra6taW1vt1+FwWNIHf1nomcsffDHeLQDAORkzb/N51exf7O+NdhADXd/bHzdOk7Ah5z//+Y86Ojrkdruj1rvdbh08eLDb95SVlWnx4sVnrc/MzOyVHgEAicu1It4d4OO89957crlcH7o9YUPO+Vi4cKGKi4vt152dnWpubtbIkSOVlJQUx87wYSKRiDIzM3XkyBF+UkwQnLPExHlLPP35nFmWpffee09er/cj6xI25IwaNUoDBgxQKBSKWh8KheTxeLp9j8PhkMPhiFqXlpbWWy0ihpxOZ7/7R5zoOGeJifOWePrrOfuoEZwuyRegj16RkpKiSZMmqbq62l7X2dmp6upq+Xy+OHYGAAD6goQdyZGk4uJizZkzR5MnT9aUKVO0YsUKnThxQnfccUe8WwMAAHGW0CHnW9/6lt59912VlpYqGAzqiiuu0Pbt28+6GBmJy+Fw6MEHHzzrZ0b0XZyzxMR5Szycs4+X0PPkAAAAfJiEvSYHAADgoxByAACAkQg5AADASIQcAABgJEIO+qRHHnlEV199tYYMGfKhEzY2NjYqLy9PQ4YMUUZGhu677z6dOnXqwjaKs6xatUoXX3yxBg8erJycHL366qvxbgn/r6amRjfffLO8Xq+SkpK0devWqO2WZam0tFSjR49WamqqcnNz9fbbb8enWUj64HFEV111lYYPH66MjAzNnDlTDQ0NUTUnT55UIBDQyJEjNWzYMM2aNeusiXL7K0IO+qS2tjZ94xvf0N13393t9o6ODuXl5amtrU27du3SunXrtHbtWpWWll7gTnG6jRs3qri4WA8++KBef/11TZw4UX6/X01NTfFuDZJOnDihiRMnatWqVd1uLy8v18qVK1VRUaG6ujoNHTpUfr9fJ0+evMCdosvOnTsVCAS0e/duVVVVqb29XTNmzNCJEyfsmnnz5un555/X5s2btXPnTh09elS33HJLHLvuQyygD3vmmWcsl8t11vo//OEPVnJyshUMBu11q1evtpxOp9Xa2noBO8TppkyZYgUCAft1R0eH5fV6rbKysjh2he5IsrZs2WK/7uzstDwej7Vs2TJ7XUtLi+VwOKzf/va3cegQ3WlqarIkWTt37rQs64NzNGjQIGvz5s12zVtvvWVJsmpra+PVZp/BSA4SUm1trSZMmBA18aPf71ckEtGBAwfi2Fn/1dbWpvr6euXm5trrkpOTlZubq9ra2jh2hp44dOiQgsFg1PlzuVzKycnh/PUh4XBYkpSeni5Jqq+vV3t7e9R5GzdunMaMGcN5Ez9XIUEFg8GzZrbueh0MBuPRUr/3n//8Rx0dHd2eF85J39d1jjh/fVdnZ6eKiop0zTXX6PLLL5f0wXlLSUk569pFztsHCDm4YEpKSpSUlPSRy8GDB+PdJgD0SYFAQPv379ezzz4b71YSRkI/uwqJZf78+frud7/7kTWf+cxnerQvj8dz1l07XXcTeDye8+oPn8yoUaM0YMCAs+7qCIVCnJME0HWOQqGQRo8eba8PhUK64oor4tQVuhQWFmrbtm2qqanRRRddZK/3eDxqa2tTS0tL1GgO/+4+wEgOLphPfepTGjdu3EcuKSkpPdqXz+fTvn37ou7aqaqqktPpVHZ2dm8dAj5CSkqKJk2apOrqantdZ2enqqur5fP54tgZeiIrK0sejyfq/EUiEdXV1XH+4siyLBUWFmrLli3asWOHsrKyorZPmjRJgwYNijpvDQ0Namxs5LyJkRz0UY2NjWpublZjY6M6Ojq0d+9eSdIll1yiYcOGacaMGcrOztZtt92m8vJyBYNBPfDAAwoEAjyRN46Ki4s1Z84cTZ48WVOmTNGKFSt04sQJ3XHHHfFuDZKOHz+ud955x3596NAh7d27V+np6RozZoyKioq0ZMkSXXrppcrKytKiRYvk9Xo1c+bM+DXdzwUCAW3YsEG///3vNXz4cPs6G5fLpdTUVLlcLhUUFKi4uFjp6elyOp2655575PP5NHXq1Dh33wfE+/YuoDtz5syxJJ21vPzyy3bN4cOHrRtuuMFKTU21Ro0aZc2fP99qb2+PX9OwLMuynnjiCWvMmDFWSkqKNWXKFGv37t3xbgn/7+WXX+7239WcOXMsy/rgNvJFixZZbrfbcjgc1vTp062Ghob4Nt3PdXe+JFnPPPOMXfP+++9bP/zhD60RI0ZYQ4YMsb72ta9Z//73v+PXdB+SZFmWdeGjFQAAQO/imhwAAGAkQg4AADASIQcAABiJkAMAAIxEyAEAAEYi5AAAACMRcgAAgJEIOQAAwEiEHAAAYCRCDgAAMBIhBwAAGImQAwAAjPR/+s478L10K4IAAAAASUVORK5CYII=",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# plot a histogram of the reward difference\n",
    "import matplotlib.pyplot as plt\n",
    "plt.hist(reward_diff, bins=100)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "top_quantile_threshold = reward_diff.abs().quantile(0.9)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "top_reward_diff_ind = (reward_diff.abs() > top_quantile_threshold).nonzero().squeeze()\n",
    "top_reward_diff_ind\n",
    "torch.save(top_reward_diff_ind, \"rm_save_files/top_reward_diff_ind.pt\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "logan",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}