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
 

python/mujoco/sysid/_src/trajectory.py

@@ -102,18 +102,19 @@ class SystemTrajectory:
     if self.state is not None:
       save_dict["state_times"] = self.state.times
       save_dict["state_data"] = self.state.data
-      save_dict["state_signal_mapping"] = np.array(
-          self.state.signal_mapping, dtype=object
-      )
+      if self.state.signal_mapping:
+        save_dict["state_signal_mapping"] = np.asarray(
+            timeseries.encode_signal_mapping(self.state.signal_mapping)
+        )
 
     if self.control.signal_mapping:
-      save_dict["control_signal_mapping"] = np.array(
-          self.control.signal_mapping, dtype=object
+      save_dict["control_signal_mapping"] = np.asarray(
+          timeseries.encode_signal_mapping(self.control.signal_mapping)
       )
 
     if self.sensordata.signal_mapping:
-      save_dict["sensordata_signal_mapping"] = np.array(
-          self.sensordata.signal_mapping, dtype=object
+      save_dict["sensordata_signal_mapping"] = np.asarray(
+          timeseries.encode_signal_mapping(self.sensordata.signal_mapping)
       )
 
     np.savez(path, **save_dict)  # type: ignore
@@ -126,7 +127,7 @@ class SystemTrajectory:
       allow_missing_sensors: bool = False,
   ) -> SystemTrajectory:
     """Load a trajectory from a compressed NumPy archive."""
-    with np.load(path, allow_pickle=True) as npz:
+    with np.load(path, allow_pickle=False) as npz:
       control_times = npz["control_times"]
       control_data = npz["control_data"]
       sensordata_times = npz["sensordata_times"]
@@ -137,15 +138,21 @@ class SystemTrajectory:
 
       control_signal_mapping = None
       if "control_signal_mapping" in npz:
-        control_signal_mapping = npz["control_signal_mapping"].item()
+        control_signal_mapping = timeseries.decode_signal_mapping(
+            str(npz["control_signal_mapping"])
+        )
 
       sensordata_signal_mapping = None
       if "sensordata_signal_mapping" in npz:
-        sensordata_signal_mapping = npz["sensordata_signal_mapping"].item()
+        sensordata_signal_mapping = timeseries.decode_signal_mapping(
+            str(npz["sensordata_signal_mapping"])
+        )
 
       state_signal_mapping = None
       if "state_signal_mapping" in npz:
-        state_signal_mapping = npz["state_signal_mapping"].item()
+        state_signal_mapping = timeseries.decode_signal_mapping(
+            str(npz["state_signal_mapping"])
+        )
 
     predicted_rollout = cls(
         model=model,