diff --git a/projects/MuJoCo-Torch/MT06_RL_PPO.ipynb b/projects/MuJoCo-Torch/MT06_RL_PPO.ipynb index cf31782..75d2702 100644 --- a/projects/MuJoCo-Torch/MT06_RL_PPO.ipynb +++ b/projects/MuJoCo-Torch/MT06_RL_PPO.ipynb @@ -36,7 +36,9 @@ }, { "cell_type": "code", + "execution_count": null, "metadata": {}, + "outputs": [], "source": [ "import os\n", "os.environ[\"MUJOCO_GL\"] = \"egl\"\n", @@ -71,9 +73,7 @@ "os.makedirs(\"output/videos\", exist_ok=True)\n", "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n", "print(\"device:\", device, \"| torch\", torch.__version__, \"| mujoco\", mujoco.__version__)" - ], - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", @@ -86,7 +86,9 @@ }, { "cell_type": "code", + "execution_count": null, "metadata": {}, + "outputs": [], "source": [ "def layer_init(layer, std=np.sqrt(2), bias_const=0.0):\n", " torch.nn.init.orthogonal_(layer.weight, std)\n", @@ -111,12 +113,11 @@ " activation(),\n", " layer_init(nn.Linear(hidden_size, 1), std=1.0),\n", " )" - ], - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", + "id": "8ad9aa3e", "metadata": {}, "source": [ "### Generalized Advantage Estimation (GAE)\n", @@ -126,12 +127,14 @@ "\\[ \\delta_t = r_t + \\gamma\\, V(s_{t+1})\\,(1-\\text{done}_t) - V(s_t) \\]\n", "\n", "with an exponential `gae_lambda` weighting to trade off bias vs. variance. Returns are then just `advantages + values`." - ], - "id": "8ad9aa3e" + ] }, { "cell_type": "code", + "execution_count": null, + "id": "e551fbaa", "metadata": {}, + "outputs": [], "source": [ "def compute_gae(rewards, values, next_obs_values, dones, gamma, gae_lambda):\n", " \"\"\"Generalized Advantage Estimation over one rollout.\n", @@ -148,24 +151,24 @@ " advantages[t] = lastgaelam = delta + gamma * gae_lambda * (1.0 - dones[t]) * lastgaelam\n", " returns = advantages + values\n", " return advantages, returns" - ], - "execution_count": null, - "outputs": [], - "id": "e551fbaa" + ] }, { "cell_type": "markdown", + "id": "7499b820", "metadata": {}, "source": [ "### The clipped-surrogate PPO update\n", "\n", "`ppo_update` runs several epochs of minibatch SGD over the collected rollout. For each minibatch it recomputes the new log-probabilities, forms the probability `ratio` against the old policy, and optimizes the **clipped surrogate objective** (so a single update cannot move the policy too far), plus a (optionally clipped) value-function loss and an optional entropy bonus. It tracks the approximate KL and clip fraction, applies global gradient-norm clipping, and can early-stop when the approximate KL exceeds `target_kl`." - ], - "id": "7499b820" + ] }, { "cell_type": "code", + "execution_count": null, + "id": "837b707c", "metadata": {}, + "outputs": [], "source": [ "def ppo_update(agent, advantages, returns):\n", " \"\"\"Run the clipped-surrogate PPO optimization on `agent` for one rollout.\n", @@ -247,10 +250,7 @@ " \"b_values\": b_values,\n", " \"b_returns\": b_returns,\n", " }" - ], - "execution_count": null, - "outputs": [], - "id": "837b707c" + ] }, { "cell_type": "markdown", @@ -263,7 +263,9 @@ }, { "cell_type": "code", + "execution_count": null, "metadata": {}, + "outputs": [], "source": [ "class PPO(object):\n", " def __init__(self, args, obs_dim, act_dim, nn_act, device):\n", @@ -468,9 +470,7 @@ " writer.add_scalar(\"charts/SPS\", metrics[\"SPS\"], step)\n", "\n", " return metrics" - ], - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", @@ -483,7 +483,9 @@ }, { "cell_type": "code", + "execution_count": null, "metadata": {}, + "outputs": [], "source": [ "args_dict = {\n", " \"seed\": 0,\n", @@ -518,9 +520,7 @@ "log_path = f\"output/logs/PPO/{args.env_name}/{args.seed}\"\n", "os.makedirs(os.path.join(log_path, \"models\"), exist_ok=True)\n", "print(\"demo steps:\", args.total_timesteps, \"| baked checkpoint:\", PPO_CKPT)" - ], - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", @@ -533,9 +533,11 @@ }, { "cell_type": "code", + "execution_count": null, "metadata": {}, + "outputs": [], "source": [ - "envs = construct_env(args.env_name)[0]\n", + "envs = construct_env(args.env_name, args.seed)[0]\n", "\n", "agent = PPO(\n", " args=args,\n", @@ -545,9 +547,7 @@ " device=device,\n", ")\n", "print(\"obs/act:\", envs.observation_space.shape, envs.action_space.shape)" - ], - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", @@ -560,17 +560,20 @@ }, { "cell_type": "markdown", + "id": "7a78eed8", "metadata": {}, "source": [ "### Collecting an on-policy rollout\n", "\n", "PPO is **on-policy**: each update uses a fresh batch of transitions gathered by the *current* policy. `collect_rollout` fills the buffer with exactly `num_steps` transitions — sampling an action and value, stepping the env, bootstrapping the next-state value, and storing everything — resetting whenever an episode ends. It returns the running observation so the next iteration continues seamlessly." - ], - "id": "7a78eed8" + ] }, { "cell_type": "code", + "execution_count": null, + "id": "353f5203", "metadata": {}, + "outputs": [], "source": [ "def collect_rollout(agent, envs, obs, args):\n", " \"\"\"Gather one on-policy rollout of `args.num_steps` transitions into\n", @@ -596,19 +599,18 @@ "\n", " obs = next_obs_tensor\n", " return obs" - ], - "execution_count": null, - "outputs": [], - "id": "353f5203" + ] }, { "cell_type": "code", + "execution_count": null, "metadata": {}, + "outputs": [], "source": [ "def run_training(agent, args, envs, log_path):\n", " seed_everything(args.seed)\n", " envs.action_space.seed(args.seed)\n", - " eval_env = EvalManager(construct_env(args.env_name)[0], device=device, reset_seed=args.seed, eval_episodes=5)\n", + " eval_env = EvalManager(construct_env(args.env_name, args.seed)[0], device=device, reset_seed=args.seed, eval_episodes=5)\n", " obs, _ = envs.reset(seed=args.seed)\n", " obs = torch.tensor(obs, dtype=torch.float32, device=device).unsqueeze(0)\n", "\n", @@ -653,9 +655,7 @@ "\n", "# Run training.\n", "eval_results = run_training(agent, args, envs, log_path)" - ], - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", @@ -668,7 +668,9 @@ }, { "cell_type": "code", + "execution_count": null, "metadata": {}, + "outputs": [], "source": [ "import matplotlib.pyplot as plt\n", "\n", @@ -727,9 +729,7 @@ "\n", "\n", "plot_ppo_curves(agent, eval_results)" - ], - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", @@ -742,10 +742,12 @@ }, { "cell_type": "code", + "execution_count": null, "metadata": {}, + "outputs": [], "source": [ "def rollout_video(agent, args, out_path, max_steps=1000, fps=30):\n", - " env = construct_env(args.env_name, render_mode=\"rgb_array\")[0]\n", + " env = construct_env(args.env_name, args.seed, render_mode=\"rgb_array\")[0]\n", " obs, _ = env.reset(seed=args.seed)\n", " obs = torch.as_tensor(obs, dtype=torch.float32, device=device)\n", " frames, total_reward = [], 0.0\n", @@ -775,22 +777,21 @@ " print(f\"Baked checkpoint not found: {PPO_CKPT}\")\n", " print(\"Run the appendix cell (train_full_ppo(PPO_CKPT)) once to produce it, \"\n", " \"then re-run this cell.\")" - ], - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "code", + "execution_count": null, + "id": "6ac3ece9", "metadata": {}, + "outputs": [], "source": [ "Video(url=\"output/videos/mt06_ppo_cheetah.mp4\")" - ], - "execution_count": null, - "outputs": [], - "id": "6ac3ece9" + ] }, { "cell_type": "markdown", + "id": "58b3d0e7", "metadata": {}, "source": [ "## Appendix: reproduce the fully-trained 300k checkpoint\n", @@ -801,12 +802,14 @@ "normal top-to-bottom pass (it takes ~5 min on the GPU) — call `train_full_ppo(PPO_CKPT)`\n", "once to (re)generate the checkpoint, then re-run the rollout cell. After regenerating,\n", "rebuild the course image so the new checkpoint is baked in." - ], - "id": "58b3d0e7" + ] }, { "cell_type": "code", + "execution_count": null, + "id": "8a645727", "metadata": {}, + "outputs": [], "source": [ "def train_full_ppo(save_path, timesteps=FULL_TIMESTEPS):\n", " \"\"\"Train PPO from scratch for `timesteps` and save the checkpoint to `save_path`.\n", @@ -815,7 +818,7 @@ " full_log = \"output/logs/PPO/full\"\n", " os.makedirs(os.path.join(full_log, \"models\"), exist_ok=True)\n", " full_args = Namespace(**{**args_dict, \"total_timesteps\": timesteps, \"save_model_interval\": []})\n", - " env = construct_env(full_args.env_name)[0]\n", + " env = construct_env(full_args.env_name, full_args.seed)[0]\n", " full_agent = PPO(full_args, env.observation_space.shape[0], env.action_space.shape[0],\n", " nn.Tanh, device)\n", " run_training(full_agent, full_args, env, full_log)\n", @@ -826,10 +829,7 @@ "\n", "# Uncomment to (re)train the baked 300k checkpoint (~5 min on GPU):\n", "# train_full_ppo(PPO_CKPT)" - ], - "execution_count": null, - "outputs": [], - "id": "8a645727" + ] }, { "cell_type": "markdown", @@ -842,13 +842,13 @@ }, { "cell_type": "markdown", + "id": "86b146ab", "metadata": {}, "source": [ "## Acknowledgements\n", "\n", "This notebook was developed in collaboration with the lab of Professor Ping-Chun Hsieh (謝秉均), Department of Computer Science, National Yang Ming Chiao Tung University (NYCU) — [pinghsieh.github.io](https://pinghsieh.github.io/)." - ], - "id": "86b146ab" + ] }, { "cell_type": "markdown", @@ -874,4 +874,4 @@ }, "nbformat": 4, "nbformat_minor": 5 -} \ No newline at end of file +} diff --git a/projects/MuJoCo-Torch/MT07_Cross_Domain_RL.ipynb b/projects/MuJoCo-Torch/MT07_Cross_Domain_RL.ipynb index e9250b3..b8f7a25 100644 --- a/projects/MuJoCo-Torch/MT07_Cross_Domain_RL.ipynb +++ b/projects/MuJoCo-Torch/MT07_Cross_Domain_RL.ipynb @@ -38,7 +38,9 @@ }, { "cell_type": "code", + "execution_count": null, "metadata": {}, + "outputs": [], "source": [ "import os\n", "os.environ[\"MUJOCO_GL\"] = \"egl\"\n", @@ -70,9 +72,7 @@ "os.makedirs(\"output/videos\", exist_ok=True)\n", "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n", "print(\"device:\", device, \"| torch\", torch.__version__, \"| mujoco\", mujoco.__version__)" - ], - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", @@ -87,14 +87,16 @@ }, { "cell_type": "code", + "execution_count": null, "metadata": {}, + "outputs": [], "source": [ "def run_demo_training(agent, args, envs, demo_timesteps):\n", " \"\"\"Short in-notebook training — mechanics + partial curves only.\n", " Full target policies are loaded from baked checkpoints for the rollouts.\"\"\"\n", " seed_everything(args.seed)\n", " envs.action_space.seed(args.seed)\n", - " eval_env = EvalManager(construct_env(args.tar_env)[0], device=device,\n", + " eval_env = EvalManager(construct_env(args.tar_env, args.seed)[0], device=device,\n", " reset_seed=args.seed, eval_episodes=2)\n", " obs, _ = envs.reset(seed=args.seed)\n", " obs = torch.as_tensor(obs, dtype=torch.float32, device=device)\n", @@ -116,21 +118,22 @@ " if step % args.evaluate_freq == 0:\n", " eval_env.evaluate(agent, step, verbose=False)\n", " return eval_env.results" - ], - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", + "id": "962fc189", "metadata": {}, "source": [ "**Helper 2 — training & evaluation curves.** `plot_curves` draws the six standard SAC/QAvatar diagnostics (Q loss, actor loss, alpha, alpha loss, and the periodic evaluation reward/length) for the two agents on shared axes." - ], - "id": "962fc189" + ] }, { "cell_type": "code", + "execution_count": null, + "id": "3aa0f7c2", "metadata": {}, + "outputs": [], "source": [ "def plot_curves(agent_q, agent_s, q_res, s_res, suffix=\"\"):\n", " plt.figure(figsize=(16, 8))\n", @@ -153,22 +156,22 @@ " plt.title(\"Eval Length\"); plt.xlabel(\"env step\"); plt.legend()\n", " plt.suptitle(\"Short in-notebook demo curves\" + suffix)\n", " plt.tight_layout(); plt.show(); plt.close()" - ], - "execution_count": null, - "outputs": [], - "id": "3aa0f7c2" + ] }, { "cell_type": "markdown", + "id": "22b4ec41", "metadata": {}, "source": [ "**Helper 3 — side-by-side rollout video.** `rollout_compare` rolls both agents in the target env and writes a side-by-side `.mp4`. MuJoCo uses the gym renderer; robosuite is rendered through a direct `mujoco.Renderer` on geom group 1 (`_make_rs_renderer` / `_rs_frame`) to avoid the corrupted frames the built-in camera path produces on this ROCm/gfx1151 stack. `_label` overlays the running return on each panel." - ], - "id": "22b4ec41" + ] }, { "cell_type": "code", + "execution_count": null, + "id": "c70a10f1", "metadata": {}, + "outputs": [], "source": [ "def _label(frame, text):\n", " img = Image.fromarray(frame)\n", @@ -192,8 +195,8 @@ "\n", "\n", "def rollout_compare(agent_q, agent_s, args, out_path, max_steps=300, fps=30, camera=\"frontview\"):\n", - " q_env, env_type = construct_env(args.tar_env, render_mode=\"rgb_array\")\n", - " s_env, _ = construct_env(args.tar_env, render_mode=\"rgb_array\")\n", + " q_env, env_type = construct_env(args.tar_env, args.seed, render_mode=\"rgb_array\")\n", + " s_env, _ = construct_env(args.tar_env, args.seed, render_mode=\"rgb_array\")\n", " rs = env_type == \"robosuite\"\n", " if rs:\n", " qR, qopt, qsim = _make_rs_renderer(q_env)\n", @@ -231,10 +234,7 @@ " imageio.mimsave(out_path, frames, fps=fps)\n", " print(f\"saved {out_path} | QAvatar return={q_ret:.1f} | SAC return={s_ret:.1f}\")\n", " return out_path" - ], - "execution_count": null, - "outputs": [], - "id": "c70a10f1" + ] }, { "cell_type": "markdown", @@ -251,7 +251,9 @@ }, { "cell_type": "code", + "execution_count": null, "metadata": {}, + "outputs": [], "source": [ "ant_args = Namespace(\n", " seed=0, total_timesteps=300000, buffer_size=20000, learning_starts=100,\n", @@ -263,15 +265,13 @@ ")\n", "DEMO_STEPS_ANT = 1500 # short demo; the real policies come from baked checkpoints\n", "\n", - "ant_env, _ = construct_env(ant_args.tar_env)\n", - "ant_src = [construct_env(e)[0] for e in ant_args.src_envs]\n", + "ant_env, _ = construct_env(ant_args.tar_env, ant_args.seed)\n", + "ant_src = [construct_env(e, ant_args.seed)[0] for e in ant_args.src_envs]\n", "ant_q = QAvatar(ant_args, ant_src, ant_env, device)\n", "ant_s = SAC(ant_args, ant_env.observation_space.shape[0], ant_env.action_space.shape[0],\n", " ant_env.action_space.low, ant_env.action_space.high, device)\n", "print(\"target obs/act:\", ant_env.observation_space.shape, ant_env.action_space.shape)" - ], - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", @@ -284,14 +284,14 @@ }, { "cell_type": "code", + "execution_count": null, "metadata": {}, + "outputs": [], "source": [ "ant_q_res = run_demo_training(ant_q, ant_args, ant_env, DEMO_STEPS_ANT)\n", "ant_s_res = run_demo_training(ant_s, ant_args, ant_env, DEMO_STEPS_ANT)\n", "plot_curves(ant_q, ant_s, ant_q_res, ant_s_res, \" — Ant 4→5 leg\")" - ], - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", @@ -304,23 +304,23 @@ }, { "cell_type": "code", + "execution_count": null, "metadata": {}, + "outputs": [], "source": [ "ant_q.load(os.path.join(CDRL_ASSETS, \"checkpoints\", \"ant_4leg_to_5leg\", \"QAvatar\", \"steps_300000.pt\"))\n", "ant_s.load(os.path.join(CDRL_ASSETS, \"checkpoints\", \"ant_4leg_to_5leg\", \"SAC\", \"steps_300000.pt\"))\n", "rollout_compare(ant_q, ant_s, ant_args, \"output/videos/mt07_ant_4to5.mp4\", max_steps=300)" - ], - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "code", + "execution_count": null, "metadata": {}, + "outputs": [], "source": [ "Video(url=\"output/videos/mt07_ant_4to5.mp4\")" - ], - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", @@ -333,10 +333,12 @@ }, { "cell_type": "code", + "execution_count": null, "metadata": {}, + "outputs": [], "source": [ "door_args = Namespace(\n", - " seed=0, total_timesteps=500000, buffer_size=20000, learning_starts=200,\n", + " seed=4, total_timesteps=500000, buffer_size=20000, learning_starts=200,\n", " q_lr=1e-4, policy_lr=1e-4, decoder_lr=1e-4, nn_hidden_size=128, batch_size=256,\n", " gamma=0.99, temperature_opt=True, evaluate_freq=400, tau=0.005,\n", " save_model_interval=[], cuda=True,\n", @@ -345,26 +347,24 @@ ")\n", "DEMO_STEPS_DOOR = 800\n", "\n", - "door_env, _ = construct_env(door_args.tar_env)\n", - "door_src = [construct_env(e)[0] for e in door_args.src_envs]\n", + "door_env, _ = construct_env(door_args.tar_env, door_args.seed)\n", + "door_src = [construct_env(e, door_args.seed)[0] for e in door_args.src_envs]\n", "door_q = QAvatar(door_args, door_src, door_env, device)\n", "door_s = SAC(door_args, door_env.observation_space.shape[0], door_env.action_space.shape[0],\n", " door_env.action_space.low, door_env.action_space.high, device)\n", "print(\"target obs/act:\", door_env.observation_space.shape, door_env.action_space.shape)" - ], - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "code", + "execution_count": null, "metadata": {}, + "outputs": [], "source": [ "door_q_res = run_demo_training(door_q, door_args, door_env, DEMO_STEPS_DOOR)\n", "door_s_res = run_demo_training(door_s, door_args, door_env, DEMO_STEPS_DOOR)\n", "plot_curves(door_q, door_s, door_q_res, door_s_res, \" — Door Panda→UR5e\")" - ], - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", @@ -377,24 +377,24 @@ }, { "cell_type": "code", + "execution_count": null, "metadata": {}, + "outputs": [], "source": [ "door_q.load(os.path.join(CDRL_ASSETS, \"checkpoints\", \"door_panda_to_ur5e\", \"QAvatar\", \"steps_500000.pt\"))\n", "door_s.load(os.path.join(CDRL_ASSETS, \"checkpoints\", \"door_panda_to_ur5e\", \"SAC\", \"steps_500000.pt\"))\n", "rollout_compare(door_q, door_s, door_args, \"output/videos/mt07_door_panda_to_ur5e.mp4\",\n", - " max_steps=150, camera=\"frontview\")" - ], - "execution_count": null, - "outputs": [] + " max_steps=500, camera=\"frontview\")" + ] }, { "cell_type": "code", + "execution_count": null, "metadata": {}, + "outputs": [], "source": [ "Video(url=\"output/videos/mt07_door_panda_to_ur5e.mp4\")" - ], - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", @@ -407,13 +407,13 @@ }, { "cell_type": "markdown", + "id": "53b14051", "metadata": {}, "source": [ "## Acknowledgements\n", "\n", "This notebook was developed in collaboration with the lab of Professor Ping-Chun Hsieh (謝秉均), Department of Computer Science, National Yang Ming Chiao Tung University (NYCU) — [pinghsieh.github.io](https://pinghsieh.github.io/)." - ], - "id": "53b14051" + ] }, { "cell_type": "markdown", @@ -438,4 +438,4 @@ }, "nbformat": 4, "nbformat_minor": 5 -} \ No newline at end of file +} diff --git a/projects/MuJoCo-Torch/cdrl_code/utils.py b/projects/MuJoCo-Torch/cdrl_code/utils.py index ee495de..0219dc8 100644 --- a/projects/MuJoCo-Torch/cdrl_code/utils.py +++ b/projects/MuJoCo-Torch/cdrl_code/utils.py @@ -169,7 +169,7 @@ def evaluate(self, agent, steps, writer=None, verbose=True): tqdm.write(f"[Eval] step={steps} | rew={eval_rews:.2f} | len={eval_lens:.1f}") return eval_rews, eval_lens, sucess -def construct_env(env_name, render_mode=None): +def construct_env(env_name, seed=0, render_mode=None): import os import gymnasium as gym @@ -227,6 +227,7 @@ def step(self, action): suite.make( task, robots=robot, + seed=seed, reward_shaping=True, use_camera_obs=False, use_object_obs=True,