Skip to content

chg MuJoCo-Torch/MT06, MT07, and utils.py #172

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
sonyyang-tw merged 1 commit into from
Jul 10, 2026
+115 −114
Merged

chg MuJoCo-Torch/MT06, MT07, and utils.py #172

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
118 changes: 59 additions & 59 deletions projects/MuJoCo-Torch/MT06_RL_PPO.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,9 @@
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"os.environ[\"MUJOCO_GL\"] = \"egl\"\n",
Expand Down Expand Up @@ -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",
Expand All @@ -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",
Expand All @@ -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",
Expand All @@ -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",
Expand All @@ -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",
Expand Down Expand Up @@ -247,10 +250,7 @@
" \"b_values\": b_values,\n",
" \"b_returns\": b_returns,\n",
" }"
],
"execution_count": null,
"outputs": [],
"id": "837b707c"
]
},
{
"cell_type": "markdown",
Expand All @@ -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",
Expand Down Expand Up @@ -468,9 +470,7 @@
" writer.add_scalar(\"charts/SPS\", metrics[\"SPS\"], step)\n",
"\n",
" return metrics"
],
"execution_count": null,
"outputs": []
]
},
{
"cell_type": "markdown",
Expand All @@ -483,7 +483,9 @@
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"args_dict = {\n",
" \"seed\": 0,\n",
Expand Down Expand Up @@ -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",
Expand All @@ -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",
Expand All @@ -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",
Expand All @@ -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",
Expand All @@ -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",
Expand Down Expand Up @@ -653,9 +655,7 @@
"\n",
"# Run training.\n",
"eval_results = run_training(agent, args, envs, log_path)"
],
"execution_count": null,
"outputs": []
]
},
{
"cell_type": "markdown",
Expand All @@ -668,7 +668,9 @@
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import matplotlib.pyplot as plt\n",
"\n",
Expand Down Expand Up @@ -727,9 +729,7 @@
"\n",
"\n",
"plot_ppo_curves(agent, eval_results)"
],
"execution_count": null,
"outputs": []
]
},
{
"cell_type": "markdown",
Expand All @@ -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",
Expand Down Expand Up @@ -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",
Expand All @@ -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",
Expand All @@ -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",
Expand All @@ -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",
Expand All @@ -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",
Expand All @@ -874,4 +874,4 @@
},
"nbformat": 4,
"nbformat_minor": 5
}
}
Loading