examples/custom_environment.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9f9addeb-6e78-4439-ae4f-98418f40d39a",
   "metadata": {},
   "source": [
    "# 1. Environment definition\n",
    "To use custom environments in **RLLTE**, it suffices to follow the [gymnasium](https://gymnasium.farama.org/) interface and prepare your environment following [Tutorials: Make Your Own Custom Environment](https://gymnasium.farama.org/tutorials/gymnasium_basics/environment_creation/#). A example is:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e514431d-24cd-4895-a4f3-18e5a0d27551",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "pygame 2.4.0 (SDL 2.26.4, Python 3.8.16)\n",
      "Hello from the pygame community. https://www.pygame.org/contribute.html\n"
     ]
    }
   ],
   "source": [
    "import gymnasium as gym\n",
    "import numpy as np\n",
    "\n",
    "class CustomEnv(gym.Env):\n",
    "    def __init__(self, total_length) -> None:\n",
    "        super().__init__()\n",
    "        self.observation_space = gym.spaces.Box(\n",
    "            shape=(9, 84, 84),\n",
    "            high=255.0,\n",
    "            low=0.,\n",
    "            dtype=np.uint8\n",
    "        )\n",
    "        self.action_space = gym.spaces.Box(\n",
    "            shape=(7,),\n",
    "            high=1.,\n",
    "            low=-1.,\n",
    "            dtype=np.float32\n",
    "        )\n",
    "        self.total_length = total_length\n",
    "        self.count = 0\n",
    "\n",
    "    def step(self, action):\n",
    "        obs = self.observation_space.sample()\n",
    "        reward = np.random.rand()\n",
    "        if self.count < self.total_length:\n",
    "            terminated = truncated = False\n",
    "        else:\n",
    "            terminated = truncated = True\n",
    "        info = {\"discount\": 0.99}\n",
    "        self.count += 1\n",
    "\n",
    "        return obs, reward, terminated, truncated, info\n",
    "\n",
    "    def reset(self, seed=None, options=None):\n",
    "        self.count = 0\n",
    "        return self.observation_space.sample(), {\"discount\": 0.99}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "acb2df0b-358e-416e-8330-3cf079640730",
   "metadata": {},
   "source": [
    "# 2. Use `make_rllte_env`\n",
    "In **RLLTE**, the environments are assumed to be ***vectorized*** and a `make_rllte_env` function is used to warp the environments:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "8fcf01f9-baf8-410d-9417-7f21a37de5e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from rllte.env import make_rllte_env\n",
    "# create vectorized environments\n",
    "env = make_rllte_env(env_id=CustomEnv, \n",
    "                     device=\"cuda\", \n",
    "                     env_kwargs={'total_length': 499} # set env arguments\n",
    "                     )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "377c7fa3-1f68-407c-bfa0-56ef716521c2",
   "metadata": {},
   "source": [
    "After that, you can use the custom environment in application directly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "6f5d4932-b6f0-4005-8b1c-9ed4d18bce39",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[08/29/2023 12:12:21 PM] - [\u001b[1m\u001b[34mINFO.\u001b[0m] - Invoking RLLTE Engine...\n",
      "[08/29/2023 12:12:21 PM] - [\u001b[1m\u001b[34mINFO.\u001b[0m] - ================================================================================\n",
      "[08/29/2023 12:12:21 PM] - [\u001b[1m\u001b[34mINFO.\u001b[0m] - Tag               : drqv2_dmc_pixel\n",
      "[08/29/2023 12:12:21 PM] - [\u001b[1m\u001b[34mINFO.\u001b[0m] - Device            : NVIDIA GeForce RTX 3090\n",
      "[08/29/2023 12:12:21 PM] - [\u001b[1m\u001b[33mDEBUG\u001b[0m] - Agent             : DrQv2\n",
      "[08/29/2023 12:12:21 PM] - [\u001b[1m\u001b[33mDEBUG\u001b[0m] - Encoder           : TassaCnnEncoder\n",
      "[08/29/2023 12:12:21 PM] - [\u001b[1m\u001b[33mDEBUG\u001b[0m] - Policy            : OffPolicyDetActorDoubleCritic\n",
      "[08/29/2023 12:12:21 PM] - [\u001b[1m\u001b[33mDEBUG\u001b[0m] - Storage           : NStepReplayStorage\n",
      "[08/29/2023 12:12:21 PM] - [\u001b[1m\u001b[33mDEBUG\u001b[0m] - Distribution      : TruncatedNormalNoise\n",
      "[08/29/2023 12:12:21 PM] - [\u001b[1m\u001b[33mDEBUG\u001b[0m] - Augmentation      : True, RandomShift\n",
      "[08/29/2023 12:12:21 PM] - [\u001b[1m\u001b[33mDEBUG\u001b[0m] - Intrinsic Reward  : False\n",
      "[08/29/2023 12:12:21 PM] - [\u001b[1m\u001b[33mDEBUG\u001b[0m] - ================================================================================\n",
      "[08/29/2023 12:12:35 PM] - [\u001b[1m\u001b[32mEVAL.\u001b[0m] - S: 0           | E: 0           | L: 500         | R: 249.162     | T: 0:00:15    \n",
      "[08/29/2023 12:12:38 PM] - [\u001b[1m\u001b[31mTRAIN\u001b[0m] - S: 1000        | E: 2           | L: 500         | R: 255.789     | FPS: 54.461    | T: 0:00:18    \n",
      "[08/29/2023 12:12:42 PM] - [\u001b[1m\u001b[31mTRAIN\u001b[0m] - S: 2000        | E: 4           | L: 500         | R: 251.800     | FPS: 91.451    | T: 0:00:21    \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[2023-08-29 12:12:55,944] torch._dynamo.convert_frame: [WARNING] torch._dynamo hit config.cache_size_limit (64)\n",
      "   function: 'forward' (/export/yuanmingqi/code/rllte/rllte/xploit/policy/off_policy_det_actor_double_critic.py:141)\n",
      "   reasons:  step == 0\n",
      "to diagnose recompilation issues, see https://pytorch.org/docs/master/dynamo/troubleshooting.html.\n",
      "[2023-08-29 12:12:56,044] torch._dynamo.convert_frame: [WARNING] torch._dynamo hit config.cache_size_limit (64)\n",
      "   function: 'get_dist' (/export/yuanmingqi/code/rllte/rllte/xploit/policy/off_policy_det_actor_double_critic.py:162)\n",
      "   reasons:  step == 0\n",
      "to diagnose recompilation issues, see https://pytorch.org/docs/master/dynamo/troubleshooting.html.\n",
      "[2023-08-29 12:12:56,140] torch._dynamo.convert_frame: [WARNING] torch._dynamo hit config.cache_size_limit (64)\n",
      "   function: 'reset' (/export/yuanmingqi/code/rllte/rllte/xplore/distribution/truncated_normal_noise.py:100)\n",
      "   reasons:  step == 0\n",
      "to diagnose recompilation issues, see https://pytorch.org/docs/master/dynamo/troubleshooting.html.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[08/29/2023 12:13:08 PM] - [\u001b[1m\u001b[31mTRAIN\u001b[0m] - S: 3000        | E: 6           | L: 500         | R: 248.395     | FPS: 62.378    | T: 0:00:48    \n",
      "[08/29/2023 12:13:20 PM] - [\u001b[1m\u001b[31mTRAIN\u001b[0m] - S: 4000        | E: 8           | L: 500         | R: 247.984     | FPS: 65.796    | T: 0:01:00    \n",
      "[08/29/2023 12:13:33 PM] - [\u001b[1m\u001b[31mTRAIN\u001b[0m] - S: 5000        | E: 10          | L: 500         | R: 249.162     | FPS: 68.033    | T: 0:01:13    \n",
      "[08/29/2023 12:13:45 PM] - [\u001b[1m\u001b[32mEVAL.\u001b[0m] - S: 5000        | E: 10          | L: 500         | R: 250.176     | T: 0:01:25    \n",
      "[08/29/2023 12:13:45 PM] - [\u001b[1m\u001b[34mINFO.\u001b[0m] - Training Accomplished!\n",
      "[08/29/2023 12:13:45 PM] - [\u001b[1m\u001b[34mINFO.\u001b[0m] - Model saved at: /export/yuanmingqi/code/rllte/examples/logs/drqv2_dmc_pixel/2023-08-29-12-12-20/model\n"
     ]
    }
   ],
   "source": [
    "from rllte.agent import DrQv2\n",
    "from rllte.env.utils import make_rllte_env\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    # env setup\n",
    "    device = \"cuda:0\"\n",
    "    env = make_rllte_env(env_id=CustomEnv, \n",
    "                        device=device, \n",
    "                        env_kwargs={'total_length': 499} # set env arguments\n",
    "                        )\n",
    "    eval_env = make_rllte_env(env_id=CustomEnv, \n",
    "                            device=device, \n",
    "                            env_kwargs={'total_length': 499} # set env arguments\n",
    "                            )\n",
    "    agent = DrQv2(env=env, \n",
    "                eval_env=eval_env, \n",
    "                device=device,\n",
    "                tag=\"drqv2_dmc_pixel\")\n",
    "    agent.train(num_train_steps=5000, log_interval=1000)"
   ]
  }
 ],
 "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.8.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}