[Tunix] improve trajectory logging to suppor numpy array and scalar types.

PiperOrigin-RevId: 875760462

Author: lc5211

tests/utils/trajectory_logger_test.py

@@ -3,6 +3,7 @@
 from typing import TypedDict
 
 from absl.testing import absltest
+import numpy as np
 import pandas as pd
 from tunix.utils import trajectory_logger
 
@@ -52,6 +53,7 @@ def test_log_item_creates_and_writes_to_file(self):
         trajectory_id='t0',
         completion='c0',
         prompt='p0',
+        value=np.int32(0),
     )
     trajectory_logger.log_item(temp_dir, item1)
 
@@ -63,6 +65,7 @@ def test_log_item_creates_and_writes_to_file(self):
         trajectory_id='t1',
         completion='c1|pipe',
         prompt='p1</reasoning>',
+        value=np.int32(1),
     )
     trajectory_logger.log_item(temp_dir, item2)
 
@@ -71,6 +74,7 @@ def test_log_item_creates_and_writes_to_file(self):
         trajectory_id='t2',
         completion='a, "b", c',
         prompt='a prompt with\na newline',
+        value=np.int32(2),
     )
     trajectory_logger.log_item(temp_dir, item3)
 
@@ -80,6 +84,7 @@ def test_log_item_creates_and_writes_to_file(self):
       self.assertEqual(df['trajectory_id'].tolist(), ['t0', 't1', 't2'])
       self.assertEqual(df['completion'][2], 'a, "b", c')
       self.assertEqual(df['prompt'][2], 'a prompt with\na newline')
+      self.assertEqual(df['value'].tolist(), [0, 1, 2])
 
 
 if __name__ == '__main__':

tunix/utils/trajectory_logger.py

@@ -25,6 +25,7 @@
 from etils import epath
 from google.protobuf import json_format
 from google.protobuf import message
+import numpy as np
 import pandas as pd
 
 def _make_serializable(item: Any) -> Any:
@@ -39,6 +40,16 @@ def _make_serializable(item: Any) -> Any:
     return _make_serializable(dataclasses.asdict(item))
   elif isinstance(item, message.Message):
     return json_format.MessageToDict(item)
+  elif isinstance(item, np.ndarray):
+    return _make_serializable(item.tolist())
+  elif isinstance(item, np.integer):
+    return int(item)
+  elif isinstance(item, np.floating):
+    return float(item)
+  elif isinstance(item, np.bool_):
+    return bool(item)
+  elif isinstance(item, np.str_):
+    return str(item)
   elif isinstance(item, (float, int, bool, str)):
     return item
   else: