Add history encoder ablation option

BotRunner64 · BotRunner64 · commit 4470ff7e67eb · 2026-07-06T11:35:56.000+08:00
diff --git a/tests/test_train_script.py b/tests/test_train_script.py
@@ -12,7 +12,13 @@
 import numpy as np
 import pytest
 
-from train_mimic.app import DEFAULT_TASK, validate_checkpoint_path, validate_motion_file
+from train_mimic.app import (
+    DEFAULT_TASK,
+    HISTORY_ENCODER_TEMPORAL_CNN,
+    apply_history_encoder_config,
+    validate_checkpoint_path,
+    validate_motion_file,
+)
 from train_mimic.data.dataset_lib import PRECOMPUTED_MOTION_VERSION
 from train_mimic.scripts import train
 from train_mimic.tasks.tracking.config.rl import make_general_tracking_ppo_runner_cfg
@@ -42,6 +48,7 @@ def _args(**overrides: object) -> argparse.Namespace:
         "rewind_prob": None,
         "rewind_min_steps": None,
         "rewind_max_steps": None,
+        "history_encoder": HISTORY_ENCODER_TEMPORAL_CNN,
         "device": None,
         "gpu_ids": None,
         "master_port": 29500,
@@ -57,6 +64,7 @@ class TestTrainLauncherHelpers:
     def test_parse_args_defaults_to_tensorboard_logger(self) -> None:
         args = train.parse_args([])
         assert args.logger == "tensorboard"
+        assert args.history_encoder == HISTORY_ENCODER_TEMPORAL_CNN
 
     def test_parse_args_accepts_logger_choice(self) -> None:
         args = train.parse_args(["--logger", "swanlab"])
@@ -75,6 +83,10 @@ def test_parse_args_accepts_robot_xml(self) -> None:
         args = train.parse_args(["--robot_xml", "assets/robots/unitree_g1/g1_29dof_dex3.xml"])
         assert args.robot_xml == "assets/robots/unitree_g1/g1_29dof_dex3.xml"
 
+    def test_parse_args_accepts_history_encoder_none(self) -> None:
+        args = train.parse_args(["--history_encoder", "none"])
+        assert args.history_encoder == "none"
+
     def test_should_launch_multi_gpu(self) -> None:
         args = _args(gpu_ids=[0, 1, 2, 3])
         assert train._should_launch_multi_gpu(args, env={"WORLD_SIZE": "1"}) is True
@@ -219,6 +231,7 @@ def test_configure_swanlab_logger_syncs_tensorboard(self, monkeypatch: pytest.Mo
                         "rewind_prob": 0.8,
                         "rewind_min_steps": 25,
                         "rewind_max_steps": 75,
+                        "history_encoder": HISTORY_ENCODER_TEMPORAL_CNN,
                     },
                 },
             ),
@@ -250,6 +263,50 @@ def test_tracking_runner_configs_disable_model_upload() -> None:
     assert make_general_tracking_ppo_runner_cfg().upload_model is False
 
 
+def test_apply_history_encoder_none_switches_runner_to_current_frame_mlp() -> None:
+    actor_marker = object()
+    critic_marker = object()
+    env_cfg = types.SimpleNamespace(
+        observations={
+            "actor": actor_marker,
+            "critic": critic_marker,
+            "actor_history": object(),
+            "critic_history": object(),
+        }
+    )
+    agent_cfg = make_general_tracking_ppo_runner_cfg()
+
+    apply_history_encoder_config(env_cfg, agent_cfg, "none")
+
+    assert env_cfg.observations == {
+        "actor": actor_marker,
+        "critic": critic_marker,
+    }
+    assert agent_cfg.history_encoder == "none"
+    assert agent_cfg.obs_groups == {
+        "actor": ("actor",),
+        "critic": ("critic",),
+    }
+    assert agent_cfg.actor.class_name == "rsl_rl.models.mlp_model:MLPModel"
+    assert agent_cfg.actor.cnn_cfg is None
+    assert agent_cfg.critic.class_name == "rsl_rl.models.mlp_model:MLPModel"
+    assert agent_cfg.critic.cnn_cfg is None
+
+
+def test_apply_history_encoder_temporal_cnn_keeps_default_runner_cfg() -> None:
+    env_cfg = types.SimpleNamespace(observations={"actor_history": object()})
+    agent_cfg = make_general_tracking_ppo_runner_cfg()
+    actor_class_name = agent_cfg.actor.class_name
+    obs_groups = agent_cfg.obs_groups
+
+    apply_history_encoder_config(env_cfg, agent_cfg, HISTORY_ENCODER_TEMPORAL_CNN)
+
+    assert env_cfg.observations.keys() == {"actor_history"}
+    assert agent_cfg.history_encoder == HISTORY_ENCODER_TEMPORAL_CNN
+    assert agent_cfg.actor.class_name == actor_class_name
+    assert agent_cfg.obs_groups == obs_groups
+
+
 def test_make_g1_training_robot_cfg_uses_requested_xml() -> None:
     from train_mimic.tasks.tracking.config.env import make_g1_training_robot_cfg
 
diff --git a/train_mimic/app.py b/train_mimic/app.py
@@ -14,6 +14,10 @@
 from train_mimic.data.dataset_lib import find_precomputed_motion_shards, validate_precomputed_motion_dataset
 
 DEFAULT_TASK = GENERAL_TRACKING_TASK
+HISTORY_ENCODER_TEMPORAL_CNN = "temporal_cnn"
+HISTORY_ENCODER_NONE = "none"
+HISTORY_ENCODER_CHOICES = (HISTORY_ENCODER_TEMPORAL_CNN, HISTORY_ENCODER_NONE)
+_MLP_MODEL_CLASS = "rsl_rl.models.mlp_model:MLPModel"
 
 
 def validate_motion_file(motion_file: str) -> None:
@@ -93,6 +97,37 @@ def build_runner_cfg_dict(agent_cfg: Any, *, force_tensorboard: bool = False) ->
     return agent_dict
 
 
+def apply_history_encoder_config(
+    env_cfg: Any,
+    agent_cfg: Any,
+    history_encoder: str,
+) -> None:
+    """Apply the training/eval observation-model variant for history ablations."""
+    if history_encoder not in HISTORY_ENCODER_CHOICES:
+        raise ValueError(
+            f"Unsupported history_encoder={history_encoder!r}. "
+            f"Supported values are: {', '.join(HISTORY_ENCODER_CHOICES)}."
+        )
+
+    agent_cfg.history_encoder = history_encoder
+    if history_encoder == HISTORY_ENCODER_TEMPORAL_CNN:
+        return
+
+    observations = getattr(env_cfg, "observations", None)
+    if observations is not None:
+        observations.pop("actor_history", None)
+        observations.pop("critic_history", None)
+
+    agent_cfg.obs_groups = {
+        "actor": ("actor",),
+        "critic": ("critic",),
+    }
+    agent_cfg.actor.class_name = _MLP_MODEL_CLASS
+    agent_cfg.actor.cnn_cfg = None
+    agent_cfg.critic.class_name = _MLP_MODEL_CLASS
+    agent_cfg.critic.cnn_cfg = None
+
+
 def resolve_device(requested_device: str | None, torch_module: Any) -> str:
     if requested_device is not None:
         return requested_device
diff --git a/train_mimic/scripts/benchmark.py b/train_mimic/scripts/benchmark.py
@@ -31,6 +31,9 @@
 
 from train_mimic.app import (
     DEFAULT_TASK,
+    HISTORY_ENCODER_CHOICES,
+    HISTORY_ENCODER_TEMPORAL_CNN,
+    apply_history_encoder_config,
     build_runner_cfg_dict,
     import_training_stack,
     load_task_components,
@@ -71,6 +74,16 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
         default=DEFAULT_TASK,
         help="Task id to benchmark (default: %(default)s)",
     )
+    parser.add_argument(
+        "--history_encoder",
+        type=str,
+        default=HISTORY_ENCODER_TEMPORAL_CNN,
+        choices=HISTORY_ENCODER_CHOICES,
+        help=(
+            "Policy history encoder variant used by the checkpoint. Use none for "
+            "MLP checkpoints trained without actor_history/critic_history."
+        ),
+    )
     return parser.parse_args(argv)
 
 
@@ -359,6 +372,7 @@ def main(argv: Sequence[str] | None = None) -> int:
         load_rl_cfg=_load_rl_cfg,
         load_runner_cls=_load_runner_cls,
     )
+    apply_history_encoder_config(base_env_cfg, agent_cfg, args.history_encoder)
     base_env_cfg.commands["motion"].motion_file = args.motion_file
     benchmark_env_cfg = _configure_benchmark_env_cfg(
         base_env_cfg,
@@ -429,6 +443,7 @@ def main(argv: Sequence[str] | None = None) -> int:
             "motion_file": args.motion_file,
             "seed": args.seed,
             "num_envs": args.num_envs,
+            "history_encoder": args.history_encoder,
         },
         plan=plan,
         results=results,
diff --git a/train_mimic/scripts/play.py b/train_mimic/scripts/play.py
@@ -28,10 +28,19 @@
 
 import argparse
 import os
+import sys
+from pathlib import Path
+
+REPO_ROOT = Path(__file__).resolve().parents[2]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
 
 from mjlab.viewer import NativeMujocoViewer, ViserPlayViewer
 from train_mimic.app import (
     DEFAULT_TASK,
+    HISTORY_ENCODER_CHOICES,
+    HISTORY_ENCODER_TEMPORAL_CNN,
+    apply_history_encoder_config,
     build_runner_cfg_dict,
     import_training_stack,
     load_task_components,
@@ -54,6 +63,16 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--device", type=str, default=None)
     parser.add_argument("--task", type=str, default=DEFAULT_TASK,
                         help="Task id to play (default: %(default)s)")
+    parser.add_argument(
+        "--history_encoder",
+        type=str,
+        default=HISTORY_ENCODER_TEMPORAL_CNN,
+        choices=HISTORY_ENCODER_CHOICES,
+        help=(
+            "Policy history encoder variant used by the checkpoint. Use none for "
+            "MLP checkpoints trained without actor_history/critic_history."
+        ),
+    )
     return parser.parse_args()
 
 
@@ -90,6 +109,7 @@ def main() -> None:
     )
 
     # Override for playback
+    apply_history_encoder_config(env_cfg, agent_cfg, args.history_encoder)
     env_cfg.scene.num_envs = args.num_envs
     env_cfg.commands["motion"].motion_file = args.motion_file
 
diff --git a/train_mimic/scripts/train.py b/train_mimic/scripts/train.py
@@ -23,6 +23,11 @@
         --motion_file data/datasets_precomputed \
         --logger swanlab
 
+    # Ablate temporal history encoder
+    python train_mimic/scripts/train.py \
+        --history_encoder none \
+        --motion_file data/datasets_precomputed
+
     # Resume for additional iterations
     python train_mimic/scripts/train.py \
         --resume logs/rsl_rl/g1_general_tracking/<run>/model_12000.pt \
@@ -40,10 +45,18 @@
 import sys
 import time
 from datetime import datetime
+from pathlib import Path
 from typing import Any, Sequence
 
+REPO_ROOT = Path(__file__).resolve().parents[2]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
 from train_mimic.app import (
     DEFAULT_TASK,
+    HISTORY_ENCODER_CHOICES,
+    HISTORY_ENCODER_TEMPORAL_CNN,
+    apply_history_encoder_config,
     build_runner_cfg_dict,
     import_training_stack,
     load_task_components,
@@ -106,6 +119,17 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
                         help="Minimum policy steps to rewind for rewind sampling")
     parser.add_argument("--rewind_max_steps", type=int, default=None,
                         help="Maximum policy steps to rewind for rewind sampling")
+    parser.add_argument(
+        "--history_encoder",
+        type=str,
+        default=HISTORY_ENCODER_TEMPORAL_CNN,
+        choices=HISTORY_ENCODER_CHOICES,
+        help=(
+            "History observation encoder variant. temporal_cnn uses actor_history/"
+            "critic_history with the TemporalCNN model; none uses current-frame "
+            "actor/critic observations with a single-input MLP."
+        ),
+    )
     parser.add_argument("--device", type=str, default=None)
     parser.add_argument(
         "--gpu_ids",
@@ -307,6 +331,9 @@ def _configure_experiment_logger(
             "rewind_prob": env_cfg.commands["motion"].rewind_prob,
             "rewind_min_steps": env_cfg.commands["motion"].rewind_min_steps,
             "rewind_max_steps": env_cfg.commands["motion"].rewind_max_steps,
+            "history_encoder": getattr(
+                agent_cfg, "history_encoder", HISTORY_ENCODER_TEMPORAL_CNN
+            ),
         },
     )
     swanlab.sync_tensorboard_torch(types=["scalar", "scalars", "image", "text"])
@@ -404,6 +431,7 @@ def _handle_shutdown(signum: int, _frame: Any) -> None:
         agent_cfg.max_iterations = args.max_iterations
     if args.experiment_name is not None:
         agent_cfg.experiment_name = args.experiment_name
+    apply_history_encoder_config(env_cfg, agent_cfg, args.history_encoder)
 
     device = _resolve_device(args, torch)
 
