Preserve NaN log probabilities in trajectory JSON

src/art/trajectories/__init__.py

@@ -76,6 +76,7 @@
 from ..types import Messages, MessagesAndChoices, Tools
 from ._serialization import (
     _CompactModel,
+    _equal_with_nan,
     _rebind_history_sources,
     _StringInterningModel,
     _StringPool,
@@ -1330,9 +1331,8 @@ def _bind_source_graph(self) -> TokenizedTrajectoryGroup[TokenizedTrajectoryT]:
         for tokenized, trajectory in zip(
             self.trajectories, self.trajectory_group.trajectories, strict=True
         ):
-            if (
-                tokenized.trajectory is not trajectory
-                and tokenized.trajectory.model_dump() != trajectory.model_dump()
+            if tokenized.trajectory is not trajectory and not _equal_with_nan(
+                tokenized.trajectory.model_dump(), trajectory.model_dump()
             ):
                 raise ValueError("Tokenized trajectory does not match its source group")
             tokenized.trajectory = trajectory

src/art/trajectories/_compact.py

@@ -20,7 +20,7 @@
     TrajectoryGroup,
     _load_tensors,
 )
-from ._serialization import _rebind_history_sources
+from ._serialization import _equal_with_nan, _rebind_history_sources
 
 _FORMAT = "art.trajectories"
 _VERSION = 1
@@ -405,7 +405,7 @@ def _compact_group_data(data: dict[str, pydantic.JsonValue]) -> None:
             dict[str, pydantic.JsonValue],
             dict(_mapping(item, "Tokenized trajectory")),
         )
-        if child.pop("trajectory", None) != source:
+        if not _equal_with_nan(child.pop("trajectory", None), source):
             raise ValueError("Tokenized trajectory does not match its source group")
         registry = _ExchangeDataRegistry.from_trajectory_data(source)
         if "history" in child:

src/art/trajectories/_serialization.py

@@ -3,6 +3,7 @@
 from collections.abc import Callable
 from contextlib import contextmanager
 from dataclasses import fields, is_dataclass
+import math
 import threading
 from typing import Any, Literal, SupportsIndex, cast
 
@@ -31,6 +32,8 @@ def _without_pickle_string_interning():
 class _StringInterningModel(BaseModel):
     """Intern strings once, immediately before this graph is pickled."""
 
+    model_config = pydantic.ConfigDict(ser_json_inf_nan="strings")
+
     # Process-local optimization state: omitting it from Pydantic private state keeps
     # equality and serialization unchanged, and lets a receiving process prepare the
     # graph again after local mutation.
@@ -232,6 +235,26 @@ def validate_history(value: object) -> object:
     return model.model_validate(data)
 
 
+def _equal_with_nan(left: Any, right: Any) -> bool:
+    """Compare serialized source data, treating unknown log probabilities as equal."""
+    if left is right:
+        return True
+    if isinstance(left, dict) and isinstance(right, dict):
+        return left.keys() == right.keys() and all(
+            _equal_with_nan(value, right[key]) for key, value in left.items()
+        )
+    if isinstance(left, (list, tuple)) and type(left) is type(right):
+        return len(left) == len(right) and all(
+            _equal_with_nan(a, b) for a, b in zip(left, right, strict=True)
+        )
+    return left == right or (
+        isinstance(left, float)
+        and isinstance(right, float)
+        and math.isnan(left)
+        and math.isnan(right)
+    )
+
+
 def _rebind_history_sources(
     history: object,
     trajectory: object | None = None,
@@ -246,6 +269,7 @@ def _rebind_history_sources(
         MessagesExchange,
         ResponsesExchange,
         Trajectory,
+        _Exchange,
     )
 
     exchange_types = (
@@ -255,7 +279,7 @@ def _rebind_history_sources(
         MessagesExchange,
     )
 
-    def exchanges(value: object) -> list[object]:
+    def exchanges(value: object) -> list[_Exchange]:
         if not isinstance(value, Trajectory):
             return []
         return [
@@ -296,10 +320,17 @@ def visit(value: object) -> None:
                         if replacement is not item:
                             object.__setattr__(value, name, replacement)
                         continue
+                    # Reject unrelated calls before dumping their full payloads.
                     matches = [
                         exchange
                         for exchange in canonical
-                        if type(exchange) is type(item) and exchange == item
+                        if type(exchange) is type(item)
+                        and exchange.start_time == item.start_time
+                        and exchange.end_time == item.end_time
+                        and (
+                            exchange == item
+                            or _equal_with_nan(exchange.model_dump(), item.model_dump())
+                        )
                     ]
                     if matches:
                         object.__setattr__(value, name, matches[0])

src/art/trajectories/tensors.py

@@ -28,6 +28,7 @@
     _StringInterningModel,
 )
 from ._serialization import (
+    _equal_with_nan,
     _rebind_history_sources,
     serialize_history,
     validate_history,
@@ -286,9 +287,8 @@ def bind_source_group(self) -> Self:
         for tensorized, trajectory in zip(
             self.trajectories, self.trajectory_group.trajectories, strict=True
         ):
-            if (
-                tensorized.trajectory is not trajectory
-                and tensorized.trajectory.model_dump() != trajectory.model_dump()
+            if tensorized.trajectory is not trajectory and not _equal_with_nan(
+                tensorized.trajectory.model_dump(), trajectory.model_dump()
             ):
                 raise ValueError(
                     "Tensorized trajectory does not match its source group"

tests/unit/trajectories/test_tokenize.py

@@ -4380,6 +4380,30 @@ def test_json_round_trip_preserves_exchange_types() -> None:
     assert isinstance(restored.exchanges.chat_completions[0].response, ChatCompletion)
 
 
+@pytest.mark.parametrize("grouped", [False, True])
+def test_json_round_trip_preserves_nan_response_logprobs(grouped: bool) -> None:
+    exchange = _chat_exchange([1], [2, 3])
+    logprobs = exchange.response.choices[0].logprobs
+    assert logprobs is not None and logprobs.content is not None
+    logprobs.content[1].logprob = math.nan
+    trajectory = art.Trajectory(
+        exchanges=TrajectoryExchanges(chat_completions=[exchange])
+    )
+    expected = trajectory.tokenize()
+    value = art.TrajectoryGroup([trajectory]) if grouped else trajectory
+    for _ in range(2):
+        value = type(value).model_validate_json(value.model_dump_json())
+        restored = (
+            value.trajectories[0] if isinstance(value, art.TrajectoryGroup) else value
+        )
+        actual = restored.tokenize()
+        assert actual.tokens == expected.tokens
+        assert actual.flags == expected.flags
+        assert actual.logprobs[1] == expected.logprobs[1]
+        assert math.isnan(actual.logprobs[2])
+        assert actual.flags[2] & tr.TokenFlag.SAMPLED
+
+
 def _response_exchange(
     response_id: str,
     output_id: int,

tests/unit/trajectories/test_tokenized_models.py

@@ -238,7 +238,46 @@ def unexpected_dump(*_: object, **__: object) -> object:
     assert group.trajectories[0].trajectory is source
 
 
-def test_public_group_tokenization_nan_json_round_trip() -> None:
+@pytest.mark.parametrize("logprob", [-0.25, math.nan])
+def test_source_rebinding_does_not_dump_unrelated_exchanges(
+    logprob: float, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    from test_tokenize import _chat_exchange
+
+    exchanges = [
+        _chat_exchange(list(range(1, 2 * i + 2)), [2 * i + 2], offset=i)
+        for i in range(12)
+    ]
+    for exchange in exchanges:
+        old = exchange.response.choices[0].logprobs
+        assert old and old.content
+        old.content[0].logprob = logprob
+    tokenized = art.Trajectory(
+        exchanges=tr.TrajectoryExchanges(chat_completions=exchanges)
+    ).tokenize()
+    payload = tokenized.model_dump_json()
+    dumps = 0
+    original_dump = tr.ChatCompletionsExchange.model_dump
+
+    def counted_dump(self: tr.ChatCompletionsExchange, *args: Any, **kwargs: Any):
+        nonlocal dumps
+        dumps += 1
+        return original_dump(self, *args, **kwargs)
+
+    monkeypatch.setattr(tr.ChatCompletionsExchange, "model_dump", counted_dump)
+    restored = tr.TokenizedTrajectory.model_validate_json(payload)
+    assert isinstance(restored.history, tr.ChatCompletionsHistory)
+    sources = [s for s in restored.history.message_sources if s is not None]
+    assert dumps <= (2 * len(sources) if math.isnan(logprob) else 0)
+    assert all(
+        any(s.exchange is e for e in restored.trajectory.exchanges.chat_completions)
+        for s in sources
+    )
+    assert restored.model_dump_json() == payload
+
+
+@pytest.mark.parametrize("logprob", [-0.25, math.nan])
+def test_public_group_tokenization_nan_json_round_trip(logprob: float) -> None:
     from datetime import datetime
 
     from openai.types.chat import ChatCompletion
@@ -271,7 +310,7 @@ def test_public_group_tokenization_nan_json_round_trip() -> None:
                             "content": [
                                 {
                                     "token": f"token_id:{token_id}",
-                                    "logprob": -0.1 * token_id,
+                                    "logprob": logprob,
                                     "bytes": [],
                                     "top_logprobs": [],
                                 }
@@ -323,6 +362,37 @@ def test_public_group_tokenization_nan_json_round_trip() -> None:
     tr.TokenizedTrajectoryGroup[tr.TokenizedMultiHistoryTrajectory].model_validate_json(
         multi_json
     )
+    for group in (single, multi, single.tensorize(), multi.tensorize()):
+        for restored in (
+            type(group).model_validate_json(group.model_dump_json()),
+            tr.compact_validate(group.compact_dump(), type=type(group)),
+        ):
+            assert restored.model_dump_json() == group.model_dump_json()
+            child = restored.trajectories[0]
+            assert child.trajectory is restored.trajectory_group.trajectories[0]
+            histories = (
+                child.histories
+                if isinstance(
+                    child,
+                    (
+                        tr.TokenizedMultiHistoryTrajectory,
+                        tr.TensorizedMultiHistoryTrajectory,
+                    ),
+                )
+                else [child]
+            )
+            for history in histories:
+                assert isinstance(history.history, tr.ChatCompletionsHistory)
+                for source in history.history.message_sources:
+                    assert source is not None
+                    assert (
+                        source.exchange
+                        is child.trajectory.exchanges.chat_completions[0]
+                    )
+        payload = group.model_dump(mode="json")
+        payload["trajectory_group"]["trajectories"][0]["reward"] = 123
+        with pytest.raises(ValueError, match="does not match its source group"):
+            type(group).model_validate(payload)
 
 
 def test_tokenized_compact_round_trips_retain_source_references() -> None: