fix(sc): preserve DatumSpec loss_multiplier through rollouts

nemo_rl/experience/interfaces.py

@@ -58,3 +58,5 @@ class PromptGroupRecord:
     metadata: dict[str, Any]
     completions: list["Completion"]
     rollout_metrics: dict[str, Any]
+    # Prompt-level DatumSpec weight; defaults preserve manually-built records.
+    loss_multiplier: float = 1.0

nemo_rl/experience/payload.py

@@ -107,7 +107,9 @@ def record_to_train_batch(
     total_reward = torch.tensor(
         [float(c.reward) for c in completions], dtype=torch.float32
     )
-    sample_mask = torch.ones(n, dtype=torch.float32)
+    sample_mask = torch.full(
+        (n,), float(record.loss_multiplier), dtype=torch.float32
+    )
 
     train_data: dict[str, Any] = {
         "input_ids": flat["token_ids"],

nemo_rl/experience/rollout_manager.py

@@ -433,6 +433,7 @@ async def run_rollout(self, input_sample: DatumSpec) -> PromptGroupRecord:
             prompt=input_sample["message_log"],
             extra_env_info=input_sample["extra_env_info"],
             metadata={"task_name": input_sample["task_name"]},
+            loss_multiplier=float(input_sample["loss_multiplier"]),
             completions=completions,
             rollout_metrics=rollout_metrics,
         )
@@ -824,6 +825,7 @@ async def run_rollout(self, input_sample: DatumSpec) -> PromptGroupRecord:
             prompt=prompt_message_log,
             extra_env_info=input_sample["extra_env_info"],
             metadata={"task_name": "nemo_gym"},
+            loss_multiplier=float(input_sample["loss_multiplier"]),
             completions=completions,
             rollout_metrics=rollout_metrics,
         )

tests/unit/experience/test_payload.py

@@ -71,7 +71,9 @@ def _completion(
     )
 
 
-def _record(completions: list[Completion]) -> PromptGroupRecord:
+def _record(
+    completions: list[Completion], *, loss_multiplier: float = 1.0
+) -> PromptGroupRecord:
     return PromptGroupRecord(
         prompt_idx=0,
         prompt=[
@@ -85,6 +87,7 @@ def _record(completions: list[Completion]) -> PromptGroupRecord:
         metadata={"task_name": "test"},
         completions=completions,
         rollout_metrics={},
+        loss_multiplier=loss_multiplier,
     )
 
 
@@ -252,3 +255,27 @@ def test_pack_payload_stamps_violation_counts_on_tags() -> None:
             "num_assistant_messages": 0,
         },
     ]
+
+
+def test_record_to_train_batch_preserves_fractional_loss_multiplier() -> None:
+    record = _record(
+        [
+            _completion(route_start=10, reward=1.0),
+            _completion(route_start=20, reward=0.0),
+        ],
+        loss_multiplier=0.25,
+    )
+
+    batch = record_to_train_batch(record, pad_value_dict={"token_ids": 0})
+
+    torch.testing.assert_close(batch["sample_mask"], torch.full((2,), 0.25))
+
+
+def test_record_to_train_batch_drops_zero_weight_prompt() -> None:
+    record = _record(
+        [_completion(route_start=10, reward=1.0)], loss_multiplier=0.0
+    )
+
+    batch = record_to_train_batch(record, pad_value_dict={"token_ids": 0})
+
+    torch.testing.assert_close(batch["sample_mask"], torch.zeros(1))

tests/unit/experience/test_rollout_manager.py

@@ -691,7 +691,8 @@ def test_async_rollout_manager(
     - completions hold independent (not aliased) message_log objects
     """
     vllm_generation, tokenizer, task_to_env, _, _ = multi_step_setup_vllm_async
-    input_sample = single_multi_step_calculator_input_sample
+    input_sample = deepcopy(single_multi_step_calculator_input_sample)
+    input_sample["loss_multiplier"] = 0.25
     num_generations = 2
     max_seq_len = 1024
     max_rollout_turns = input_sample["extra_env_info"]["max_steps"] + 1
@@ -711,6 +712,7 @@ def test_async_rollout_manager(
     vllm_generation.finish_generation()
 
     assert isinstance(record, PromptGroupRecord)
+    assert record.loss_multiplier == 0.25
     assert len(record.completions) == num_generations, (
         f"Expected {num_generations} completions, got {len(record.completions)}"
     )
@@ -789,7 +791,8 @@ def test_async_rollout_manager_matches_original(
     TODO: remove this test together with run_async_multi_turn_rollout when the legacy path is deleted.
     """
     vllm_generation, tokenizer, task_to_env, _, _ = multi_step_setup_vllm_async
-    input_sample = single_multi_step_calculator_input_sample
+    input_sample = deepcopy(single_multi_step_calculator_input_sample)
+    input_sample["loss_multiplier"] = 0.25
     num_generations = 2
     max_seq_len = 1024
     max_rollout_turns = input_sample["extra_env_info"]["max_steps"] + 1
@@ -835,6 +838,7 @@ def test_async_rollout_manager_matches_original(
     # Both should produce N results
     assert len(original_batch["message_log"]) == num_generations
     assert len(record.completions) == num_generations
+    assert record.loss_multiplier == 0.25
 
     for i in range(num_generations):
         orig_msg_log = original_batch["message_log"][i]
@@ -950,7 +954,7 @@ def test_async_nemo_gym_rollout_manager(
         "extra_env_info": input_batch["extra_env_info"][0],
         "task_name": "nemo_gym",
         "idx": 0,
-        "loss_multiplier": float(input_batch["loss_multiplier"][0]),
+        "loss_multiplier": 0.25,
     }
     num_generations = 2
 
@@ -965,6 +969,7 @@ def test_async_nemo_gym_rollout_manager(
     record = asyncio.run(manager.run_rollout(single_prompt))
 
     assert isinstance(record, PromptGroupRecord)
+    assert record.loss_multiplier == 0.25
     assert len(record.completions) == num_generations, (
         f"Expected {num_generations} completions, got {len(record.completions)}"
     )