code/tempsamp_r1/src/tempsamp_main.py

"""TempSamp-R1 (no-CoT) baseline on ActivityForensics.

Why this file is tiny
---------------------
forensics_grpo's `Qwen2VLGRPOTrainer_Video_GT_Soft` already implements every
algorithmic component of TempSamp-R1 paper section 3:

  - Mix-policy sampling (G-1 on-policy rollouts + 1 ground-truth off-policy
    solution): grpo_trainer_video_GT_soft.py lines ~509, 786-787.
  - Non-linear reward shaping (paper Eq. 4, τ=0.8, α1=0.01, α2=1):
    `transform_rewards` at grpo_trainer_video_GT_soft.py:692-704.
  - Group-relative advantage with normalization: lines 1102-1106 (default
    path when no forensics-specific decomposition is enabled).

So this baseline = forensics_grpo trainer + no-CoT prompt + only IoU reward
+ all forensics-specific extras disabled.

This file only:
  1. Sets FORENSICS_COT=false so the trainer picks QUESTION_TEMPLATE_NOCOT.
  2. Asserts that no forensics-specific advantage modifier is enabled.
  3. Delegates to forensics_grpo's `grpo_forensics.main`.

Reward stays whatever is passed via --reward_funcs on the CLI (the launch
script restricts to a single IoU reward).
"""
import os
import sys

# Lock in no-CoT BEFORE importing the trainer so the prompt template selector
# (grpo_trainer_video_GT_soft.get_question_template) reads the right value.
os.environ.setdefault("FORENSICS_COT", "false")

# Hard-fail if any forensics-specific extension is enabled — those break the
# TempSamp-R1 algorithmic equivalence we are trying to preserve.
_FORBIDDEN = [
    "FORENSICS_DECOMP_ADV",
    "FORENSICS_HUNGARIAN_DECOMP",
    "FORENSICS_CTA_GRPO",
    "FORENSICS_FORGERY_AWARE",
    "FORENSICS_USE_EXTERNAL_VERIFIER",
    "FORENSICS_VERIFIER_CONTEXT",
]
for _k in _FORBIDDEN:
    _v = os.environ.get(_k, "false").lower()
    if _v in ("true", "1", "yes"):
        raise RuntimeError(
            f"TempSamp-R1 baseline requires {_k}=false (got {_v!r}). "
            "Forensics-specific advantage/reward extensions would diverge "
            "from the paper algorithm."
        )

FORENSICS_GRPO_ROOT = os.environ.get(
    "FORENSICS_GRPO_ROOT", "/mnt/local-fast/zhangt/forensics_grpo"
)
sys.path.insert(0, FORENSICS_GRPO_ROOT)

from trl import GRPOConfig, ModelConfig, TrlParser  # noqa: E402
from src.open_r1.grpo_forensics import GRPOScriptArguments, main  # noqa: E402


if __name__ == "__main__":
    parser = TrlParser((GRPOScriptArguments, GRPOConfig, ModelConfig))
    script_args, training_args, model_args = parser.parse_args_and_config()
    main(script_args, training_args, model_args)