Make Terra PPO actor unmasked by default

Author: Idate96

README.md

@@ -96,7 +96,7 @@ action, max-step, and reward settings.
 | Entropy schedule | `ent_schedule_start=0.15`, `ent_schedule_end=0.005`, `ent_schedule_steps=9500` |
 | Logging/eval/checkpoints | `--log_train_interval 1`, `--log_eval_interval 100`, `--checkpoint_interval 100`, `--eval_episodes 100` |
 | Reward multipliers | `dump_bonus_mult=0.5`, `excavator_relocate_dumped_mult=1.5`, `excavator_relocate_dug_dirt_mult=1.5` |
-| Action mask | Disabled by default; opt in only with `--enable_action_mask` for the coarse action-availability mask while full-shape masked PPO is still experimental |
+| Action mask | Disabled in PPO training/eval; `action_mask` remains an observation/debug signal for diagnostics |
 
 ### Available Presets
 

docs/EXPERIMENTS_LOG.md

@@ -4,11 +4,62 @@
 
 Active Slurm jobs:
 
+- `66756388`: running 120h-queue four-GPU ResMap64 R1/R2 terminal-reward-fix replacement from paired `codex/mask-speedup-wip` worktrees.
+  - Script: `scripts/euler/terra_train_resmap64_phase_4gpu.sbatch`
+  - Shape: `4` GPUs requested, `1024` envs/GPU, `4096` envs total, `num_steps=32`, `num_minibatches=32`, `total_timesteps=50000000000`.
+  - Architecture: `--map_encoder resnet_delayed --map_feature_dim 128 --use_map_derived_channels --separate_actor_critic_trunks`.
+  - Masking: PPO actor is unmasked by construction; current code has no action-mask train/eval switch.
+  - R1/R2 flags: `--use_critic_affordances --include_episode_progress --edge_features_dim 10`.
+  - Timeout semantics: pre-reset `final_observation` bootstrap for max-step timeouts; true task terminals do not bootstrap; recursive GAE stops at reset boundaries; initial `env_steps`/`episode_progress` are randomized once after startup reset to desynchronize max-step timeouts; terminal success reward is paid only on `task_done`, not on max-step timeout.
+  - W&B: online full run `04e8dada`, `https://wandb.ai/aless-weber-eth/mixed-agents/runs/04e8dada`, entity `aless-weber-eth`, project `mixed-agents`.
+  - Local gates on `2026-05-16 12:07 CEST`: `py_compile`, CPU `env-episode-progress` high-completion timeout terminal-zero validation, CPU `training-accounting`, CPU `synthetic-env-action-mask`, and two-agent review passed.
+  - Local 4090 smoke result: one update, `1024` envs, `num_steps=32`, `num_minibatches=32`, ResMap64 + critic affordances, `432.00s`, `79.36` steps/s.
+  - Current status on `2026-05-16 12:34 CEST`: running on `eu-g6-072` in `gpuhe.120h`, `5-00:00:00`, with `4 x NVIDIA GeForce RTX 4090`.
+  - Euler gates: GPU guard passed, CUDA/cuDNN/NCCL preflight passed, and the W&B-disabled full-shape smoke completed one update in `416.07s` at `315.15` steps/s.
+  - First online row, W&B step `0`: `entropy=2.0771589` (`ln(8)=2.07944`, not collapsed), `value_loss=0.0018877`, `explained_variance=0.0001058`, `sched/entropy_coef=0.15`, `train/timeout_rate=train/done_rate=0.002037`, `train/task_done_rate=0`, `eval/success_rate=0`, `eval/positive_terminations=0`, `eval/failure/rewards/terminal=0`, and `eval/rewards/terminal=0`.
+  - Update `100` health on `2026-05-16 12:46 CEST`: `entropy=1.9799`, `value_loss=0.00231`, `explained_variance=0.954`, `train/timeout_rate=0.0017166`, `train/task_done_rate=0`, `eval/success_rate=0`, `eval/max_reward=0.86875`, `eval/rewards=-0.00295`, and terminal reward on eval failures remains `0`.
+  - Update `100` local rollout probe: `docs/policy_diagnostics/04e8dada_update100_rollout_32x550_seed0_1.json`, `32` envs, `550` steps, seeds `0` and `1`, unmasked plus forced-mask. Result: `0/128` total successes across all mode/seed rollouts, invalid selections are low or zero, forced masking does not change success, final dig coverage is about `0.07-0.08`, final dump coverage is about `0.22-0.28`, and the policy still uses sparse `DO` with high `DO_NOTHING`.
+  - Update `500` health on `2026-05-16 13:36 CEST`: W&B around step `536` has `entropy=1.8695`, `value_loss=0.00244`, `explained_variance=0.9907`, `train/timeout_rate=0.001724`, `train/task_done_rate=0`, `eval/success_rate=0`, `eval/max_reward=0.823`, `eval/rewards=-0.00046`, `eval/DO=0.0725`, `eval/DO_NOTHING %=0.156`, and terminal reward on eval failures remains `0`.
+  - Update `500` local rollout probe: `docs/policy_diagnostics/04e8dada_update500_rollout_32x550_seed0_1.json`. Result: still `0` successes, but terrain changed in `30/32` envs in every seed/mode, terrain-change/action rose to `0.0146-0.0159`, final dig coverage rose to about `0.50`, final dump coverage rose to about `0.93-0.94`, returns improved versus update `100`, and max returns became positive up to `2.46`. Unmasked logits use more `DO` (`4.6-5.6%`) but include `4-5%` invalid selections; forced mask removes invalid selections but pushes more mass to `DO_NOTHING`.
+  - Update `900-1000` health on `2026-05-16 14:52 CEST`: W&B now has tiny nonzero eval success rows at step `600` (`0.000244`), step `900` (`0.000244`, `eval/max_reward=4.59`), and step `1000` (`0.000977`, `eval/max_reward=5.22`).
+  - Update `1100` eval on `2026-05-16 15:08 CEST`: success returned to `0`, `eval/max_reward=0.847`, `eval/rewards=0.000135`, `eval/DO=0.0886`, and `eval/DO_NOTHING %=0.155`. Completion still trends up (`completion=0.6467`, `core=0.8069`, `edge=0.5855` on the eval row; latest summary completion `0.666`). Treat the `600/900/1000` successes as rare discovery, not persistent learning yet.
+  - Update `1200` eval on `2026-05-16 15:18 CEST`: tiny success returned (`eval/success_rate=0.000244`, `eval/max_reward=5.45`, `eval/rewards=0.000303`, `eval/DO=0.0829`, `eval/DO_NOTHING %=0.153`).
+  - Update `1300` eval/latest summary at step `1320`: tiny success rose to `0.001709`, `eval/max_reward=5.54`, `eval/rewards=0.000518`, `eval/DO=0.0820`, `eval/DO_NOTHING %=0.159`, `entropy=1.8642`, `value_loss=0.00148`, `explained_variance=0.9956`, `train/timeout_rate=0.00182`, `train/task_done_rate=0`, `progress/completion=0.6711`, `core=0.8164`, and `edge=0.6129`.
+  - Latest W&B summary at step `1474`: tiny success remains nonzero at `0.000732`, `eval/max_reward=5.44`, `eval/rewards=0.000463`, `eval/DO=0.0880`, `eval/DO_NOTHING %=0.155`, `entropy=1.8593`, `value_loss=0.00142`, `explained_variance=0.9953`, `train/timeout_rate=0.00168`, `train/task_done_rate=0`, `progress/completion=0.6883`, `core=0.8257`, and `edge=0.6370`.
+  - Value/termination history through W&B step `1053`: no return of the old synchronized timeout spike pattern. `value_loss` median is `0.00185`, p95 `0.00264`, max `0.00627`; `explained_variance` median is `0.992`; `train/timeout_rate` stays near `1/550` (`0.00153-0.00208`); `train/task_done_rate` is mostly `0` with rare `1.5e-5` rows.
+  - Update `1000` local rollout probe: `docs/policy_diagnostics/04e8dada_update1000_rollout_32x550_seed0_1.json`. Result: still `0/128` successes across unmasked/forced-mask seeds, so local deterministic completions are not reliable yet. Several behavior metrics moved in the useful direction versus update `500`: unmasked `DO` is about `7.4%`, invalid raw selections are down to about `1.9-2.6%`, terrain-change/action rose to `0.021-0.025`, and final dig coverage rose to about `0.59-0.63`; terrain changed in `27-29/32` envs and max returns remain positive up to `1.56`. Forced masking still removes invalid selections but lowers `DO` to about `4.9%` and increases `DO_NOTHING` to about `20-21%`, so re-enabling a naive action mask is not yet the obvious rescue.
+  - Update `1000` failure summary: final undug fraction is still high at about `0.37-0.41`. Positive moved dirt exists in `26/32` seed-0 failures and `29/32` seed-1 failures; within those positive-dirt failures, only about `1.0-2.7%` of moved dirt is off the dump zones, and all positive-dirt failures are below the simple `<=20%` off-dump threshold. `8-10/32` envs are still loaded at timeout. Current local evidence points more to unfinished digging/sequence completion than to wrong-place dumping, but it does not prove every failed env completed dumping.
+  - Latest rolling-checkpoint probe at W&B step about `1258`: `docs/policy_diagnostics/04e8dada_latest_step1258_rollout_32x550_seed0_1.json`, `0/128` successes. It improves the final state versus update `1000`: avg returns are about `-0.34` to `-0.06`, final dig coverage is about `0.65-0.70`, final dump coverage about `0.89-0.97`, final undug fraction about `0.30-0.35`, off-dump moved dirt about `0.4-1.4%`, and loaded-at-timeout is `0-1/32`. The concern is action distribution: unmasked `DO` fell to `4.1-4.8%`, forced-mask `DO` fell to `2.2-2.3%`, and `DO_NOTHING` rose to `27-43%`, so the policy is improving final maps but may be drifting toward too much waiting.
+  - Stochastic rolling-checkpoint probe: `docs/policy_diagnostics/04e8dada_latest_step1258_stochastic_rollout_32x550_seed0_1.json`, unmasked only, `0/64` successes. This better matches W&B eval action sampling: `DO=8.0-8.1%`, `DO_NOTHING=15.2-15.7%`, entropy `1.88`, invalid selections about `7.1-7.3%`. Final maps are closer than argmax: final dig coverage `0.74-0.77`, dump coverage `0.995`, final undug fraction `0.23-0.26`, no loaded-at-timeout failures, and off-dump moved dirt `0.4-0.5%`. The W&B-comparable bottleneck is still finishing the remaining dig, not dumping placement or entropy collapse.
+  - Clean-baseline stochastic comparison: `docs/policy_diagnostics/clean_baseline_stochastic_rollout_32x550_seed0_1.json` succeeds `12-14/32` per seed on the same local probe. It reaches final dig coverage `0.89-0.90`, final undug fraction `0.095-0.112`, terrain-change/action `0.058-0.064`, low invalid selection `0.3-0.8%`, and near-deterministic entropy `0.03-0.04`. Current `04e8dada` has similar raw `DO` in stochastic mode, but only about half the productive terrain-change/action and much more remaining dig.
+  - Early checkpoint preservation: `train_mixed.py` now keeps additional snapshots at updates `0`, `100`, `500`, `1000`, `2000`, `5000`, and `10000` in addition to the rolling checkpoint.
+  - Logs once started: `/cluster/home/lterenzi/codex_terra_edge_validation/logs/66756388_resmap64_phase_4gpu.out`.
+  - Next health checks: monitor update `2000`, whether the tiny `600/900/1000` success rows become persistent, value-loss/explained-variance spikes, entropy trend, `DO` rising toward historical healthy levels, `DO_NOTHING` falling, and early checkpoints at updates `2000`, `5000`, and `10000`.
+
+Recently cancelled:
+
+- `66670662`: online W&B four-GPU ResMap64 R1/R2 staggered-timeout run, W&B `xxf7eoap`.
+  - Cancelled on `2026-05-16 12:05 CEST` after `12:58:11` on `eu-g4-019`.
+  - Reason: latest-checkpoint diagnosis found zero success despite timeout desync; W&B showed nonzero `eval/failure/rewards/terminal` while `eval/success_rate=0`, and code paid terminal reward on `done` where `done` includes timeouts. Replaced by `66756388` with terminal reward gated on `task_done`.
+
+- `66591675`: online W&B four-GPU ResMap64 R1/R2 combined-fix run, W&B `faen4uin`.
+  - Cancelled on `2026-05-15 21:41 CEST` after `06:02:00` on `eu-g4-027`.
+  - Reason: superseded by the staggered-start timeout fix. Diagnosis before cancellation: value-loss/explained-variance spikes remained timeout-synchronized; spike updates had `timeout_rate ~= 1/32`, `task_done_rate=0`, and high error in `train/value_loss_timeout_bucket_75_100`.
+
+- `66660010`: first staggered-timeout 120h replacement.
+  - Failed on `2026-05-15 22:11 CEST` after `00:00:48` on `eu-g4-027`.
+  - Reason: partial code sync left remote `train.py` older than `train_mixed.py`, causing `ImportError: cannot import name 'eval_component_metrics' from 'train'`.
+  - Fix: resynced both paired WIP trees to Euler and submitted replacement `66670662`.
+
+- `66438404`: `terra-clean-ma-4x4090`, W&B/checkpoint name `terra-clean-multiagent-4x4090-autotune0-euler-pr-2026-05-13-19-49-55`.
+  - Cancelled on `2026-05-15 15:16 CEST` after `1-19:26:22` on `eu-g4-031`.
+  - Reason: run appeared saturated and freeing its four RTX 3090 GPUs/fairshare should help the pending R1/R2 launch.
+
 - `66574910`: online W&B four-GPU residual-map architecture run from paired `codex/mask-speedup-wip` worktrees.
   - Script: `scripts/euler/terra_train_resmap64_phase_4gpu.sbatch`
   - Shape: `4` GPUs requested, `1024` envs/GPU, `4096` envs total, `num_steps=32`, `total_timesteps=50000000000`.
   - Architecture: `--map_encoder resnet_delayed --map_feature_dim 128 --use_map_derived_channels --separate_actor_critic_trunks`.
-  - Masking: explicit `--disable_action_mask`; this is an architecture run, not a mask run.
+  - Masking: unmasked PPO actor; this is an architecture run, not a mask run.
   - PPO minibatching: `num_minibatches=32` to keep the residual-map gradient batch at `1024` map frames instead of the failed `2048`.
   - Eval/checkpointing: every `100` updates, `eval_episodes=100`, `num_rollouts_eval=550`.
   - W&B: online, entity `aless-weber-eth`, project `mixed-agents`.
@@ -22,6 +73,20 @@ Active Slurm jobs:
   - W&B run: `d1dkjojl`, `https://wandb.ai/aless-weber-eth/mixed-agents/runs/d1dkjojl`.
   - Euler health: GPU guard passed, CUDA/cuDNN/NCCL preflight passed, full-shape W&B-disabled smoke passed in `422.93s` at `310.04` steps/s, and the full online run is past update `300`.
   - Log note: W&B emits `ERROR:root:Driver not initialized (amdgpu not found in modules)` from system monitoring; this is not a training failure.
+  - Cancelled on `2026-05-15 11:48 CEST` after `03:01:32` on `eu-g6-061`; Slurm state `CANCELLED by 542598`.
+  - Last diagnostic read before cancellation: W&B was around update `998`; periodic value-loss spikes aligned with the episode horizon, eval terminations were still zero, and the run was too early to judge learning.
+
+Combined-fix run lineage:
+
+- Detailed plan: `docs/RESMAP_AFFORDANCE_TIMEOUT_PLAN.md`.
+- Keep the ResMap64 spatial architecture, derived map channels, separate actor/critic trunks, and `num_minibatches=32`.
+- Feed existing `edge_features` plus a new `episode_progress = env_steps / max_steps_in_episode` scalar to the critic trunk only.
+- Add training diagnostics for `task_done_rate`, `timeout_rate`, value loss by timeout bucket, and explained variance by timeout bucket.
+- Fix time-limit handling as a separate training semantics change: bootstrap through max-step truncations, but still zero bootstrap for true task success.
+- Keep actor-visible affordances as a later deployability-gated ablation; anything fed to the actor must be computable on the real robot with the same semantics.
+- `66591675` tested this combined R1/R2 setup without initial timeout staggering. It reached online training but was cancelled when W&B diagnostics showed value-loss spikes phase-locked to max-step timeouts.
+- `66670662` added startup `env_steps`/`episode_progress` randomization and fixed timeout synchronization, but still paid timeout terminal reward.
+- `66756388` is the active replacement and adds the terminal-reward gate fix on top of the same architecture and PPO shape.
 
 Recent failed predecessor:
 

docs/EXPERIMENTS_RUNNING.md

@@ -2,11 +2,7 @@
 import jax
 import math
 import sys
-from utils.models import (
-    infer_edge_features_dim_from_model_params,
-    infer_use_action_mask_from_train_config,
-    load_neural_network,
-)
+from utils.models import load_neural_network, restore_checkpoint_model_config
 from utils.helpers import load_pkl_object
 from terra.env import TerraEnvBatch
 from terra.actions import (
@@ -16,7 +12,6 @@
     TrackedActionType,
 )
 import jax.numpy as jnp
-from utils.action_masking import apply_action_mask
 from utils.utils_ppo import obs_to_model_input, wrap_action
 from train_mixed import MixedAgentTrainConfig
 
@@ -96,8 +91,6 @@ def rollout_episode(
         if model is not None:
             obs_model = obs_to_model_input(timestep.observation, prev_actions, rl_config)
             v, logits_pi = model.apply(model_params, obs_model)
-            if getattr(rl_config, "use_action_mask", False):
-                logits_pi = apply_action_mask(logits_pi, obs_model[22])
             if deterministic:
                 action = np.argmax(logits_pi, axis=-1)
             else:
@@ -268,8 +261,7 @@ def print_stats(
 
     log = load_pkl_object(f"{args.run_name}")
     config = log["train_config"]
-    config.edge_features_dim = infer_edge_features_dim_from_model_params(log["model"])
-    config.use_action_mask = infer_use_action_mask_from_train_config(config, default=False)
+    restore_checkpoint_model_config(config, log["model"])
     # from utils.helpers import load_config
     # config = load_config("agents/Terra/ppo.yaml", 22333, 33222, 5e-04, True, "")["train_config"]
 

eval.py

@@ -22,13 +22,8 @@
 import jax.random as jrandom
 from tensorflow_probability.substrates import jax as tfp
 
-from utils.models import (
-    infer_edge_features_dim_from_model_params,
-    infer_use_action_mask_from_train_config,
-    load_neural_network,
-)
+from utils.models import load_neural_network, restore_checkpoint_model_config
 from utils.helpers import load_pkl_object
-from utils.action_masking import apply_action_mask
 from utils.utils_ppo import obs_to_model_input, wrap_action
 from terra.env import TerraEnvBatch
 from terra.config import BatchConfig
@@ -95,8 +90,6 @@ def make_mcts_step_fn(model, env, config):
 
     def apply_model(params, inp):
         val, logits_pi = model.apply(params, inp)
-        if getattr(config, "use_action_mask", False):
-            logits_pi = apply_action_mask(logits_pi, inp[22])
         pi = tfp.distributions.Categorical(logits=logits_pi)
         return val, pi
 
@@ -319,8 +312,6 @@ def rollout_episode(
             rng, rng_act, rng_step = jrandom.split(rng, 3)
             obs_model = obs_to_model_input(timestep.observation, prev_actions, rl_config)
             v, logits_pi = model.apply(model_params, obs_model)
-            if getattr(rl_config, "use_action_mask", False):
-                logits_pi = apply_action_mask(logits_pi, obs_model[22])
             if deterministic:
                 action = jnp.argmax(logits_pi, axis=-1)
             else:
@@ -708,8 +699,7 @@ def _fmt(x):
     n_envs = args.n_envs
     log = load_pkl_object(f"{args.run_name}")
     config = log["train_config"]
-    config.edge_features_dim = infer_edge_features_dim_from_model_params(log["model"])
-    config.use_action_mask = infer_use_action_mask_from_train_config(config, default=False)
+    restore_checkpoint_model_config(config, log["model"])
     config.num_test_rollouts = n_envs
     config.num_devices = 1
     # MCTS params (only used when --use-mcts)

eval_mcts.py

@@ -7,11 +7,7 @@
 import jax
 import math
 from tqdm import tqdm
-from utils.models import (
-    infer_edge_features_dim_from_model_params,
-    infer_use_action_mask_from_train_config,
-    load_neural_network,
-)
+from utils.models import load_neural_network, restore_checkpoint_model_config
 from utils.helpers import load_pkl_object
 from terra.env import TerraEnvBatch
 from terra.actions import (
@@ -22,7 +18,6 @@
 )
 from terra.config import BatchConfig
 import jax.numpy as jnp
-from utils.action_masking import apply_action_mask
 from utils.utils_ppo import obs_to_model_input, wrap_action
 from terra.state import State
 from tensorflow_probability.substrates import jax as tfp
@@ -149,8 +144,6 @@ def rollout_episode(
         if model is not None:
             obs_model = obs_to_model_input(timestep.observation, prev_actions, rl_config)
             v, logits_pi = model.apply(model_params, obs_model)
-            if getattr(rl_config, "use_action_mask", False):
-                logits_pi = apply_action_mask(logits_pi, obs_model[22])
             if deterministic:
                 action = np.argmax(logits_pi, axis=-1)
             else:
@@ -621,8 +614,7 @@ def _fmt(x):
 
     log = load_pkl_object(f"{args.run_name}")
     config = log["train_config"]
-    config.edge_features_dim = infer_edge_features_dim_from_model_params(log["model"])
-    config.use_action_mask = infer_use_action_mask_from_train_config(config, default=False)
+    restore_checkpoint_model_config(config, log["model"])
     config.num_test_rollouts = n_envs
     config.num_devices = 1