feat: example multi turn env

[cmunley1]

Example multi turn environment with scripted follow ups, where each question depends on the last turn. Uses Gymnasium interface.

example input

{
  "responses_create_params": {
    "input": [
      {
        "role": "user",
        "content": "What programming language was created by Guido van Rossum?"
      }
    ]
  },
  "follow_ups": [
    "In what year was it first released?"
  ],
  "expected_answer": "1991",
  "agent_ref": {
    "type": "responses_api_agents",
    "name": "example_multi_turn_gymnasium_agent"
  }
}

design

class ExampleMultiTurnEnv(GymnasiumServer):
    session_turns: Dict[str, int] = Field(default_factory=dict)

    async def reset(self, metadata: dict, session_id: Optional[str] = None) -> tuple[Optional[str], dict]:
        """Returns (observation, info)."""
        self.session_turns[session_id] = 0
        return None, {}

    async def step(
        self, action: NeMoGymResponse, metadata: dict, session_id: Optional[str] = None
    ) -> tuple[Optional[str], float, bool, bool, dict]:
        """Returns (observation, reward, terminated, truncated, info)."""
        follow_ups = metadata.get("follow_ups", [])
        turn = self.session_turns.get(session_id, 0)

        if turn < len(follow_ups):
            self.session_turns[session_id] = turn + 1
            return follow_ups[turn], 0.0, False, False, {}

        expected = metadata.get("expected_answer", "")
        text = extract_text(action)
        reward = 1.0 if expected and expected.lower() in text.lower() else 0.0
        return None, reward, True, False, {}

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.