TD3_view_random_points.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da49dba5-0b24-4fcd-af74-8277fa52fe67",
   "metadata": {},
   "outputs": [],
   "source": [
    "## This can leave open processes if you don't keep track of them, be sure to clean up after"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69ccc51b-ef01-46ab-a267-e13a14b93a66",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "import gym\n",
    "import pybullet_envs\n",
    "import os\n",
    "import time\n",
    "\n",
    "import utils\n",
    "import TD3\n",
    "from numpngw import write_apng"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69d4210a-84f8-4342-ba2a-0ba85a083254",
   "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": "b580302c-4df3-4dbf-8bc6-fc59dc7174a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "register(id='MyAntBulletEnv-v0',\n",
    "         entry_point='override_ant_random_points:MyAntBulletEnv',\n",
    "         max_episode_steps=2000,\n",
    "         reward_threshold=2500.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ca51466-69d2-42d0-87e3-4a89d9d9e72b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def eval_policy_render(policy, env_name, seed, eval_episodes=5):\n",
    "    eval_env = gym.make(env_name, render=True)\n",
    "    eval_env.seed(seed + 100)\n",
    "\n",
    "    avg_reward = 0.\n",
    "    for i 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",
    "        images = [eval_env.render('rgb_array')]\n",
    "        time_step_counter = 1\n",
    "        while not done:\n",
    "            if time_step_counter % 500 == 0:\n",
    "                rand_deg = np.random.randint(0,360) # degrees here for reader clarity, rather than directly in 2pi\n",
    "                rand_x = r*np.cos(np.pi/180 * rand_deg)\n",
    "                rand_y = r*np.sin(np.pi/180 * rand_deg)\n",
    "                eval_env.robot.walk_target_x = rand_x\n",
    "                eval_env.robot.walk_target_y = rand_y\n",
    "            time.sleep(1. / 60.)\n",
    "            action = policy.select_action(np.array(state))\n",
    "            state, reward, done, _ = eval_env.step(action)\n",
    "            avg_reward += reward\n",
    "            images.append(eval_env.render('rgb_array'))\n",
    "            time_step_counter = time_step_counter + 1\n",
    "\n",
    "        print(f'Saving animation: anim_{i}.png, lenght: {len(images)} frames.')\n",
    "        #write_apng(f'anim_{i}.png', images[::2], delay=50)  #uncomment this line to save animations\n",
    "        print('Save file complete')\n",
    "            \n",
    "    avg_reward /= eval_episodes\n",
    "    return avg_reward"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1816f732-d7bf-4e41-9149-278fa9ee3193",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_policy(env_name_var):\n",
    "    args = {\n",
    "            \"policy\" : \"TD3\",                  # Policy name (TD3, DDPG or OurDDPG)\n",
    "            \"env\" : env_name_var,              # 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\" : \"default\",           # 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",
    "    return policy\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6984e32c-1c87-430a-b81a-7b4da2b677ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "policy = load_policy(\"MyAntBulletEnv-v0\")\n",
    "eval_policy_render(policy, \"MyAntBulletEnv-v0\", 0)"
   ]
  }
 ],
 "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
}