nemo_rl.experience.payload

Producer-side payload helpers for the async-RL TQ path.

Module Contents

Functions

record_to_train_batch	Convert one prompt group’s record into a packed BatchedDataDict of N rows.
pack_payload	Pack a producer batch into (sample_ids, fields, tags) for put_samples.

API

nemo_rl.experience.payload.record_to_train_batch(

record: nemo_rl.experience.interfaces.PromptGroupRecord,

*,

pad_value_dict: collections.abc.Mapping[str, int],

) → nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any]

Convert one prompt group’s record into a packed BatchedDataDict of N rows.

Parameters:

• record – Rollout’s PromptGroupRecord with N completions to flatten into rows.

• pad_value_dict – Field-name → pad value used by batched_message_log_to_flat_message.

Returns:

BatchedDataDict with input_ids, input_lengths, generation_logprobs, token_mask, sample_mask, prompt_ids_for_adv, total_reward, and optional routed_experts.

nemo_rl.experience.payload.pack_payload(

train_batch: collections.abc.Mapping[str, Any],

*,

weight_version: int,

group_id: str,

) → tuple[list[str], tensordict.TensorDict, list[dict[str, Any]]]

Pack a producer batch into (sample_ids, fields, tags) for put_samples.

Parameters:

• train_batch – Mapping with at least input_lengths plus the tensor/object fields to send.

• weight_version – Trainer weight version stamped on every row’s tag.

• group_id – Per-group identifier used as the sample_id prefix; the caller owns uniqueness.

Returns:

sample_ids of the form {group_id}_g{i}, a jagged-packed TensorDict, and per-row tags.