Copybara import of the project:

--
6ab68a057443f038a3d844807f205b1c99cab944 by Kevin Zakka <kevinarmandzakka@gmail.com>:

Fix pickle code-execution vulnerability in sysid loaders

Serialize signal_mapping as JSON so the trajectory and time series loaders can use allow_pickle=False, preventing arbitrary code execution from untrusted .npz files.

COPYBARA_INTEGRATE_REVIEW=https://github.com/google-deepmind/mujoco/pull/3353 from kevinzakka:sysid-disable-pickle-load 6ab68a057443f038a3d844807f205b1c99cab944
PiperOrigin-RevId: 935400242
Change-Id: Iecd907174441fbcfd02c106b912a7ff375c9098c

python/mujoco/sysid/_src/timeseries.py

@@ -20,6 +20,7 @@ from __future__ import annotations
 from collections.abc import Sequence
 from dataclasses import dataclass
 from enum import Enum
+import json
 import pathlib
 from typing import Literal, TypeAlias
 
@@ -39,6 +40,27 @@ class SignalType(Enum):
 
 SignalMappingType: TypeAlias = dict[str, tuple[SignalType, np.ndarray]]
 
+
+def encode_signal_mapping(signal_mapping: SignalMappingType) -> str:
+  """Serialize a signal mapping to a JSON string.
+
+  This avoids storing object arrays (which require ``allow_pickle=True`` to
+  load) on disk. The mapping is plain data: enum values and integer indices.
+  """
+  return json.dumps({
+      name: [sig_type.value, np.asarray(indices).astype(int).tolist()]
+      for name, (sig_type, indices) in signal_mapping.items()
+  })
+
+
+def decode_signal_mapping(encoded: str) -> SignalMappingType:
+  """Inverse of :func:`encode_signal_mapping`."""
+  raw = json.loads(encoded)
+  return {
+      name: (SignalType(sig_type), np.asarray(indices, dtype=int))
+      for name, (sig_type, indices) in raw.items()
+  }
+
 InterpolationMethod = Literal[
     "linear", "cubic", "quadratic", "quintic", "zero_order_hold", "zoh"
 ]
@@ -569,12 +591,12 @@ class TimeSeries:
     Args:
       path: Path where the data will be saved.
     """
-    np.savez(
-        path,
-        times=self.times,
-        data=self.data,
-        signal_mapping=np.array(self.signal_mapping, dtype=object),
-    )
+    save_dict = {"times": self.times, "data": self.data}
+    if self.signal_mapping:
+      save_dict["signal_mapping"] = np.asarray(
+          encode_signal_mapping(self.signal_mapping)
+      )
+    np.savez(path, **save_dict)
 
   def save_to_csv(self, path: str | pathlib.Path) -> None:
     """Save the time series data to a CSV file.
@@ -598,11 +620,11 @@ class TimeSeries:
     Returns:
       A new TimeSeries object.
     """
-    with np.load(path, allow_pickle=True) as npz:
+    with np.load(path, allow_pickle=False) as npz:
       times = npz["times"]
       data = npz["data"]
       if "signal_mapping" in npz:
-        signal_mapping = npz["signal_mapping"].item()
+        signal_mapping = decode_signal_mapping(str(npz["signal_mapping"]))
       else:
         signal_mapping = None