nemo_gym.comparison.loading

Reading *_aggregate_metrics.json for gym eval compare, and picking which agent to compare.

All filesystem I/O for the compare feature lives here. The rollouts JSONL a user points at is never opened: it is the run’s identity and the handle its _aggregate_metrics.json sibling is derived from.

Module Contents

Classes

Name	Description
`AgentSelection`	Which agent to read from each side for one comparison block.
`LoadedRun`	One side of the comparison, narrowed to a single agent.

Functions

Name	Description
`_derive_num_repeats`	How many repeats the run collected: the most any single task has.
`_read_agent_entries`	-
`_require_agent`	-
`_sole_agent`	The file’s only agent, or an error asking which of several to use.
`build_loaded_run`	Narrow a parsed run file to one agent.
`load_agg_metrics_file`	Resolve and parse one run’s `*_aggregate_metrics.json`.
`resolve_agent_selections`	Decide which agent to read from each side.
`resolve_aggregate_metrics_fpath`	The `*_aggregate_metrics.json` to read: the explicit override, else the derived sibling.

Data

RunRole

API

class nemo_gym.comparison.loading.AgentSelection(
    baseline_agent: str,
    candidate_agents: typing.Tuple[str, ...]
)

Dataclass

Which agent to read from each side for one comparison block.

baseline_agent

str

candidate_agents

Tuple[str, ...]

class nemo_gym.comparison.loading.LoadedRun(
    agent_name: str,
    agent_metrics: typing.Dict[str, typing.Any],
    key_metrics: typing.Dict[str, typing.Any],
    group_level_metrics: typing.List[typing.Dict[str, typing.Any]],
    num_tasks: int = 0,
    num_repeats: typing.Optional[int] = None,
    has_repeat_cis: bool = False
)

Dataclass

One side of the comparison, narrowed to a single agent.

Only what the diff actually consumes. Run identity (role, paths, label) stays on RunFile, which the report carries directly, so it is deliberately not duplicated here.

agent_metrics

Dict[str, Any]

agent_name

str

group_level_metrics

List[Dict[str, Any]]

has_repeat_cis

bool = False

key_metrics

Dict[str, Any]

num_repeats

Optional[int] = None

num_tasks

int = 0

nemo_gym.comparison.loading._derive_num_repeats(
    group_level_metrics: typing.List[typing.Dict[str, typing.Any]],
    repeat_level_metrics: typing.List[typing.Dict[str, typing.Any]]
) -> typing.Optional[int]

How many repeats the run collected: the most any single task has.

expected_num_rollouts is per task, and a partially recovered run leaves some tasks short of the rest, so the max is the run’s repeat count. repeat_level_metrics has exactly one entry per repeat but is absent from single-repeat runs and from files written before it existed.

nemo_gym.comparison.loading._read_agent_entries(
    metrics_fpath: pathlib.Path
) -> typing.Dict[str, typing.Dict[str, typing.Any]]

nemo_gym.comparison.loading._require_agent(
    run_file: nemo_gym.comparison.schema.RunFile,
    agent_name: str
) -> None

nemo_gym.comparison.loading._sole_agent(
    run_file: nemo_gym.comparison.schema.RunFile
) -> str

The file’s only agent, or an error asking which of several to use.

nemo_gym.comparison.loading.build_loaded_run(
    run_file: nemo_gym.comparison.schema.RunFile,
    agent_name: str
) -> nemo_gym.comparison.loading.LoadedRun

Narrow a parsed run file to one agent.

nemo_gym.comparison.loading.load_agg_metrics_file(
    rollouts_jsonl_fpath: str,
    role: nemo_gym.comparison.loading.RunRole,
    index: int = 0,
    aggregate_metrics_fpath_override: typing.Optional[str] = None
) -> nemo_gym.comparison.schema.RunFile

Resolve and parse one run’s *_aggregate_metrics.json.

nemo_gym.comparison.loading.resolve_agent_selections(
    baseline_file: nemo_gym.comparison.schema.RunFile,
    candidate_files: typing.Sequence[nemo_gym.comparison.schema.RunFile],
    agent_name: typing.Optional[str] = None,
    baseline_agent_name: typing.Optional[str] = None,
    candidate_agent_names: typing.Optional[typing.Sequence[str]] = None
) -> typing.Tuple[typing.List[nemo_gym.comparison.loading.AgentSelection], typing.List[str], typing.Dict[str, typing.List[str]]]

Decide which agent to read from each side.

Precedence is most-specific-first: a per-side name beats agent_name, which beats the default full join over agent names common to every run. Returns the selections, any warnings, and the agents that were present but not compared (keyed by run label).

nemo_gym.comparison.loading.resolve_aggregate_metrics_fpath(
    rollouts_jsonl_fpath: str,
    override: typing.Optional[str]
) -> pathlib.Path

The *_aggregate_metrics.json to read: the explicit override, else the derived sibling.

nemo_gym.comparison.loading.RunRole = Literal['baseline', 'candidate']