LLM Probe Environments

See also

Debugging Reinforcement Learning for a guide on how to use probe environments to debug RL and LLM agents.

Lightweight string-observation probe environments for multi-turn LLM debugging.

These are ordered by increasing complexity for debugging LLM PPO training pipelines: fixed target, input-dependent target, multi-digit arithmetic target, then multi-turn spatial navigation.

class agilerl.utils.probe_envs_llm.ConditionalTargetEnv(digits: tuple[int, ...] = (1, 2, 3), seed: int = 42)

Single-turn probe: observation is one digit; target is (digit % 3) + 1.

reset(seed: int | None = None) → tuple[str, dict]

step(action: str) → tuple[str, float, bool, bool, dict]

class agilerl.utils.probe_envs_llm.ConstantTargetEnv(target_digit: str = '3', prompt: str = '11', seed: int = 42)

Single-turn probe with a fixed prompt and a fixed correct digit.

reset(seed: int | None = None) → tuple[str, dict]

step(action: str) → tuple[str, float, bool, bool, dict]

class agilerl.utils.probe_envs_llm.GridNavigationEnv(grid_size: int = 4, max_turns: int = 5, step_cost: float = -0.1, seed: int = 42)

1D grid navigation: move left/right to reach a target position.

The agent sees its starting position and target once (initial obs), then receives only its current position as feedback after each move. It must remember the target from context and navigate toward it.

Actions map to: “1” = left, “2” = stay, “3” = right. Positions are clamped to [0, grid_size-1].

reset(seed=None)

step(action)

class agilerl.utils.probe_envs_llm.MultiInputConditionalEnv(digits: tuple[int, ...] = (1, 2, 3), seed: int = 42)

Single-turn probe: two-digit prompt; target is (sum(digits) % 3) + 1.

reset(seed: int | None = None) → tuple[str, dict]

step(action: str) → tuple[str, float, bool, bool, dict]