nemo_rl.environments.vlm_environment#
Module Contents#
Classes#
Functions#
API#
- class nemo_rl.environments.vlm_environment.VLMEnvConfig#
Bases:
typing.TypedDict- num_workers: int#
None
- stop_strings: Optional[list[str]]#
None
- reward_functions: List[dict[str, Any]]#
None
- nemo_rl.environments.vlm_environment._mute_output()#
- class nemo_rl.environments.vlm_environment.VLMVerifyWorker( )#
Initialization
- verify(
- pred_responses: list[str],
- ground_truths: list[str],
Verify the correctness of the predicted responses against the ground truth.
- Parameters:
pred_responses – list[str]. The predicted responses from the LLM.
ground_truths – list[str]. The ground truth responses.
- Returns:
list[float]. The rewards for each predicted response.
- class nemo_rl.environments.vlm_environment.VLMEnvironmentMetadata#
Bases:
typing.TypedDict- ground_truth: str#
None
- class nemo_rl.environments.vlm_environment.VLMEnvironment(cfg: nemo_rl.environments.vlm_environment.VLMEnvConfig)#
Bases:
nemo_rl.environments.interfaces.EnvironmentInterface- shutdown() None#
- step(
- message_log_batch: list[list[dict[str, str]]],
- metadata: list[nemo_rl.environments.vlm_environment.VLMEnvironmentMetadata],
Runs a step in the vlm environment.
- Parameters:
message_log – list[list[dict[str, str]]]. A batch of OpenAI-API-like message logs that represent interactions with the VLM.
metadata – list[VLMEnvironmentMetadata]. The grader will use the ‘ground_truth’ key to evaluate correctness.
- Returns:
A tuple containing: - list[dict[str, str]]: Observations/responses batch - list[dict]: Updated metadata - list[str]: Next stop strings for the next turn - Tensor: Rewards tensor - Tensor: Done flags tensor
- Return type:
- global_post_process_and_metrics( ) tuple[nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any], dict[str, float | int]]#
Computes metrics for this environment given a global rollout batch.
Every rank will run this function, so you’re free to use distributed calculations if you’d prefer for heavy metrics.