Hardening: serialize.load() runs jsonpickle.decode() on the infos column (code execution when loading an untrusted trajectory dataset)

[mtholmquist]

Summary

imitation.data.serialize.load() round-trips a saved trajectory dataset. On the HuggingFace datasets (directory) path, the per-trajectory infos column is decoded with jsonpickle.decode() with no restriction, so loading a trajectory-dataset directory produced by an untrusted party and then reading trajectory.infos executes arbitrary Python code embedded in that column.

I recognize serialize.load() is documented as loading trajectories "saved by save()", and the sibling .npz/.pkl branch already calls np.load(path, allow_pickle=True) — i.e. the API already accepts pickle-grade code execution by design. So this is not a claim that a security boundary was bypassed. It's a defense-in-depth / documentation request: the code-execution surface on the datasets path is undocumented, hides in a string column of an artifact (an Arrow dataset directory) that users increasingly treat as inert shareable data, and it fires lazily on read — and the two "obvious" fixes both fail, so I've done the work to identify one that doesn't.

Affected

• Package: imitation (PyPI), all versions since the HuggingFace datasets backend was introduced (observed on 1.0.1, the current release; master is unchanged).
• Sink: src/imitation/data/huggingface_utils.py — _LazyDecodedList.__getitem__ →

  self._decoded_cache[idx] = jsonpickle.decode(self._encoded_list[idx])

  No safe=/classes=/backend restriction is passed.
• Reached from the documented API: serialize.load(dir) → datasets.load_from_disk(dir) → TrajectoryDatasetSequence, which wraps the infos column in _LazyDecodedList. Decoding is deferred to the first read of trajectory.infos[i].

Reachability (fires on normal preprocessing, not just manual indexing)

The decode is lazy, so load() alone is inert — but any normal consumption of infos triggers it. Notably, the standard preprocessing step rollout.flatten_trajectories(...) (the canonical prep before BC/GAIL/AIRL/DAgger/preference training) coerces the lazy list via np.concatenate, which decodes every element. So a victim who merely runs:

from imitation.data import serialize, rollout
trajs = serialize.load(untrusted_dir)          # inert
transitions = rollout.flatten_trajectories(trajs)   # <-- decodes infos → code runs

executes attacker code without ever touching .infos[i] directly.

Proof of concept (benign marker only)

A trajectory dataset whose infos column contains a jsonpickle py/reduce gadget runs code on decode. Illustrated with a benign payload that just prints (no weaponized command):

import datasets, json

# Benign illustration: py/reduce gadget that calls builtins.print on decode.
payload = json.dumps({"py/reduce": [{"py/type": "builtins.print"},
                                    {"py/tuple": ["jsonpickle-decode-executed"]}]})
datasets.Dataset.from_dict({
    "obs":      [[[0.0], [1.0], [0.0]]],
    "acts":     [[[0.0], [0.0]]],
    "infos":    [[payload, payload]],
    "terminal": [True],
}).save_to_disk("demo_dataset/")

Loading demo_dataset/ and reading the trajectory's infos (or running flatten_trajectories) executes the gadget. A real attacker substitutes any callable.

The two obvious fixes do NOT work

I verified both empirically against jsonpickle==4.1.2:

1. jsonpickle.decode(..., safe=True) is a no-op here. safe=True is already the default in jsonpickle 4.x, and it only gates the legacy py/repr (eval) path — py/reduce, py/object, py/type, py/function, py/module all reconstruct regardless. The py/reduce gadget above still executes under safe=True. (See the jsonpickle safe-mode bypass documented in CVE-2026-20251 / Splunk SVD-2026-0601.)
2. json.loads breaks functionality. infos are gym/env step() info dicts that legitimately hold numpy arrays — which is exactly why save() uses jsonpickle.encode. Plain json.loads would silently drop the numpy content (leaving plain dicts), so it isn't a drop-in replacement. Note also that jsonpickle's classes= parameter only widens scope; it does not restrict py/reduce.

Suggested remediation

Because infos may carry numpy objects, the fix needs a restricted decoder, not a jsonpickle flag:

• Preferred: decode infos with json.loads plus explicit, allowlisted reconstruction of only the expected shapes (e.g. numpy ndarrays via a known {"dtype", "shape", "values"} envelope), rejecting arbitrary py/* tags. This preserves the numpy round-trip while eliminating arbitrary code execution.
• Minimum: add an explicit trusted-input warning to the serialize.load() / serialize.save() docstrings (and ideally near flatten_trajectories), stating that loading a trajectory dataset from an untrusted source can execute code — consistent with how torch.load / pickle / np.load(allow_pickle=True) are documented. Optionally gate the reconstructing decode behind an opt-in trusted= / weights_only-style parameter that defaults to the safe path.

Severity / framing

This is the same trust class as torch.load / pickle.load / np.load(allow_pickle=True) — code execution when loading an untrusted artifact. I'm not requesting a CVE or claiming an authentication/authorization bypass. The ask is documentation + an optional restricted-decode path, because (a) the surface is undocumented, (b) it hides in a datasets Arrow directory that reads as inert data, (c) it triggers on ordinary preprocessing, and (d) the ecosystem is moving toward safe-by-default artifact loading (e.g. torch flipping weights_only=True).

Happy to send a PR for either the docstring warning or the restricted infos decoder if that's welcome.

Credit

Reported by Michael Holmquist (mtholmquist). Surfaced during automated dependency security review; the chain above was independently reproduced against imitation 1.0.1 / jsonpickle 4.1.2 / datasets 5.0.0.