MemGym/memgym-rm-scenario-ood-webarena

MemGym-RM Scenario-OOD — WebArena V2

Description

MemGym-RM-Scenario-OOD-WebArena is a held-out evaluation set designed to test whether MemRM generalizes to a completely different agent domain (WebArena browser tasks) not seen during training (which used SWE-Gym software engineering tasks). Each row is a trajectory step from a WebArena agent running under one of three memory cohorts, with a scenario_swap perturbation applied.

OOD axis: Agent scenario / task domain (SWE-Gym → WebArena).

Ships two files (use `pairs_paper_eval.jsonl` to reproduce paper metrics):

The 61-row delta is the contents of one inference shard (1 of 8 planned shards failed to land during the paper run). The split is sticky and label-unbiased — see docs/release/d13_487_vs_426.md (source-only) for the per-perturbation / per-app / per-label / per-prompt-length breakdown.

CRITICAL — which checkpoint backs these numbers: Metrics come from 1.7B QLoRA checkpoint-500. Eval file: training_output/lightweight_comparison/eval_results_1p7b_ckpt500_webarena_v2_ood.json

\* See "Known limitations — AUROC reporting" below for caveats.

Do NOT use reward_model_v2_run1/eval_results.json — that file is the 8B model and does not contain WebArena OOD results.

Schema (post-sanitization)

All fields below reflect the field set after sanitize_hf_jsonl.py removes or redacts the private training_path sub-field from provenance.

Filter Recipe (487 union → 426 paper-eval)

The eval-results JSON has no row_id field; the canonical join key is the (instance_id, step, source_dir) tuple, which is unique across all 487 source rows.

import json

eval_d = json.load(open(
    "training_output/lightweight_comparison/eval_results_1p7b_ckpt500_webarena_v2_ood.json"
))
keep = {(r["instance_id"], int(r["step"]), r["source_dir"])
        for r in eval_d["per_row_predictions"]}
# len(keep) == 426

with open("pairs_union.jsonl") as f_in, \
     open("pairs_paper_eval.jsonl", "w") as f_out:
    for line in f_in:
        d = json.loads(line)
        if (d["instance_id"], int(d["step"]), d["source_dir"]) in keep:
            f_out.write(line)
# Output line count: 426

The 61 dropped rows belong to the single inference shard that failed during evaluation (1 of 8 planned shards). They are structurally valid but were never scored. The shard mapping is sticky — the same 61 rows would drop on a re-run of the same shard split. The pre-upload sanitize sweep (docs/release/sanitize_hf_jsonl.py) strips private training_path / replay_path values from provenance for both files; row counts are preserved.

License

MIT. See the MemGym repository's docs/licenses.md for the full asset-license matrix.

Citation

@inproceedings{xu2026memgym,
  title     = {MemGym: a Long-Horizon Memory Environment for LLM Agents},
  author    = {Anonymous Authors},
  booktitle = {Advances in Neural Information Processing Systems},
  year      = {2026},
  note      = {Under review}
}

Loading the Dataset

from datasets import load_dataset

# Paper-reproduction (426 rows scored by the 1.7B ckpt-500 evaluator):
ds_paper = load_dataset(
    "MemGym/memgym-rm-scenario-ood-webarena",
    data_files="pairs_paper_eval.jsonl", split="train",
)
print(len(ds_paper))  # 426

# Full union (487 rows) — re-score on your own checkpoint:
ds_full = load_dataset(
    "MemGym/memgym-rm-scenario-ood-webarena",
    data_files="pairs_union.jsonl", split="train",
)
print(len(ds_full))  # 487

Regeneration Recipe

See docs/release/provenance.md artifact #8 and docs/release/d13_487_vs_426.md for the full provenance trace.

Step 1 — rebuild the 487-row union from raw WebArena trajectories:

# Inputs (from the Stage 0.5 EC2 backup mirror or original EC2 host):
#   <wa-root>/grid_idfix/<app>_none/         (5 baseline cohorts)
#   <wa-root>/prompt_fix_v2/<app>_<perturbation>/  (15 OOD cohorts)

python -m memgym.training.scripts.build_webarena_long_context_pairs \
    --baseline-root <wa-root>/grid_idfix \
    --ood-root      <wa-root>/prompt_fix_v2 \
    --baseline-cohort-template "{app}_none" \
    --ood-cohorts gitlab_struct_ms10 gitlab_summ_ms10 gitlab_summ_ms15 \
                  gmail_struct_ms10  gmail_summ_ms10  gmail_summ_ms15 \
                  linear_struct_ms10 linear_summ_ms10 linear_summ_ms15 \
                  paypal_struct_ms10 paypal_summ_ms10 paypal_summ_ms15 \
                  superhuman_struct_ms10 superhuman_summ_ms10 superhuman_summ_ms15 \
    --out pairs_union.jsonl
# Expect: 487 rows; HARD RULE pass when checked with probe_v7_hardrule.

Step 2 — derive the 426-row paper-eval subset from the union + eval JSON (see the filter recipe block above).

Step 3 — sanitize provenance paths before HF upload (docs/release/sanitize_hf_jsonl.py).

Known Limitations

• —61-row delta (Stage 1 M1): 1 of 8 planned inference shards did not land. The 61 un-evaluated rows are kept in pairs_union.jsonl but excluded from pairs_paper_eval.jsonl. Re-running the same shard split would drop the same 61 rows.
• —HARMFUL class is minority (~20%). Both files share the same class imbalance: 99/487 = 20.3% HARMFUL in the union; 87/426 = 20.4% HARMFUL in the paper-eval. Use class-weighted CE / threshold sweep when re-scoring; do not interpret raw accuracy without the per-class F1.
• —AUROC reporting nuance. The eval results file reports auroc=0.4248 over all 426 rows. Earlier drafts of this README cited Covered AUROC = 0.748 based on a covered-subset filter; that number could not be reproduced from the eval JSON alone and is therefore omitted pending paper-cross-check. Use the per-row predictions in the eval JSON to recompute under whichever subset you intend to cite.
• —Private path in raw file (Stage 3 D01): The pre-sanitize 487-row source file contained an absolute deploy-host path under provenance.training_path in every row. Both files shipped here are post-sanitize — the absolute paths are rewritten to repo-relative paths and the HF-hosted files contain no private paths.
• —EC2-only inputs: Raw WebArena trajectories live under prompt_fix_v2/ and grid_idfix_*/ on the original EC2 host. A full backup mirror is maintained off-EC2 (Stage 0.5); see the regeneration recipe block for the precise input layout.