TD3notebook-random_points.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abefc8b4-e11b-4db9-9c41-35ea901d17a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "import gym\n",
    "import pybullet_envs\n",
    "import os\n",
    "\n",
    "import utils\n",
    "import TD3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a267aa0c-028a-43b2-b366-cd4b6ca7fcd1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from gym.envs.registration import registry, make, spec\n",
    "\n",
    "\n",
    "def register(id, *args, **kvargs):\n",
    "  if id in registry.env_specs:\n",
    "    return\n",
    "  else:\n",
    "    return gym.envs.registration.register(id, *args, **kvargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5cb77905-d358-43e0-b8d7-cb3926d132e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "register(id='MyAntBulletEnv-v0',\n",
    "         entry_point='override_ant_random_points:MyAntBulletEnv',\n",
    "         max_episode_steps=1000,\n",
    "         reward_threshold=2500.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c00a92a-87aa-4ce8-b9e1-7823302879ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Runs policy for X episodes and returns average reward\n",
    "# A fixed seed is used for the eval environment\n",
    "def eval_policy(policy, env_name, seed, eval_episodes=10):\n",
    "    eval_env = gym.make(env_name)\n",
    "    eval_env.seed(seed + 100)\n",
    "\n",
    "    avg_reward = 0.\n",
    "    for _ in range(eval_episodes):\n",
    "        r = np.linalg.norm([20,20])\n",
    "        rand_deg = np.random.randint(0,360) # degrees here for reader clarity, rather than directly in 2pi\n",
    "\n",
    "        rand_x = r*np.cos(np.pi/180 * rand_deg)\n",
    "        rand_y = r*np.sin(np.pi/180 * rand_deg)\n",
    "\n",
    "        eval_env.robot.walk_target_x = rand_x\n",
    "        eval_env.robot.walk_target_y = rand_y\n",
    "        state, done = eval_env.reset(), False\n",
    "        while not done:\n",
    "            action = policy.select_action(np.array(state))\n",
    "            state, reward, done, _ = eval_env.step(action)\n",
    "            avg_reward += reward\n",
    "\n",
    "    avg_reward /= eval_episodes\n",
    "\n",
    "    print(\"---------------------------------------\")\n",
    "    print(f\"Evaluation over {eval_episodes} episodes: {avg_reward:.3f}\")\n",
    "    print(\"---------------------------------------\")\n",
    "    return avg_reward\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2ee5a5d-dfb9-468d-b08d-df957641d12b",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.set_num_threads(4) # the training will eat all available cores by default and does not scale well, you can play around with this setting for your own machine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30a181cc-363c-458f-9ea5-4f712af935dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.get_num_threads()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0c2c01f-89bc-455e-9f24-fe424a17a5c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def main():\n",
    "    args = {\n",
    "            \"policy\" : \"TD3\",                  # Policy name (TD3, DDPG or OurDDPG)\n",
    "            \"env\" : \"MyAntBulletEnv-v0\",         # OpenAI gym environment name\n",
    "            \"seed\" : 0,                        # Sets Gym, PyTorch and Numpy seeds\n",
    "            \"start_timesteps\" : 25e3,          # Time steps initial random policy is used\n",
    "            \"eval_freq\" : 5e3,                 # How often (time steps) we evaluate\n",
    "            \"max_timesteps\" : 2e6,             # Max time steps to run environment\n",
    "            \"expl_noise\" : 0.1,                # Std of Gaussian exploration noise\n",
    "            \"batch_size\" : 256,                # Batch size for both actor and critic\n",
    "            \"discount\" : 0.99,                 # Discount factor\n",
    "            \"tau\" : 0.007,                     # Target network update rate\n",
    "            \"policy_noise\" : 0.2,              # Noise added to target policy during critic update\n",
    "            \"noise_clip\" : 0.5,                # Range to clip target policy noise\n",
    "            \"policy_freq\" : 2,                 # Frequency of delayed policy updates\n",
    "            \"save_model\" : \"store_true\",       # Save model and optimizer parameters\n",
    "            \"load_model\" : \"\",                 # Model load file name, \"\" doesn't load, \"default\" uses file_name\n",
    "           }\n",
    "\n",
    "    file_name = f\"{args['policy']}_{args['env']}_{args['seed']}_{args['tau']}\"\n",
    "    print(\"---------------------------------------\")\n",
    "    print(f\"Policy: {args['policy']}, Env: {args['env']}, Seed: {args['seed']}\")\n",
    "    print(\"---------------------------------------\")\n",
    "\n",
    "    if not os.path.exists(\"./results\"):\n",
    "        os.makedirs(\"./results\")\n",
    "\n",
    "    if args['save_model'] and not os.path.exists(\"./models\"):\n",
    "        os.makedirs(\"./models\")\n",
    "\n",
    "    env = gym.make(args['env'])\n",
    "\n",
    "    # Set seeds\n",
    "    env.seed(args['seed'])\n",
    "    env.action_space.seed(args['seed'])\n",
    "    torch.manual_seed(args['seed'])\n",
    "    np.random.seed(args['seed'])\n",
    "\n",
    "    state_dim = env.observation_space.shape[0]\n",
    "    action_dim = env.action_space.shape[0] \n",
    "    max_action = float(env.action_space.high[0])\n",
    "\n",
    "    kwargs = {\n",
    "        \"state_dim\": state_dim,\n",
    "        \"action_dim\": action_dim,\n",
    "        \"max_action\": max_action,\n",
    "        \"discount\": args['discount'],\n",
    "        \"tau\": args['tau'],\n",
    "    }\n",
    "\n",
    "    # Initialize policy\n",
    "    if args['policy'] == \"TD3\":\n",
    "        # Target policy smoothing is scaled wrt the action scale\n",
    "        kwargs[\"policy_noise\"] = args['policy_noise'] * max_action\n",
    "        kwargs[\"noise_clip\"] = args['noise_clip'] * max_action\n",
    "        kwargs[\"policy_freq\"] = args['policy_freq']\n",
    "        policy = TD3.TD3(**kwargs)\n",
    "\n",
    "    if args['load_model'] != \"\":\n",
    "        policy_file = file_name if args['load_model'] == \"default\" else args['load_model']\n",
    "        policy.load(f\"./models/{policy_file}\")\n",
    "\n",
    "    replay_buffer = utils.ReplayBuffer(state_dim, action_dim)\n",
    "\n",
    "    # Evaluate untrained policy\n",
    "    evaluations = [eval_policy(policy, args['env'], args['seed'])]\n",
    "\n",
    "    state, done = env.reset(), False\n",
    "    episode_reward = 0\n",
    "    episode_timesteps = 0\n",
    "    episode_num = 0\n",
    "\n",
    "    for t in range(int(args['max_timesteps'])):\n",
    "\n",
    "        episode_timesteps += 1\n",
    "\n",
    "        # Select action randomly or according to policy\n",
    "        if t < args['start_timesteps']:\n",
    "            action = env.action_space.sample()\n",
    "        else:\n",
    "            action = (\n",
    "                policy.select_action(np.array(state))\n",
    "                + np.random.normal(0, max_action * args['expl_noise'], size=action_dim)\n",
    "            ).clip(-max_action, max_action)\n",
    "\n",
    "        # Perform action\n",
    "        next_state, reward, done, _ = env.step(action) \n",
    "        done_bool = float(done) if episode_timesteps < env._max_episode_steps else 0\n",
    "\n",
    "        # Store data in replay buffer\n",
    "        replay_buffer.add(state, action, next_state, reward, done_bool)\n",
    "\n",
    "        state = next_state\n",
    "        episode_reward += reward\n",
    "\n",
    "        # Train agent after collecting sufficient data\n",
    "        if t >= args['start_timesteps']:\n",
    "            policy.train(replay_buffer, args['batch_size'])\n",
    "\n",
    "        if done: \n",
    "            # +1 to account for 0 indexing. +0 on ep_timesteps since it will increment +1 even if done=True\n",
    "            print(f\"Total T: {t+1} Episode Num: {episode_num+1} Episode T: {episode_timesteps} Reward: {episode_reward:.3f}\")\n",
    "            # Reset environment\n",
    "            \n",
    "            #samples x,y from a circle of r=sqrt(20**2+20**2)\n",
    "            r = np.linalg.norm([20,20])\n",
    "            rand_deg = np.random.randint(0,360) # degrees here for reader clarity, rather than directly in 2pi\n",
    "\n",
    "            rand_x = r*np.cos(np.pi/180 * rand_deg)\n",
    "            rand_y = r*np.sin(np.pi/180 * rand_deg)\n",
    "            \n",
    "            env.robot.walk_target_x = rand_x\n",
    "            env.robot.walk_target_y = rand_y\n",
    "            state, done = env.reset(), False\n",
    "            episode_reward = 0\n",
    "            episode_timesteps = 0\n",
    "            episode_num += 1 \n",
    "\n",
    "        # Evaluate episode\n",
    "        if (t + 1) % args['eval_freq'] == 0:\n",
    "            evaluations.append(eval_policy(policy, args['env'], args['seed']))\n",
    "            np.save(f\"./results/{file_name}\", evaluations)\n",
    "            if args['save_model']: policy.save(f\"./models/{file_name}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1ee9d19-9907-43d5-97e3-5f61521ec431",
   "metadata": {},
   "outputs": [],
   "source": [
    "main()"
   ]
  }
 ],
 "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}