Suppress env frame errors; serialize logprob

ldp/alg/callbacks.py

@@ -4,6 +4,7 @@
 import time
 from collections import defaultdict
 from collections.abc import Collection, Iterable, Sequence
+from contextlib import suppress
 from pathlib import Path
 from typing import Any, cast
 
@@ -173,7 +174,11 @@ async def after_transition(
         # TODO: make this async?
         traj.to_jsonl(self.out_files[traj_id])
         if transition.done:
-            with Path(self.env_files[traj_id]).open("w") as f:
+            with (
+                # Do not fail if the environment didn't implement export_frame().
+                suppress(NotImplementedError),
+                Path(self.env_files[traj_id]).open("w") as f,
+            ):
                 f.write(env.export_frame().model_dump_json(exclude={"state"}, indent=2))
 
 

ldp/data_structures.py

@@ -4,6 +4,7 @@
 import logging
 import os
 from collections.abc import Callable, Hashable, Iterable
+from contextlib import suppress
 from typing import Any, ClassVar, Self, cast
 from uuid import UUID
 
@@ -121,7 +122,12 @@ def from_jsonl(cls, filename: str | os.PathLike) -> Self:
             reader = iter(f)
             traj = cls(traj_id=json.loads(next(reader)))
             for json_line in reader:
-                traj.steps.append(Transition(**json.loads(json_line)))
+                data = json.loads(json_line)
+                # logprob may have been serialized, but cannot be passed to
+                # OpResult, so remove it here.
+                with suppress(KeyError):
+                    data["action"].pop("logprob")
+                traj.steps.append(Transition(**data))
         return traj
 
     def compute_discounted_returns(self, discount: float = 1.0) -> list[float]:

ldp/graph/async_torch.py

@@ -115,8 +115,8 @@ async def __call__(self, **kwargs):
             # Sleep, to let another coroutine take over if it needs to
             await asyncio.sleep(0.0)
 
-    async def _maybe_process_batch(self):
-        """If the buffer is >= batch size or we have been waiting long enough, process the old batch.
+    async def _maybe_process_batch(self) -> None:
+        """If the buffer is >= batch size or we have been waiting long enough, process the oldest batch.
 
         If neither condition is met, do nothing.
         """
@@ -134,7 +134,7 @@ async def _maybe_process_batch(self):
             batch = self._work_buffer[: self.batch_size]
             self._work_buffer = self._work_buffer[self.batch_size :]
 
-            # Construct the batch tensors
+            # Construct the batch inputs
             sample_kwargs = [x[2] for x in batch]
             batch_kwargs = self.collate_fn(sample_kwargs)
 
@@ -144,7 +144,7 @@ async def _maybe_process_batch(self):
             self._result_buffer.update(zip(request_ids, results, strict=True))
 
     @abstractmethod
-    async def _batched_call(self, batch_kwargs: dict[str, Any]):
+    async def _batched_call(self, batch_kwargs: dict[str, Any]) -> Any:
         """Logic to call the worker on a batch of inputs."""
 
 

ldp/graph/ops.py

@@ -61,6 +61,7 @@ def to_dict(self) -> dict[str, Any]:
             "op_name": self.op_name,
             "op_class_name": self.op_class_name,
             "value": value_dump,
+            "logprob": self.logprob,
         }
 
     @classmethod