Allow environments to declare preserved state columns

[hubert-marek]

Two fixes for deferred group scoring compatibility.

Description

1. Allow environments to declare preserved state columns

When environments write custom fields into vf.State during post_rollout (e.g. state["eval_score"]), those fields are dropped during state_to_output() serialization unless explicitly requested via state_columns. Currently, the only way to preserve them is for the orchestrator to hardcode environment-specific field names in REQUIRED_STATE_COLUMNS, which couples the orchestrator to individual environment implementations.

This adds an optional state_columns parameter to Environment.__init__ so environments can declare their own required state columns at construction time:

return MyEnv(
    dataset=dataset,
    rubric=rubric,
    state_columns=["eval_score", "eval_error"],
)

The orchestrator (e.g. prime-rl) can then read env.state_columns and include them automatically, without needing to know the field names itself.

Defaults to [].

2. Handle serialized messages in tool monitor rubrics

ToolMonitorRubric used attribute access (msg.role) which fails when completion messages are serialized to plain dicts during deferred group scoring. Added normalize_messages() to convert any serialized dicts back to Message objects before processing, using the existing message_utils helper.

Fixes:

Type of Change

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Documentation update
• Test improvement

Testing

• All existing tests pass when running uv run pytest locally.
• New tests have been added to cover the changes

Checklist

• My code follows the style guidelines of this project as outlined in AGENTS.md
• I have performed a self-review of my own code
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• Any dependent changes have been merged and published

Additional Notes

Companion PR: PrimeIntellect-ai/prime-rl#2057

---

Note

Low Risk
Low risk, backwards-compatible additions: state_columns is optional and defaults to empty, and message normalization only broadens accepted input shapes for tool-monitor metrics.

Overview
Adds an optional state_columns parameter to Environment construction and stores it on env.state_columns, so environments can declare which JSON-serializable State fields should be preserved into RolloutOutput without orchestrator hardcoding.

Fixes deferred/group scoring compatibility for ToolMonitorRubric by normalizing completion messages (e.g., serialized dicts) before counting tool calls. Documentation and API reference are updated to reflect both changes.

^{Written by Cursor Bugbot for commit bd016db. This will update automatically on new commits. Configure here.}

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

^{Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}