{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introduction\n",
    "\n",
    "This notebook aims to give you an extremely simple example of how to use Qlib RL. We will show you how to build up a simulator, a policy, and reward function. Then, we will show you how to run RL training & backtest workflows based on the aforementioned componenets.\n",
    "\n",
    "This notebook assumes that you are aware of basic concepts of reinforcement learning. If you are unfamiliar with RL, please search for materials (for example, [Wikipedia](https://en.wikipedia.org/wiki/Reinforcement_learning)) to get basic understanding of RL.\n",
    "\n",
    "Let's start with the simulator."
   ]
  },
  {
   "attachments": {
    "image.png": {
     "image/png": "[encoded data, 57500 characters]"
    }
   },
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simulator\n",
    "\n",
    "![image.png](attachment:image.png)\n",
    "\n",
    "As stated in the above figure, environment is one of the two core componenets in the RL ecosystem. The operating logic of the environment is implemented in the **simulator**. In this notebook, we define a primitive simulator based on `qlib.rl.simulator.Simulator`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import namedtuple\n",
    "from typing import Any\n",
    "from qlib.rl.simulator import Simulator\n",
    "\n",
    "State = namedtuple(\"State\", [\"value\", \"last_action\"])\n",
    "\n",
    "class SimpleSimulator(Simulator[float, State, float]):\n",
    "    def __init__(self, initial: float, nsteps: int, **kwargs: Any) -> None:\n",
    "        super().__init__(initial)\n",
    "\n",
    "        self.value = initial\n",
    "        self.last_action = 0.0\n",
    "        self.remain_steps = nsteps\n",
    "\n",
    "    def step(self, action: float) -> None:\n",
    "        assert 0.0 <= action <= self.value\n",
    "        self.last_action = action\n",
    "        self.remain_steps -= 1\n",
    "\n",
    "    def get_state(self) -> State:\n",
    "        return State(self.value, self.last_action)\n",
    "\n",
    "    def done(self) -> bool:\n",
    "        return self.remain_steps == 0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This simulator will run `nsteps` steps fixedly. It has a internal `value`, but this value will not be affected by the action. We will talk about what this property does later. The expected action from the agent should be a float number in the range of `[0.0, self.value]`. Each time the `step()` method is called, the simulator will record the action value, and wrap it into the simulator's state (see in `get_state()` method)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Intepreters & Reward\n",
    "\n",
    "There are three items in the process of interaction between environment and agent: **state**, **action**, and **reward**. Let's talk about state and action first. \n",
    "\n",
    "Environment and agent may represent state/action in different formats, so a \"format converter\" is needed between them. In Qlib RL, it is called **interpreter**. State interpreter takes states from the environment, and convert it to a format that agent could understand. Action interpreter works in the opposite direction. It takes actions from agent and convert it to a format that environment accepts.\n",
    "\n",
    "As we mentioned when introducing the simulator, a state generated by the simulator is a tuple of two float numbers. The state interpreter is defined as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Tuple\n",
    "import numpy as np\n",
    "from gym import spaces\n",
    "from qlib.rl.interpreter import StateInterpreter\n",
    "\n",
    "class SimpleStateInterpreter(StateInterpreter[Tuple[float, float], np.ndarray]):\n",
    "    def interpret(self, state: State) -> np.ndarray:\n",
    "        # Convert state.value to a 1D Numpy array\n",
    "        # last_action is not used by agents.\n",
    "        return np.array([state.value], dtype=np.float32)\n",
    "\n",
    "    @property\n",
    "    def observation_space(self) -> spaces.Box:\n",
    "        return spaces.Box(0, np.inf, shape=(1,), dtype=np.float32)\n",
    "    \n",
    "state_interpreter = SimpleStateInterpreter()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As for the action interpreter, in this notebook, we assume that the agent will not directly generate a float action. Instead, it generates a bucketized value (an integer) because we want to use a categorical model to build the policy. To be concrete, we divide `simulator.value` into N parts evenly, and the agent's action should be an integer in `[0, N]`. The value of this integer will be converted to the float number as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from qlib.rl.interpreter import ActionInterpreter\n",
    "\n",
    "class SimpleActionInterpreter(ActionInterpreter[State, int, float]):\n",
    "    def __init__(self, n_value: int) -> None:\n",
    "        self.n_value = n_value\n",
    "\n",
    "    @property\n",
    "    def action_space(self) -> spaces.Discrete:\n",
    "        return spaces.Discrete(self.n_value + 1)\n",
    "\n",
    "    def interpret(self, simulator_state: State, action: int) -> float:\n",
    "        assert 0 <= action <= self.n_value\n",
    "        # simulator_state.value is used as the denominator\n",
    "        return simulator_state.value * (action / self.n_value)\n",
    "    \n",
    "action_interpreter = SimpleActionInterpreter(n_value=10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, let's define the reward function. In this example, we use the proportion of `action` in `value` as the reward. In other words, the larger the action, the larger the reward:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from qlib.rl.reward import Reward\n",
    "\n",
    "class SimpleReward(Reward[State]):\n",
    "    def reward(self, simulator_state: State) -> float:\n",
    "        # Use last_action to calculate reward. This is why it should be in the state.\n",
    "        rew = simulator_state.last_action / simulator_state.value\n",
    "        return rew\n",
    "    \n",
    "reward = SimpleReward()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Agent & Policy\n",
    "\n",
    "The last thing we haven't talk about in the RL ecosystem is the agent. In RL, agents take actions follow particular policies, so defining agents is actually equivalent to defining policies. In this example, we define a PPO policy with a naive categorical neural network:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import torch\n",
    "from torch import nn\n",
    "from qlib.rl.order_execution import PPO\n",
    "\n",
    "class SimpleFullyConnect(nn.Module):\n",
    "    def __init__(self, dims: List[int]) -> None:\n",
    "        super().__init__()\n",
    "\n",
    "        self.dims = [1] + dims\n",
    "        self.output_dim = dims[-1]\n",
    "\n",
    "        layers = []\n",
    "        for in_dim, out_dim in zip(self.dims[:-1], self.dims[1:]):\n",
    "            layers.append(nn.Linear(in_dim, out_dim))\n",
    "            layers.append(nn.ReLU())\n",
    "        self.fc = nn.Sequential(*layers)\n",
    "\n",
    "    def forward(self, x: torch.Tensor) -> torch.Tensor:\n",
    "        return self.fc(x)\n",
    "    \n",
    "policy = PPO(\n",
    "    network=SimpleFullyConnect(dims=[16, 8]),\n",
    "    obs_space=state_interpreter.observation_space,\n",
    "    action_space=action_interpreter.action_space,\n",
    "    lr=0.01,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset\n",
    "\n",
    "In order to get sufficient training data, usually we need to build multiple environments in parallel, and collect data from all these environments. To achieve that, we need to create multiple simulators with different initial settings. In Qlib RL, we could use a data set to implement this function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import Dataset\n",
    "\n",
    "class SimpleDataset(Dataset):\n",
    "    def __init__(self, positions: List[float]) -> None:\n",
    "        self.positions = positions\n",
    "\n",
    "    def __len__(self) -> int:\n",
    "        return len(self.positions)\n",
    "\n",
    "    def __getitem__(self, index: int) -> float:\n",
    "        return self.positions[index]\n",
    "    \n",
    "dataset = SimpleDataset(positions=[10.0, 50.0, 100.0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Each number in `positions` will be used to create a unique simulator. The creation of simulators will be handle by Qlib RL automatically. The only thing that developers need to do is defining the data set class and create a data set instance with data they want to use."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training workflow\n",
    "\n",
    "Based on all the componenets we already have, it is super easy to launch a training workflow:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "from typing import cast\n",
    "from qlib.rl.trainer import Checkpoint, train\n",
    "\n",
    "NSTEPS = 10\n",
    "\n",
    "trainer_kwargs = {\n",
    "    \"max_iters\": 10,\n",
    "    \"finite_env_type\": \"dummy\",\n",
    "    \"callbacks\": [Checkpoint(\n",
    "        dirpath=Path(\"./checkpoints\"),\n",
    "        every_n_iters=1,\n",
    "        save_latest=\"copy\",\n",
    "    )],\n",
    "}\n",
    "vessel_kwargs = {\n",
    "    \"update_kwargs\": {\"batch_size\": 16, \"repeat\": 5},\n",
    "    \"episode_per_iter\": 50,\n",
    "}\n",
    "\n",
    "print(\"Training started\")\n",
    "train(\n",
    "    simulator_fn=lambda position: SimpleSimulator(position, NSTEPS),\n",
    "    state_interpreter=state_interpreter,\n",
    "    action_interpreter=action_interpreter,\n",
    "    policy=policy,\n",
    "    reward=reward,\n",
    "    initial_states=cast(List[float], SimpleDataset([10.0, 50.0, 100.0])),\n",
    "    trainer_kwargs=trainer_kwargs,\n",
    "    vessel_kwargs=vessel_kwargs,\n",
    ")\n",
    "print(\"Training finished\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Backtest workflow\n",
    "\n",
    "After the policy is trained, we could let it run in the environment and see its behaviors. This is called \"backtest\":"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tianshou.data import Batch\n",
    "\n",
    "simulator = SimpleSimulator(100.0, NSTEPS)\n",
    "state = simulator.get_state()\n",
    "obs = [{\"obs\": state_interpreter.interpret(state)}]\n",
    "policy_out = policy(Batch(obs))\n",
    "act = float(action_interpreter.interpret(state, policy_out.act))\n",
    "\n",
    "simulator.step(act)\n",
    "rew = float(reward(simulator.get_state()))\n",
    "\n",
    "print(f\"Action = {act:.6f}, Reward = {rew:.6f}.\")"
   ]
  }
 ],
 "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
