style: compact comments in language runtime

tests/policies/pi052/test_pi052_attention_masking.py

@@ -56,17 +56,9 @@ def _shifted_ce(logits, labels):
     eye = torch.eye(vocab_size, dtype=logits.dtype, device="cuda")
     return _shifted_lin_ce(logits.cuda(), eye, labels.cuda()).cpu()
 
-# ---------------------------------------------------------------------------
-# A synthetic PI052 prefix layout: [images, prompt-lang, target-lang]
-#
-#   indices 0-1  : 2 image tokens          (att = 0)
-#   indices 2-4  : 3 user-prompt lang      (att = 0)
-#   indices 5-8  : 4 supervised target lang(att = 0 from embed_prefix)
-#
-# ``text_labels`` covers the 7 language tokens; -100 on the prompt span,
-# real ids on the 4-token target span. PaliGemma's prefix has no state
-# token (unlike SmolVLA), so the lang span ends at the prefix end.
-# ---------------------------------------------------------------------------
+
+# Synthetic prefix: two image tokens, three prompt tokens, and four supervised target tokens.
+# Text labels mask the prompt with -100 and cover the target through the prefix end.
 N_IMAGE = 2
 N_PROMPT = 3
 N_TARGET = 4
@@ -95,9 +87,7 @@ def _attends(prefix_att_masks: torch.Tensor) -> torch.Tensor:
 
 def test_mark_sets_att_on_targets_only():
     """Only the supervised target language positions flip to att=1."""
-    marked = _mark_target_span_causal(
-        _embed_prefix_att_masks(), _text_labels(), LANG_START, LANG_END
-    )
+    marked = _mark_target_span_causal(_embed_prefix_att_masks(), _text_labels(), LANG_START, LANG_END)
     expected = [False] * PREFIX_LEN
     for i in range(LANG_START + N_PROMPT, LANG_END):  # target span
         expected[i] = True
@@ -107,9 +97,7 @@ def test_mark_sets_att_on_targets_only():
 def test_target_tokens_attend_causally_among_themselves():
     """A target token must NOT attend to later targets, but must attend
     to earlier ones — genuine causal next-token prediction."""
-    marked = _mark_target_span_causal(
-        _embed_prefix_att_masks(), _text_labels(), LANG_START, LANG_END
-    )
+    marked = _mark_target_span_causal(_embed_prefix_att_masks(), _text_labels(), LANG_START, LANG_END)
     attends = _attends(marked)
     tgt = range(LANG_START + N_PROMPT, LANG_END)
     for i in tgt:
@@ -122,9 +110,7 @@ def test_target_tokens_attend_causally_among_themselves():
 
 def test_target_tokens_attend_prompt_and_images_bidirectionally():
     """Targets keep full visibility of images + the user prompt."""
-    marked = _mark_target_span_causal(
-        _embed_prefix_att_masks(), _text_labels(), LANG_START, LANG_END
-    )
+    marked = _mark_target_span_causal(_embed_prefix_att_masks(), _text_labels(), LANG_START, LANG_END)
     attends = _attends(marked)
     context = list(range(0, LANG_START + N_PROMPT))  # images + prompt
     for i in range(LANG_START + N_PROMPT, LANG_END):
@@ -136,9 +122,7 @@ def test_non_target_subtask_stays_bidirectional():
     """A flow-only / non-target language span (all -100 labels) leaves the
     mask untouched — the action expert reads it bidirectionally."""
     all_ignored = torch.full((1, N_PROMPT + N_TARGET), -100, dtype=torch.long)
-    marked = _mark_target_span_causal(
-        _embed_prefix_att_masks(), all_ignored, LANG_START, LANG_END
-    )
+    marked = _mark_target_span_causal(_embed_prefix_att_masks(), all_ignored, LANG_START, LANG_END)
     assert torch.equal(marked, _embed_prefix_att_masks())
 
 

tests/policies/pi052/test_pi052_fast_action_loss.py

@@ -18,7 +18,7 @@
 
 import pytest
 import torch
-from torch.nn import functional as F
+from torch.nn import functional as F  # noqa: N812
 
 pytest.importorskip("transformers")
 pytest.importorskip("liger_kernel")
@@ -40,9 +40,7 @@ def _fast_ce(logits, action_tokens, action_code_mask, predict_actions_t):
     vocab_size = logits.shape[-1]
     eye = torch.eye(vocab_size, dtype=logits.dtype, device="cuda")
     predict = predict_actions_t.cuda() if predict_actions_t is not None else None
-    loss = _fast_lin_ce(
-        logits.cuda(), eye, action_tokens.cuda(), action_code_mask.cuda(), predict
-    )
+    loss = _fast_lin_ce(logits.cuda(), eye, action_tokens.cuda(), action_code_mask.cuda(), predict)
     return loss.cpu()
 
 
@@ -67,9 +65,7 @@ def test_fast_ce_supervises_only_discrete_action_codes():
         reduction="mean",
     )
 
-    # Looser tolerance: the fused Triton kernel (GPU) differs from CPU eager
-    # F.cross_entropy at the ~1e-7 level, which exceeds the default rtol on
-    # these very small (~1e-4) losses.
+    # Allow the fused GPU kernel's ~1e-7 difference on small losses.
     assert torch.allclose(loss, expected, atol=1e-5, rtol=1e-3)
 
 
@@ -77,9 +73,7 @@ def test_fast_ce_masks_non_action_samples():
     """Recipe samples with predict_actions=False do not contribute FAST loss."""
     vocab_size = 8
     action_tokens = torch.tensor([[1, 2, 3, 4], [1, 2, 5, 6]])
-    action_code_mask = torch.tensor(
-        [[False, False, True, True], [False, False, True, True]]
-    )
+    action_code_mask = torch.tensor([[False, False, True, True], [False, False, True, True]])
     predict_actions = torch.tensor([True, False])
 
     logits = torch.zeros(2, action_tokens.shape[1], vocab_size)
@@ -96,9 +90,7 @@ def test_fast_ce_masks_non_action_samples():
         reduction="mean",
     )
 
-    # Looser tolerance: the fused Triton kernel (GPU) differs from CPU eager
-    # F.cross_entropy at the ~1e-7 level, which exceeds the default rtol on
-    # these very small (~1e-4) losses.
+    # Allow the fused GPU kernel's ~1e-7 difference on small losses.
     assert torch.allclose(loss, expected, atol=1e-5, rtol=1e-3)
 
 

tests/policies/pi052/test_pi052_text_processor.py

@@ -63,18 +63,11 @@ def test_flatten_leaves_messages_without_tool_calls_untouched():
 
 def test_flatten_drops_non_say_tool_calls_but_keeps_content():
     weather = {"type": "function", "function": {"name": "check_weather", "arguments": {}}}
-    out = _flatten_say_tool_calls(
-        {"role": "assistant", "content": "plan only", "tool_calls": [weather]}
-    )
+    out = _flatten_say_tool_calls({"role": "assistant", "content": "plan only", "tool_calls": [weather]})
     assert out["content"] == "plan only"
     assert "tool_calls" not in out
 
 
-# ---------------------------------------------------------------------------
-# EOS-termination supervision
-# ---------------------------------------------------------------------------
-
-
 def test_format_messages_appends_eos_to_target_turns_only():
     msgs = [
         {"role": "user", "content": "pick cube"},