nemo_gym.comparison.loading

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

Input file 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 from which its _aggregate_metrics.json sibling is derived. Aggregate files written before repeat-level statistics existed are enriched from their already-recorded per-rollout summaries in memory.

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
`_compute_repeat_metrics`	Add repeat metrics to one legacy aggregate entry.
`_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]],
    repeat_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.

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

repeat_level_metrics

List[Dict[str, Any]]

nemo_gym.comparison.loading._compute_repeat_metrics(
    entry: typing.Dict[str, typing.Any]
) -> typing.Dict[str, typing.Any]

Add repeat metrics to one legacy aggregate entry.

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']