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

@@ -105,9 +105,9 @@ def record_to_train_batch(
 
     Returns:
         BatchedDataDict with input_ids, input_lengths, generation_logprobs,
-        token_mask, an all-ones sample_mask, the raw mask_sample and truncated
-        flags, prompt_ids_for_adv, total_reward, violation counts, and optional
-        routed experts and message-violation masks.
+        token_mask, a prompt-loss-weighted sample_mask, the raw mask_sample and
+        truncated flags, prompt_ids_for_adv, total_reward, violation counts,
+        and optional routed experts and message-violation masks.
     """
     # Lazy imports: grpo and llm_message_utils transitively pull
     # experience.rollouts, so importing at module top risks a cycle.
@@ -154,7 +154,7 @@ def record_to_train_batch(
     )
     mask_sample = _mask_sample_flags(c.env_extras for c in completions)
     truncated = torch.tensor([c.truncated for c in completions], dtype=torch.bool)
-    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

@@ -435,6 +435,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,
         )
@@ -828,6 +829,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

@@ -82,7 +82,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=[
@@ -96,6 +98,7 @@ def _record(completions: list[Completion]) -> PromptGroupRecord:
         metadata={"task_name": "test"},
         completions=completions,
         rollout_metrics={},
+        loss_multiplier=loss_multiplier,
     )
 
 
@@ -378,3 +381,33 @@ 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},
+        include_message_violation_fields=False,
+    )
+
+    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},
+        include_message_violation_fields=False,
+    )
+
+    torch.testing.assert_close(batch["sample_mask"], torch.zeros(1))

tests/unit/experience/test_rollout_manager.py

@@ -807,7 +807,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
@@ -827,6 +828,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)}"
     )
@@ -905,7 +907,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
@@ -951,6 +954,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]
@@ -1066,7 +1070,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
 
@@ -1081,6 +1085,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)}"
     )