rrm workspace: TRM/HRM/SRM code, Maze dataset, dynamical-analysis pipelineHEADmain

Curated export for clone-and-run Maze training (2x A6000) + diagnostics.
trm/hrm pretrain.py carry trajectory-augmentation code (backward-compatible).
Heavy artifacts (checkpoints/wandb/npz) gitignored; see PROVENANCE.md.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

3 files changed, 76 insertions, 0 deletions

diff --git a/env/pip-freeze.txt b/env/pip-freeze.txt
new file mode 100644
index 0000000..03c81b2
--- /dev/null
+++ b/env/pip-freeze.txt
@@ -0,0 +1,18 @@
+adam_atan2==0.0.3
+argdantic==1.3.3
+coolname==2.2.0
+einops==0.8.1
+flash_attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.3/flash_attn-2.8.3+cu12torch2.7cxx11abiTRUE-cp310-cp310-linux_x86_64.whl#sha256=ce91e246f21d61ad66b1a7555340dbaa28e4aa86edcf00c18f0837422939b529
+huggingface-hub==0.34.4
+hydra-core==1.3.2
+numba==0.61.2
+omegaconf==2.3.0
+pydantic==2.11.7
+pydantic_core==2.33.2
+pydantic-settings==2.14.1
+torch==2.7.0+cu126
+torchaudio==2.7.0+cu126
+torchvision==0.22.0+cu126
+tqdm==4.67.1
+triton==3.3.0
+wandb==0.21.3
diff --git a/env/requirements.txt b/env/requirements.txt
new file mode 100644
index 0000000..d03a446
--- /dev/null
+++ b/env/requirements.txt
@@ -0,0 +1,13 @@
+# Merged from HRM + TRM official; versions pinned per TRM specific_requirements.txt
+# torch handled separately (torch==2.7.0+cu126 via cu126 index)
+einops==0.8.1
+tqdm==4.67.1
+coolname==2.2.0
+pydantic==2.11.7
+pydantic-core==2.33.2
+argdantic==1.3.3
+wandb==0.21.3
+omegaconf==2.3.0
+hydra-core==1.3.2
+huggingface_hub==0.34.4
+numba==0.61.2
diff --git a/env/setup.sh b/env/setup.sh
new file mode 100755
index 0000000..6a8435c
--- /dev/null
+++ b/env/setup.sh
@@ -0,0 +1,45 @@
+#!/usr/bin/env bash
+# rrm 统一环境安装脚本
+# 复现 HRM (sapientinc/HRM) + TRM (SamsungSAILMontreal/TinyRecursiveModels) 用
+# 用法: bash env/setup.sh
+set -euo pipefail
+
+ENV_NAME=${ENV_NAME:-rrm}
+PY_VER=${PY_VER:-3.10}
+
+source "$(conda info --base)/etc/profile.d/conda.sh"
+
+if ! conda env list | awk '{print $1}' | grep -qx "$ENV_NAME"; then
+  conda create -n "$ENV_NAME" python="$PY_VER" -y
+fi
+conda activate "$ENV_NAME"
+
+pip install --upgrade pip wheel setuptools packaging ninja setuptools-scm
+
+# Torch 2.7.0 cu126 (TRM specific_requirements pinned)
+pip install torch==2.7.0+cu126 torchvision==0.22.0+cu126 torchaudio==2.7.0+cu126 \
+  --index-url https://download.pytorch.org/whl/cu126
+
+# 合并版 requirements (HRM + TRM)
+pip install -r "$(dirname "$0")/requirements.txt"
+
+# FlashAttention 2 — A6000/Ampere prebuilt wheel (cxx11abi=TRUE 对应 torch 2.7 cu126)
+pip install https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.3/flash_attn-2.8.3+cu12torch2.7cxx11abiTRUE-cp310-cp310-linux_x86_64.whl
+
+# adam-atan2 需要 nvcc 编译; conda 装 cuda-toolkit 提供 nvcc + 头文件
+conda install -y -c nvidia cuda-toolkit=12.6
+CUDA_HOME="$CONDA_PREFIX" pip install --no-cache-dir --no-build-isolation adam-atan2==0.0.3
+
+# wandb 离线 (避免 smoke test 上传)
+export WANDB_MODE=${WANDB_MODE:-offline}
+
+python - <<'PY'
+import torch, flash_attn, adam_atan2
+print("torch:", torch.__version__, "CUDA:", torch.version.cuda)
+print("GPUs:", torch.cuda.device_count(), "->", [torch.cuda.get_device_name(i) for i in range(torch.cuda.device_count())])
+print("flash_attn:", flash_attn.__version__)
+print("adam_atan2 ok")
+PY
+
+echo
+echo "==> rrm env ready. conda activate $ENV_NAME"