From 0be7dfe2dfa5f41bdbbfa6b1b6c0838bdc9d41d7 Mon Sep 17 00:00:00 2001
From: Alex Manley <amanley@nvidia.com>
Date: Tue, 14 Jul 2026 14:26:58 -0700
Subject: [PATCH 01/14] [Trajectory] add new trajectory handling

---
 src/cloudai/configurator/trajectory.py | 352 +++++++++++++++++++++++++
 tests/test_trajectory.py               | 246 +++++++++++++++++
 2 files changed, 598 insertions(+)
 create mode 100644 src/cloudai/configurator/trajectory.py
 create mode 100644 tests/test_trajectory.py

diff --git a/src/cloudai/configurator/trajectory.py b/src/cloudai/configurator/trajectory.py
new file mode 100644
index 000000000..194f9bdd8
--- /dev/null
+++ b/src/cloudai/configurator/trajectory.py
@@ -0,0 +1,352 @@
+# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Ordered trajectory steps composed from typed dataclass components."""
+
+from __future__ import annotations
+
+import csv
+import dataclasses
+import json
+import logging
+from collections.abc import Callable, Iterator, Mapping, Sequence
+from pathlib import Path
+from typing import Any, Protocol, TypeVar, cast, overload
+
+ComponentT = TypeVar("ComponentT")
+
+
+@dataclasses.dataclass(frozen=True)
+class EnvParamsSample:
+    """Environment-parameter values sampled for one trial."""
+
+    env_params: dict[str, Any]
+
+
+@dataclasses.dataclass(frozen=True)
+class TrialResult:
+    """The action and resulting values required for every trajectory step."""
+
+    action: Mapping[str, Any]
+    reward: float
+    observation: Sequence[Any]
+
+
+@dataclasses.dataclass(frozen=True)
+class TrajectoryEntry:
+    """One immutable step containing a fixed set of typed data components."""
+
+    step: int
+    components: tuple[object, ...]
+
+    def __post_init__(self) -> None:
+        """Validate the trial index and component uniqueness."""
+        if self.step < 1:
+            raise ValueError(f"trajectory step must be positive; got {self.step}")
+        _components_by_type(self.components)
+
+    def get(self, component_type: type[ComponentT]) -> ComponentT | None:
+        """Return the component with exactly the requested type, if present."""
+        for component in self.components:
+            if type(component) is component_type:
+                return cast(ComponentT, component)
+        return None
+
+
+class TrajectoryWriter(Protocol):
+    """Persistence boundary for one flattened trajectory record."""
+
+    @property
+    def output_path(self) -> Path: ...
+
+    def append(self, record: Mapping[str, object]) -> None:
+        """Persist one record."""
+
+
+class _FileTrajectoryWriter:
+    """Resolve a writer-specific filename beneath an iteration directory."""
+
+    file_name = ""
+
+    def __init__(self, iteration_dir: Path | Callable[[], Path]) -> None:
+        self._iteration_dir = iteration_dir
+
+    @property
+    def output_path(self) -> Path:
+        iteration_dir = self._iteration_dir() if callable(self._iteration_dir) else self._iteration_dir
+        return iteration_dir / self.file_name
+
+
+class CsvTrajectoryWriter(_FileTrajectoryWriter):
+    """Append trajectory records to trajectory.csv."""
+
+    file_name = "trajectory.csv"
+
+    def __init__(self, iteration_dir: Path | Callable[[], Path]) -> None:
+        super().__init__(iteration_dir)
+        self._fields: tuple[str, ...] | None = None
+
+    def append(self, record: Mapping[str, object]) -> None:
+        fields = tuple(record)
+        if self._fields is None:
+            self._fields = fields
+        elif fields != self._fields:
+            raise ValueError(f"trajectory record fields changed: expected {self._fields}, got {fields}")
+
+        path = self.output_path
+        new_file = not path.exists()
+        path.parent.mkdir(parents=True, exist_ok=True)
+        with path.open("a", newline="") as file:
+            writer = csv.DictWriter(file, fieldnames=self._fields)
+            if new_file:
+                writer.writeheader()
+            writer.writerow(record)
+        logging.debug("Wrote trajectory record to %s.", path)
+
+
+class JsonLinesTrajectoryWriter(_FileTrajectoryWriter):
+    """Append trajectory records to trajectory.jsonl as newline-delimited JSON objects."""
+
+    file_name = "trajectory.jsonl"
+
+    def append(self, record: Mapping[str, object]) -> None:
+        path = self.output_path
+        path.parent.mkdir(parents=True, exist_ok=True)
+        with path.open("a") as file:
+            file.write(json.dumps(record))
+            file.write("\n")
+        logging.debug("Wrote trajectory record to %s.", path)
+
+
+class Trajectory(Sequence[TrajectoryEntry]):
+    """
+    Ordered entries for one DSE iteration with one fixed component schema.
+
+    ``components`` declares optional data types every entry must contain;
+    :class:`TrialResult` is always included.
+    ``identity`` is the subset that affects trial equivalence;
+    informational components such as logs and metrics can be excluded.
+
+    Steps must be appended in increasing order, but gaps are permitted because
+    CloudAI does not record constraint-failed trials.
+    """
+
+    def __init__(
+        self,
+        entries: Sequence[TrajectoryEntry] = (),
+        *,
+        writer: TrajectoryWriter | None = None,
+        components: Sequence[type[object]] = (),
+        identity: Sequence[type[object]] = (),
+    ) -> None:
+        self._component_types = (TrialResult, *components)
+        self._components = frozenset(self._component_types)
+        self._identity = frozenset(identity)
+        if len(self._components) != len(self._component_types):
+            raise ValueError("components cannot contain duplicate types")
+        if len(self._identity) != len(identity):
+            raise ValueError("identity cannot contain duplicate types")
+
+        undeclared_identity_types = self._identity - self._components
+        if undeclared_identity_types:
+            names = ", ".join(sorted(component_type.__name__ for component_type in undeclared_identity_types))
+            raise ValueError(f"identity types must be declared in components: {names}")
+
+        self._writer = writer
+        self._fields_by_type = self._build_fields_by_type()
+
+        self._entries: list[TrajectoryEntry] = []
+        for entry in entries:
+            self._store_entry(entry, persist=False)
+
+        component_names = ", ".join(component_type.__name__ for component_type in self._component_types)
+        logging.debug(
+            "Initialized Trajectory with component types %s and %s entries.",
+            component_names,
+            len(self),
+        )
+
+    @overload
+    def __getitem__(self, index: int) -> TrajectoryEntry: ...
+
+    @overload
+    def __getitem__(self, index: slice) -> list[TrajectoryEntry]: ...
+
+    def __getitem__(self, index: int | slice) -> TrajectoryEntry | list[TrajectoryEntry]:
+        """Return one entry or a list containing an entry slice."""
+        return self._entries[index]
+
+    def __len__(self) -> int:
+        """Return the number of recorded entries."""
+        return len(self._entries)
+
+    def __iter__(self) -> Iterator[TrajectoryEntry]:
+        """Iterate over entries in step order."""
+        return iter(self._entries)
+
+    @property
+    def output_path(self) -> Path | None:
+        """Return the writer's current output path, if persistence is configured."""
+        return self._writer.output_path if self._writer is not None else None
+
+    def append(self, *, step: int, **values: object) -> TrajectoryEntry:
+        """Build configured components from values, then store and persist one entry."""
+        components = self._construct_components(self._component_types, values, "trajectory values")
+        entry = TrajectoryEntry(step=step, components=components)
+        self._store_entry(entry, persist=True)
+        return entry
+
+    def _store_entry(self, entry: TrajectoryEntry, *, persist: bool) -> None:
+        self._validate_entry_components(entry)
+        if self._entries and entry.step <= self._entries[-1].step:
+            raise ValueError(f"trajectory steps must increase: last step is {self._entries[-1].step}, got {entry.step}")
+        if persist and self._writer is not None:
+            self._writer.append(self._to_record(entry))
+        self._entries.append(entry)
+        logging.debug("Appended trajectory entry for step %s (total entries: %s).", entry.step, len(self))
+
+    def find(self, action: Mapping[str, Any], **identity_values: object) -> TrajectoryEntry | None:
+        identity_components = self._construct_components(
+            tuple(component_type for component_type in self._component_types if component_type in self._identity),
+            identity_values,
+            "trajectory identity values",
+        )
+        identity = self._identity_for(identity_components)
+        for entry in self._entries:
+            result = entry.get(TrialResult)
+            if result is None:
+                raise ValueError(f"trajectory entry at step {entry.step} is missing TrialResult")
+            if _values_match_exact(result.action, action) and _values_match_exact(
+                self._identity_for(entry.components), identity
+            ):
+                logging.debug("Found matching trajectory entry at step %s for action %s.", entry.step, action)
+                return entry
+        logging.debug("No matching trajectory entry found for action %s.", action)
+        return None
+
+    def _validate_entry_components(self, entry: TrajectoryEntry) -> None:
+        _validate_schema(
+            frozenset(type(component) for component in entry.components),
+            self._components,
+            "trajectory entry components",
+        )
+
+    def _identity_for(self, components: Sequence[object]) -> dict[type[object], object]:
+        return {type(component): component for component in components if type(component) in self._identity}
+
+    def _construct_components(
+        self,
+        component_types: Sequence[type[object]],
+        values: Mapping[str, object],
+        context: str,
+    ) -> tuple[object, ...]:
+        fields_by_type = {
+            component_type: tuple(field for field in self._fields_by_type[component_type] if field.init)
+            for component_type in component_types
+        }
+        expected = {field.name for fields in fields_by_type.values() for field in fields}
+        required = {
+            field.name
+            for fields in fields_by_type.values()
+            for field in fields
+            if field.default is dataclasses.MISSING and field.default_factory is dataclasses.MISSING
+        }
+        actual = set(values)
+        missing = required - actual
+        unexpected = actual - expected
+        if missing or unexpected:
+            details = []
+            if missing:
+                details.append(f"missing: {', '.join(sorted(missing))}")
+            if unexpected:
+                details.append(f"unexpected: {', '.join(sorted(unexpected))}")
+            raise TypeError(f"{context} do not match configured schema ({'; '.join(details)})")
+
+        components = []
+        for component_type, fields in fields_by_type.items():
+            kwargs = {field.name: values[field.name] for field in fields if field.name in values}
+            constructor = cast(Any, component_type)
+            components.append(constructor(**kwargs))
+        return tuple(components)
+
+    def _build_fields_by_type(self) -> dict[type[object], tuple[dataclasses.Field[Any], ...]]:
+        fields_by_type: dict[type[object], tuple[dataclasses.Field[Any], ...]] = {}
+        for component_type in self._component_types:
+            if not dataclasses.is_dataclass(component_type):
+                raise TypeError(f"trajectory component type {component_type.__name__} must be a dataclass")
+            fields_by_type[component_type] = dataclasses.fields(component_type)
+
+        fields = ("step", *(field.name for component_fields in fields_by_type.values() for field in component_fields))
+        if len(fields) != len(set(fields)):
+            raise ValueError("trajectory record fields must be unique")
+        return fields_by_type
+
+    def _to_record(self, entry: TrajectoryEntry) -> dict[str, object]:
+        record: dict[str, object] = {"step": entry.step}
+        components_by_type = {type(component): component for component in entry.components}
+        for component_type in self._component_types:
+            component = components_by_type[component_type]
+            record.update(
+                {field.name: getattr(component, field.name) for field in self._fields_by_type[component_type]}
+            )
+        return record
+
+
+def _components_by_type(components: Sequence[object]) -> dict[type[object], object]:
+    non_dataclasses = [type(component).__name__ for component in components if not dataclasses.is_dataclass(component)]
+    if non_dataclasses:
+        raise TypeError(f"trajectory components must be dataclass instances: {', '.join(non_dataclasses)}")
+    by_type = {type(component): component for component in components}
+    if len(by_type) != len(components):
+        raise ValueError("components cannot contain duplicate component types")
+    return by_type
+
+
+def _validate_schema(
+    actual: frozenset[type[object]],
+    expected: frozenset[type[object]],
+    context: str,
+) -> None:
+    if actual == expected:
+        return
+
+    details = []
+    missing = expected - actual
+    unexpected = actual - expected
+    if missing:
+        details.append(f"missing: {', '.join(sorted(component_type.__name__ for component_type in missing))}")
+    if unexpected:
+        details.append(f"unexpected: {', '.join(sorted(component_type.__name__ for component_type in unexpected))}")
+    raise TypeError(f"{context} do not match configured schema ({'; '.join(details)})")
+
+
+def _values_match_exact(left: Any, right: Any) -> bool:
+    if type(left) is not type(right):
+        return False
+    if dataclasses.is_dataclass(left) and not isinstance(left, type):
+        return all(
+            _values_match_exact(getattr(left, field.name), getattr(right, field.name))
+            for field in dataclasses.fields(left)
+        )
+    if isinstance(left, Mapping):
+        if set(left) != set(right):
+            return False
+        return all(_values_match_exact(left[key], right[key]) for key in left)
+    if isinstance(left, (list, tuple)):
+        return len(left) == len(right) and all(
+            _values_match_exact(left_item, right_item) for left_item, right_item in zip(left, right, strict=True)
+        )
+    return left == right
diff --git a/tests/test_trajectory.py b/tests/test_trajectory.py
new file mode 100644
index 000000000..e78e59747
--- /dev/null
+++ b/tests/test_trajectory.py
@@ -0,0 +1,246 @@
+# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import dataclasses
+import json
+from collections.abc import Mapping
+from pathlib import Path
+from typing import Any
+
+import pytest
+
+from cloudai.configurator.trajectory import (
+    CsvTrajectoryWriter,
+    EnvParamsSample,
+    JsonLinesTrajectoryWriter,
+    Trajectory,
+    TrajectoryEntry,
+    TrialResult,
+)
+
+
+@dataclasses.dataclass(frozen=True)
+class LoggingMetrics:
+    logging_metrics: Mapping[str, float]
+
+
+class RecordingWriter:
+    def __init__(self, output_path: Path, *, fail: bool = False) -> None:
+        self.output_path = output_path
+        self.fail = fail
+        self.records: list[Mapping[str, object]] = []
+
+    def append(self, record: Mapping[str, object]) -> None:
+        if self.fail:
+            raise OSError("write failed")
+        self.records.append(record)
+
+
+def _entry(step: int, action: Mapping[str, Any] | None = None) -> TrajectoryEntry:
+    return TrajectoryEntry(
+        step=step,
+        components=(TrialResult(action=action or {"x": step}, reward=float(step), observation=[step]),),
+    )
+
+
+def _extended_entry(step: int, speed: int, power: float = 600.0) -> TrajectoryEntry:
+    return TrajectoryEntry(
+        step=step,
+        components=(
+            TrialResult(action={"x": 1}, reward=float(step), observation=[step]),
+            EnvParamsSample({"speed": speed}),
+            LoggingMetrics({"gpu_power_watts": power}),
+        ),
+    )
+
+
+def test_trajectory_is_an_ordered_sequence() -> None:
+    trajectory = Trajectory([_entry(1), _entry(3)])
+
+    assert len(trajectory) == 2
+    assert [entry.step for entry in trajectory] == [1, 3]
+    assert trajectory[0].step == 1
+    assert [entry.step for entry in trajectory[:]] == [1, 3]
+
+
+def test_trajectory_rejects_non_increasing_steps() -> None:
+    trajectory = Trajectory([_entry(2)])
+
+    with pytest.raises(ValueError, match="steps must increase"):
+        trajectory.append(step=2, action={"x": 2}, reward=2.0, observation=[2])
+
+
+def test_writer_failure_does_not_append_entry_to_memory(tmp_path: Path) -> None:
+    trajectory = Trajectory(writer=RecordingWriter(tmp_path / "trajectory", fail=True))
+
+    with pytest.raises(OSError, match="write failed"):
+        trajectory.append(step=1, action={"x": 1}, reward=1.0, observation=[1])
+
+    assert len(trajectory) == 0
+
+
+def test_initial_entries_are_not_replayed_to_writer(tmp_path: Path) -> None:
+    writer = RecordingWriter(tmp_path / "trajectory")
+
+    trajectory = Trajectory([_entry(1)], writer=writer)
+
+    assert len(trajectory) == 1
+    assert writer.records == []
+
+
+def test_append_writes_component_values_to_csv(tmp_path: Path) -> None:
+    path = tmp_path / "trajectory.csv"
+    trajectory = Trajectory(
+        writer=CsvTrajectoryWriter(tmp_path),
+        components=(LoggingMetrics,),
+    )
+
+    trajectory.append(step=1, action={"x": 1}, reward=1.0, observation=[1], logging_metrics={"power": 600.0})
+    trajectory.append(step=2, action={"x": 2}, reward=2.0, observation=[2], logging_metrics={"power": 610.0})
+
+    assert path.read_text().splitlines() == [
+        "step,action,reward,observation,logging_metrics",
+        "1,{'x': 1},1.0,[1],{'power': 600.0}",
+        "2,{'x': 2},2.0,[2],{'power': 610.0}",
+    ]
+
+
+def test_append_writes_generic_records_as_json_lines(tmp_path: Path) -> None:
+    path = tmp_path / "trajectory.jsonl"
+    trajectory = Trajectory(
+        writer=JsonLinesTrajectoryWriter(tmp_path),
+        components=(EnvParamsSample, LoggingMetrics),
+    )
+
+    trajectory.append(
+        step=1,
+        action={"x": 1},
+        reward=1.0,
+        observation=[1],
+        env_params={"speed": 8},
+        logging_metrics={"gpu_power_watts": 610.0},
+    )
+
+    assert json.loads(path.read_text()) == {
+        "step": 1,
+        "action": {"x": 1},
+        "reward": 1.0,
+        "observation": [1],
+        "env_params": {"speed": 8},
+        "logging_metrics": {"gpu_power_watts": 610.0},
+    }
+
+
+def test_entry_contains_and_retrieves_typed_components() -> None:
+    result = TrialResult({"x": 1}, 1.0, [1])
+    env_params = EnvParamsSample({"speed": 1})
+    metrics = LoggingMetrics({"gpu_power_watts": 600.0})
+    entry = TrajectoryEntry(step=1, components=(result, env_params, metrics))
+
+    assert entry.components == (result, env_params, metrics)
+    assert entry.get(TrialResult) is result
+    assert entry.get(EnvParamsSample) is env_params
+    assert entry.get(LoggingMetrics) is metrics
+
+
+def test_entry_rejects_duplicate_component_types() -> None:
+    with pytest.raises(ValueError, match="duplicate component types"):
+        TrajectoryEntry(
+            step=1,
+            components=(EnvParamsSample({"speed": 1}), EnvParamsSample({"speed": 2})),
+        )
+
+
+def test_trajectory_validates_its_fixed_component_schema() -> None:
+    trajectory = Trajectory(components=(EnvParamsSample, LoggingMetrics))
+
+    with pytest.raises(TypeError, match="missing: logging_metrics"):
+        trajectory.append(step=1, action={"x": 1}, reward=1.0, observation=[1], env_params={"speed": 1})
+
+    with pytest.raises(TypeError, match="unexpected: logging_metrics"):
+        Trajectory().append(
+            step=1,
+            action={"x": 1},
+            reward=1.0,
+            observation=[1],
+            logging_metrics={},
+        )
+
+
+def test_identity_component_types_must_be_declared() -> None:
+    with pytest.raises(ValueError, match=r"must be declared.*EnvParamsSample"):
+        Trajectory(identity=(EnvParamsSample,))
+
+
+def test_find_uses_only_configured_identity_components() -> None:
+    trajectory = Trajectory(
+        components=(EnvParamsSample, LoggingMetrics),
+        identity=(EnvParamsSample,),
+    )
+    first = trajectory.append(
+        step=1,
+        action={"x": 1},
+        reward=1.0,
+        observation=[1],
+        env_params={"speed": 1},
+        logging_metrics={"gpu_power_watts": 600.0},
+    )
+
+    assert trajectory.find({"x": 1}, env_params={"speed": 1}) is first
+    assert trajectory.find({"x": 1}, env_params={"speed": 2}) is None
+
+
+def test_find_ignores_informational_component_values() -> None:
+    trajectory = Trajectory(components=(LoggingMetrics,))
+    first = trajectory.append(
+        step=1,
+        action={"x": 1},
+        reward=1.0,
+        observation=[1],
+        logging_metrics={"gpu_power_watts": 600.0},
+    )
+
+    assert trajectory.find({"x": 1}) is first
+
+
+def test_find_requires_the_configured_identity_component_types() -> None:
+    trajectory = Trajectory(
+        components=(EnvParamsSample,),
+        identity=(EnvParamsSample,),
+    )
+
+    with pytest.raises(TypeError, match="missing: env_params"):
+        trajectory.find({"x": 1})
+
+    with pytest.raises(TypeError, match="unexpected: logging_metrics"):
+        trajectory.find({"x": 1}, logging_metrics={})
+
+
+def test_find_preserves_exact_value_types_inside_components() -> None:
+    trajectory = Trajectory(
+        components=(EnvParamsSample,),
+        identity=(EnvParamsSample,),
+    )
+    trajectory.append(step=1, action={"x": 1}, reward=1.0, observation=[1], env_params={"speed": 1.0})
+
+    assert trajectory.find({"x": 1}, env_params={"speed": 1}) is None
+
+
+def test_find_preserves_exact_action_value_types() -> None:
+    trajectory = Trajectory([_entry(1, {"x": 1.0})])
+
+    assert trajectory.find({"x": 1}) is None
+
+
+def test_trajectory_logs_lifecycle_and_lookup(caplog: pytest.LogCaptureFixture) -> None:
+    with caplog.at_level("DEBUG"):
+        trajectory = Trajectory()
+        entry = trajectory.append(step=1, action={"x": 1}, reward=1.0, observation=[1])
+        assert trajectory.find({"x": 1}) is entry
+        assert trajectory.find({"x": 2}) is None
+
+    assert "Initialized Trajectory with component types TrialResult and 0 entries." in caplog.messages
+    assert "Appended trajectory entry for step 1 (total entries: 1)." in caplog.messages
+    assert "Found matching trajectory entry at step 1 for action {'x': 1}." in caplog.messages
+    assert "No matching trajectory entry found for action {'x': 2}." in caplog.messages

From bef543e5bf492f75a195aad53b7eb019b5076a6c Mon Sep 17 00:00:00 2001
From: Alex Manley <amanley@nvidia.com>
Date: Tue, 14 Jul 2026 14:27:34 -0700
Subject: [PATCH 02/14] [Trajectory] update gym and others to support new
 trajectory

---
 src/cloudai/configurator/__init__.py          |  17 +-
 src/cloudai/configurator/cloudai_gym.py       | 219 +++++++-----------
 src/cloudai/configurator/env_params.py        |  26 ---
 src/cloudai/models/workload.py                |   2 +-
 src/cloudai/report_generator/dse_report.py    |  23 +-
 src/cloudai/reporter.py                       |  11 +-
 .../systems/slurm/single_sbatch_runner.py     |  27 +--
 7 files changed, 142 insertions(+), 183 deletions(-)

diff --git a/src/cloudai/configurator/__init__.py b/src/cloudai/configurator/__init__.py
index a88432c41..9d00017cd 100644
--- a/src/cloudai/configurator/__init__.py
+++ b/src/cloudai/configurator/__init__.py
@@ -16,15 +16,30 @@
 
 from .base_agent import BaseAgent
 from .base_gym import BaseGym
-from .cloudai_gym import CloudAIGymEnv, TrajectoryEntry
+from .cloudai_gym import CloudAIGymEnv
 from .grid_search import GridSearchAgent
 from .gymnasium_adapter import GymnasiumAdapter
+from .trajectory import (
+    CsvTrajectoryWriter,
+    EnvParamsSample,
+    JsonLinesTrajectoryWriter,
+    Trajectory,
+    TrajectoryEntry,
+    TrajectoryWriter,
+    TrialResult,
+)
 
 __all__ = [
     "BaseAgent",
     "BaseGym",
     "CloudAIGymEnv",
+    "CsvTrajectoryWriter",
+    "EnvParamsSample",
     "GridSearchAgent",
     "GymnasiumAdapter",
+    "JsonLinesTrajectoryWriter",
+    "Trajectory",
     "TrajectoryEntry",
+    "TrajectoryWriter",
+    "TrialResult",
 ]
diff --git a/src/cloudai/configurator/cloudai_gym.py b/src/cloudai/configurator/cloudai_gym.py
index 21b3f82af..848824e26 100644
--- a/src/cloudai/configurator/cloudai_gym.py
+++ b/src/cloudai/configurator/cloudai_gym.py
@@ -15,29 +15,25 @@
 # limitations under the License.
 
 import copy
-import csv
-import dataclasses
 import logging
 from pathlib import Path
-from typing import Any, Dict, Optional, Tuple
+from typing import Any, Dict, Literal, Optional, Tuple
 
 from cloudai.core import METRIC_ERROR, BaseRunner, Registry, TestRun
 from cloudai.util.lazy_imports import lazy
 
 from .base_agent import RewardOverrides
 from .base_gym import BaseGym
-from .env_params import EnvParams, ObsLeafDescriptor, write_env_params
-
-
-@dataclasses.dataclass(frozen=True)
-class TrajectoryEntry:
-    """Represents a trajectory entry."""
-
-    step: int
-    action: dict[str, Any]
-    reward: float
-    observation: list
-    env_params: dict[str, Any] = dataclasses.field(default_factory=dict)
+from .env_params import EnvParams, ObsLeafDescriptor
+from .trajectory import (
+    CsvTrajectoryWriter,
+    EnvParamsSample,
+    JsonLinesTrajectoryWriter,
+    Trajectory,
+    TrajectoryEntry,
+    TrajectoryWriter,
+    TrialResult,
+)
 
 
 class CloudAIGymEnv(BaseGym):
@@ -47,7 +43,14 @@ class CloudAIGymEnv(BaseGym):
     Uses the TestRun object and actual runner methods to execute jobs.
     """
 
-    def __init__(self, test_run: TestRun, runner: BaseRunner, rewards: RewardOverrides):
+    def __init__(
+        self,
+        test_run: TestRun,
+        runner: BaseRunner,
+        rewards: RewardOverrides,
+        *,
+        trajectory_file_type: Literal["csv", "jsonl"] = "jsonl",
+    ):
         """
         Initialize the Gym environment using the TestRun object.
 
@@ -55,6 +58,7 @@ def __init__(self, test_run: TestRun, runner: BaseRunner, rewards: RewardOverrid
             test_run (TestRun): A test run object that encapsulates cmd_args, extra_cmd_args, etc.
             runner (BaseRunner): The runner object to execute jobs.
             rewards: Reward / observation overrides from agent config.
+            trajectory_file_type: Format used to persist trajectory records.
         """
         self.test_run = test_run
         self.original_test_run = copy.deepcopy(test_run)  # Preserve clean state for DSE
@@ -62,15 +66,10 @@ def __init__(self, test_run: TestRun, runner: BaseRunner, rewards: RewardOverrid
         self.rewards = rewards
         self.max_steps = test_run.test.agent_steps
         self.reward_function = Registry().get_reward_function(test_run.test.agent_reward_function)
-        self.trajectory: dict[int, list[TrajectoryEntry]] = {}
         self.params: EnvParams | None = EnvParams.from_test(test_run.test)
+        self.trajectory = self._new_trajectory(trajectory_file_type)
         super().__init__()
 
-    @property
-    def env_params_record_path(self) -> Path:
-        """``env.csv`` lives alongside ``trajectory.csv`` so a plain ``merge`` joins them."""
-        return self.iteration_dir / "env.csv"
-
     @property
     def upcoming_trial(self) -> int:
         """
@@ -104,7 +103,7 @@ def define_observation_space(self) -> list:
     def reset(
         self,
         seed: Optional[int] = None,
-        options: Optional[dict[str, Any]] = None,  # noqa: Vulture
+        options: Optional[dict[str, Any]] = None,
     ) -> Tuple[list, dict[str, Any]]:
         """
         Reset the environment and reinitialize the TestRun.
@@ -118,6 +117,7 @@ def reset(
                 - observation (list): Initial observation.
                 - info (dict): Additional info for debugging.
         """
+        del options
         if seed is not None:
             lazy.np.random.seed(seed)
         self.test_run.current_iteration = 0
@@ -150,61 +150,56 @@ def step(self, action: Any) -> Tuple[list, float, bool, dict]:
         cached_result = self.get_cached_trajectory_result(action, sampled_env_params)
 
         if cached_result is not None:
+            cached_trial_result = cached_result.get(TrialResult)
+            if cached_trial_result is None:
+                raise ValueError(f"cached trajectory entry at step {cached_result.step} is missing TrialResult")
             logging.info(
                 "Retrieved cached result from trajectory with reward %s (from step %s). Skipping execution.",
-                cached_result.reward,
+                cached_trial_result.reward,
                 cached_result.step,
             )
-            self.write_trajectory(
-                TrajectoryEntry(
-                    step=self.test_run.step,
-                    action=action,
-                    reward=cached_result.reward,
-                    observation=cached_result.observation,
-                    env_params=sampled_env_params,
-                )
-            )
-
-            return cached_result.observation, cached_result.reward, False, info
-
-        if not self.test_run.test.constraint_check(self.test_run, self.runner.system):
-            logging.info("Constraint check failed. Skipping step.")
-            return [-1.0], self.rewards.constraint_failure, True, info
-
-        new_tr = copy.deepcopy(self.test_run)
-        new_tr.output_path = self.runner.get_job_output_path(new_tr)
-        self.runner.test_scenario.test_runs = [new_tr]
-
-        self.runner.shutting_down = False
-        self.runner.jobs.clear()
-        self.runner.testrun_to_job_map.clear()
-
-        try:
-            self.runner.run()
-        except Exception as e:
-            logging.error(f"Error running step {self.test_run.step}: {e}")
-
-        if self.runner.test_scenario.test_runs and self.runner.test_scenario.test_runs[0].output_path.exists():
-            self.test_run = self.runner.test_scenario.test_runs[0]
+            observation = list(cached_trial_result.observation)
+            reward = cached_trial_result.reward
         else:
-            self.test_run = copy.deepcopy(self.original_test_run)
-            self.test_run.step = new_tr.step
-            self.test_run.output_path = new_tr.output_path
-
-        metrics = self.get_observation(action)
-        reward = self.compute_reward(metrics)
-
-        self.write_trajectory(
-            TrajectoryEntry(
-                step=self.test_run.step,
-                action=action,
-                reward=reward,
-                observation=metrics,
-                env_params=sampled_env_params,
-            )
+            if not self.test_run.test.constraint_check(self.test_run, self.runner.system):
+                logging.info("Constraint check failed. Skipping step.")
+                return [-1.0], self.rewards.constraint_failure, True, info
+
+            new_tr = copy.deepcopy(self.test_run)
+            new_tr.output_path = self.runner.get_job_output_path(new_tr)
+            self.runner.test_scenario.test_runs = [new_tr]
+
+            self.runner.shutting_down = False
+            self.runner.jobs.clear()
+            self.runner.testrun_to_job_map.clear()
+
+            try:
+                self.runner.run()
+            except Exception as e:
+                logging.error(f"Error running step {self.test_run.step}: {e}")
+
+            if self.runner.test_scenario.test_runs and self.runner.test_scenario.test_runs[0].output_path.exists():
+                self.test_run = self.runner.test_scenario.test_runs[0]
+            else:
+                self.test_run = copy.deepcopy(self.original_test_run)
+                self.test_run.step = new_tr.step
+                self.test_run.output_path = new_tr.output_path
+
+            observation = self.get_observation(action)
+            reward = self.compute_reward(observation)
+
+        optional_values: dict[str, object] = {}
+        if self.params is not None:
+            optional_values["env_params"] = dict(sampled_env_params)
+        self.trajectory.append(
+            step=self.test_run.step,
+            action=dict(action),
+            reward=reward,
+            observation=observation,
+            **optional_values,
         )
 
-        return metrics, reward, False, info
+        return observation, reward, False, info
 
     def render(self, mode: str = "human"):
         """
@@ -277,41 +272,35 @@ def encode_env_params(self, env_params: dict[str, Any]) -> Dict[str, Any]:
         """
         return self.params.encode(env_params) if self.params is not None else {}
 
-    def write_trajectory(self, entry: TrajectoryEntry):
-        """
-        Append the entry to the in-memory cache and trajectory.csv (plus env.csv when declared).
-
-        ``trajectory.csv`` and the ``env.csv`` projection are sunk from the same
-        ``TrajectoryEntry`` here, so a trial that never produces an entry (e.g. a
-        constraint failure returns before this call) lands in neither file and the
-        two stay 1:1 step-aligned.
-        """
-        self.current_trajectory.append(entry)
-
-        file_exists = self.trajectory_file_path.exists()
-        logging.debug(f"Writing trajectory into {self.trajectory_file_path} (exists: {file_exists})")
-        self.trajectory_file_path.parent.mkdir(parents=True, exist_ok=True)
+    def _new_trajectory(self, file_type: Literal["csv", "jsonl"]) -> Trajectory:
+        writer: TrajectoryWriter
+        if file_type == "csv":
+            writer = CsvTrajectoryWriter(lambda: self.iteration_dir)
+        elif file_type == "jsonl":
+            writer = JsonLinesTrajectoryWriter(lambda: self.iteration_dir)
+        else:
+            raise ValueError(f"Invalid file type: {file_type}")
 
-        with open(self.trajectory_file_path, mode="a", newline="") as file:
-            writer = csv.writer(file)
-            if not file_exists:
-                writer.writerow(["step", "action", "reward", "observation"])
-            writer.writerow([entry.step, entry.action, entry.reward, entry.observation])
+        if self.params is None:
+            return Trajectory(writer=writer)
 
-        write_env_params(self.env_params_record_path, entry.step, entry.env_params)
+        return Trajectory(
+            writer=writer,
+            components=(EnvParamsSample,),
+            identity=(EnvParamsSample,),
+        )
 
     @property
     def iteration_dir(self) -> Path:
-        """Per-iteration output dir; trajectory.csv and env.csv both live here, step-aligned."""
+        """Per-iteration output directory containing the trajectory output."""
         return self.runner.scenario_root / self.test_run.name / f"{self.test_run.current_iteration}"
 
     @property
     def trajectory_file_path(self) -> Path:
-        return self.iteration_dir / "trajectory.csv"
-
-    @property
-    def current_trajectory(self) -> list[TrajectoryEntry]:
-        return self.trajectory.setdefault(self.test_run.current_iteration, [])
+        path = self.trajectory.output_path
+        if path is None:
+            raise RuntimeError("trajectory persistence is not configured")
+        return path
 
     def get_cached_trajectory_result(self, action: Any, env_params: dict[str, Any]) -> TrajectoryEntry | None:
         """
@@ -324,36 +313,6 @@ def get_cached_trajectory_result(self, action: Any, env_params: dict[str, Any])
         do not declare any ``[env_params.*]`` block. The sample is passed in (a
         per-trial local owned by ``step``), exactly like ``action``.
         """
-        for entry in self.current_trajectory:
-            action_match = self._values_match_exact(entry.action, action)
-            env_params_match = self._values_match_exact(entry.env_params, env_params)
-            if action_match and env_params_match:
-                return entry
-
-        return None
-
-    @classmethod
-    def _values_match_exact(cls, left: Any, right: Any) -> bool:
-        if type(left) is not type(right):
-            return False
-
-        elif isinstance(left, dict):
-            left_keys = set(left.keys())
-            right_keys = set(right.keys())
-            if left_keys != right_keys:
-                return False
-
-            return all(cls._values_match_exact(left[key], right[key]) for key in left_keys)
-
-        elif isinstance(left, (list, tuple)):
-            if len(left) != len(right):
-                return False
-
-            for left_item, right_item in zip(left, right, strict=True):
-                if not cls._values_match_exact(left_item, right_item):
-                    return False
-
-            return True
-
-        else:
-            return left == right
+        if not env_params:
+            return self.trajectory.find(action)
+        return self.trajectory.find(action, env_params=dict(env_params))
diff --git a/src/cloudai/configurator/env_params.py b/src/cloudai/configurator/env_params.py
index a824d6205..af53b8a49 100644
--- a/src/cloudai/configurator/env_params.py
+++ b/src/cloudai/configurator/env_params.py
@@ -27,11 +27,9 @@
 
 from __future__ import annotations
 
-import csv
 import dataclasses
 import math
 import random
-from pathlib import Path
 from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Protocol, runtime_checkable
 
 from pydantic import BaseModel, ConfigDict, Field, model_validator
@@ -239,30 +237,6 @@ def observation_descriptors(self) -> Dict[str, ObsLeafDescriptor]:
         return {name: param.observation_descriptor() for name, param in self.params.items()}
 
 
-def write_env_params(path: Path, step: int, sample: Dict[str, Any]) -> None:
-    """
-    Append one trial's env_params sample to a step-aligned CSV.
-
-    The CSV mirrors how ``trajectory.csv`` serialises its ``action`` column
-    (one row per env.step(), sample dict stringified in a single cell) so the
-    two files align 1:1 on ``step`` and a plain ``merge`` joins them.
-
-    Empty samples are skipped, so a run without env_params writes nothing and
-    callers can sink every trial unconditionally.
-    """
-    if step < 1:
-        raise ValueError(f"step must be a positive trial index (cloudai DSE is 1-based); got {step}")
-    if not sample:
-        return
-    new_file = not path.exists()
-    path.parent.mkdir(parents=True, exist_ok=True)
-    with path.open("a", newline="") as f:
-        writer = csv.writer(f)
-        if new_file:
-            writer.writerow(("step", "env"))
-        writer.writerow([step, sample])
-
-
 def validate_domain_randomization_active(test_scenario: "TestScenario") -> None:
     """
     Reject prepped configs that declare domain randomization no agent will run.
diff --git a/src/cloudai/models/workload.py b/src/cloudai/models/workload.py
index 72416afaa..975237540 100644
--- a/src/cloudai/models/workload.py
+++ b/src/cloudai/models/workload.py
@@ -127,7 +127,7 @@ class TestDefinition(BaseModel, ABC):
         description=(
             "Environment parameters sampled by the env per trial. Sibling to "
             "cmd_args; not part of the agent's action space. CloudAIGymEnv samples, "
-            "persists to env.csv, and includes them in the trajectory cache key."
+            "persists them in the trajectory output, and includes them in the trajectory cache key."
         ),
     )
 
diff --git a/src/cloudai/report_generator/dse_report.py b/src/cloudai/report_generator/dse_report.py
index efb85956e..3089d09cf 100644
--- a/src/cloudai/report_generator/dse_report.py
+++ b/src/cloudai/report_generator/dse_report.py
@@ -122,12 +122,27 @@ def format_money(value: float | None) -> str:
 
 
 def _safe_literal_eval(raw: Any, default: Any) -> Any:
+    if isinstance(raw, type(default)):
+        return raw
     if isinstance(raw, str):
         with contextlib.suppress(SyntaxError, ValueError):
             return ast.literal_eval(raw)
     return default
 
 
+def load_trajectory_dataframe(iteration_dir: Path) -> tuple[Path, Any] | None:
+    """Load trajectory output."""
+    jsonl_path = iteration_dir / "trajectory.jsonl"
+    if jsonl_path.is_file():
+        return jsonl_path, lazy.pd.read_json(jsonl_path, lines=True)
+
+    csv_path = iteration_dir / "trajectory.csv"
+    if csv_path.is_file():
+        return csv_path, lazy.pd.read_csv(csv_path)
+
+    return None
+
+
 def _format_scalar(value: Any) -> str:
     if isinstance(value, float):
         return f"{value:.4f}".rstrip("0").rstrip(".")
@@ -239,12 +254,12 @@ def _build_trajectory_steps(
     test_case: TestRun,
     test_runs: list[TestRun],
 ) -> list[TrajectoryStep] | None:
-    trajectory_file = iteration_dir / "trajectory.csv"
-    if not trajectory_file.is_file():
-        logging.warning(f"No trajectory file found for {test_case.name} at {trajectory_file}")
+    loaded_trajectory = load_trajectory_dataframe(iteration_dir)
+    if loaded_trajectory is None:
+        logging.warning(f"No trajectory file found for {test_case.name} in {iteration_dir}")
         return None
 
-    df = lazy.pd.read_csv(trajectory_file)
+    trajectory_file, df = loaded_trajectory
     if df.empty:
         logging.warning(f"No trajectory data found for {test_case.name} at {trajectory_file}")
         return None
diff --git a/src/cloudai/reporter.py b/src/cloudai/reporter.py
index a897015c3..a97ef2fc3 100644
--- a/src/cloudai/reporter.py
+++ b/src/cloudai/reporter.py
@@ -27,9 +27,8 @@
 from rich.console import Console
 from rich.table import Table
 
-from cloudai.report_generator.dse_report import build_dse_summaries
+from cloudai.report_generator.dse_report import build_dse_summaries, load_trajectory_dataframe
 from cloudai.report_generator.util import load_system_metadata
-from cloudai.util.lazy_imports import lazy
 
 from .core import CommandGenStrategy, Reporter, TestRun, case_name
 from .models.scenario import TestRunDetails
@@ -182,12 +181,12 @@ def report_best_dse_config(self):
                 continue
 
             tr_root = self.results_root / tr.name / f"{tr.current_iteration}"
-            trajectory_file = tr_root / "trajectory.csv"
-            if not trajectory_file.is_file():
-                logging.warning("No trajectory file found for %s at %s", tr.name, trajectory_file)
+            loaded_trajectory = load_trajectory_dataframe(tr_root)
+            if loaded_trajectory is None:
+                logging.warning("No trajectory file found for %s in %s", tr.name, tr_root)
                 continue
 
-            df = lazy.pd.read_csv(trajectory_file)
+            _, df = loaded_trajectory
             best_step = df.loc[df["reward"].idxmax()]["step"]
             best_step_details = tr_root / f"{best_step}" / CommandGenStrategy.TEST_RUN_DUMP_FILE_NAME
             if not best_step_details.is_file():
diff --git a/src/cloudai/systems/slurm/single_sbatch_runner.py b/src/cloudai/systems/slurm/single_sbatch_runner.py
index 8a3e769d9..9f73fb2af 100644
--- a/src/cloudai/systems/slurm/single_sbatch_runner.py
+++ b/src/cloudai/systems/slurm/single_sbatch_runner.py
@@ -21,7 +21,7 @@
 from pathlib import Path
 from typing import Generator, Optional, cast
 
-from cloudai.configurator import CloudAIGymEnv, TrajectoryEntry
+from cloudai.configurator import CloudAIGymEnv
 from cloudai.core import BaseJob, JobIdRetrievalError, Registry, System, TestRun, TestScenario
 from cloudai.util import CommandShell, format_time_limit, parse_time_limit
 
@@ -205,27 +205,24 @@ def handle_dse(self):
             if not tr.is_dse_job:
                 continue
 
+            agent_class = registry.get_agent(tr.test.agent)
+            agent_config_data = tr.test.agent_config or {}
+            agent_config = agent_class.get_config_class()(**agent_config_data)
+            gym = CloudAIGymEnv(tr, self, rewards=agent_config.rewards)
+
             for idx, combination in enumerate(tr.all_combinations, start=1):
                 next_tr = tr.apply_params_set(combination)
                 next_tr.step = idx
                 next_tr.output_path = self.get_job_output_path(next_tr)
 
-                rewards = None
-                agent_class = registry.get_agent(next_tr.test.agent)
-                agent_config_data = next_tr.test.agent_config or {}
-                agent_config = agent_class.get_config_class()(**agent_config_data)
-                rewards = agent_config.rewards
-
-                gym = CloudAIGymEnv(next_tr, self, rewards=rewards)
+                gym.test_run = next_tr
                 observation = gym.get_observation({})
                 reward = gym.compute_reward(observation)
-                gym.write_trajectory(
-                    TrajectoryEntry(
-                        step=idx,
-                        action=combination,
-                        reward=reward,
-                        observation=observation,
-                    )
+                gym.trajectory.append(
+                    step=idx,
+                    action=combination,
+                    reward=reward,
+                    observation=observation,
                 )
 
     def completed_test_runs(self, job: BaseJob) -> list[TestRun]:

From 840ff5b694be27a81e7a068724fbe2cd833dc06e Mon Sep 17 00:00:00 2001
From: Alex Manley <amanley@nvidia.com>
Date: Tue, 14 Jul 2026 14:27:52 -0700
Subject: [PATCH 03/14] [Trajectory] update tests

---
 tests/test_cloudaigym.py                 | 194 +++++++++++------------
 tests/test_env_params.py                 |  23 ---
 tests/test_gymnasium_adapter_contract.py |   2 +-
 tests/test_handlers.py                   |  10 +-
 tests/test_reporter.py                   |  15 +-
 tests/test_single_sbatch_runner.py       |   4 +-
 6 files changed, 119 insertions(+), 129 deletions(-)

diff --git a/tests/test_cloudaigym.py b/tests/test_cloudaigym.py
index dfc7ea7ae..6b2660ba0 100644
--- a/tests/test_cloudaigym.py
+++ b/tests/test_cloudaigym.py
@@ -14,13 +14,21 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import json
 from pathlib import Path
 from typing import cast
 from unittest.mock import MagicMock, PropertyMock, patch
 
 import pytest
 
-from cloudai.configurator import CloudAIGymEnv, GridSearchAgent, TrajectoryEntry
+from cloudai.configurator import (
+    CloudAIGymEnv,
+    EnvParamsSample,
+    GridSearchAgent,
+    Trajectory,
+    TrajectoryEntry,
+    TrialResult,
+)
 from cloudai.configurator.env_params import EnvParamSpec, ObsLeafDescriptor
 from cloudai.core import BaseRunner, RewardOverrides, Runner, TestRun, TestScenario
 from cloudai.systems.slurm import SlurmSystem
@@ -37,6 +45,19 @@
 from tests.test_env_params import EnvVarCmdArgs, EnvVarTestDefinition
 
 
+def _trajectory_entry(
+    step: int,
+    action: dict[str, object],
+    reward: float,
+    observation: list[float] | list[int],
+    *components: object,
+) -> TrajectoryEntry:
+    return TrajectoryEntry(
+        step=step,
+        components=(TrialResult(action=action, reward=reward, observation=observation), *components),
+    )
+
+
 @pytest.fixture
 def nemorun() -> NeMoRunTestDefinition:
     return NeMoRunTestDefinition(
@@ -353,31 +374,25 @@ def test_apply_params_set__preserves_installables_state(setup_env: tuple[TestRun
 
 
 @pytest.mark.parametrize(
-    ("trajectory", "current_iteration", "action", "expected_step"),
+    ("entries", "action", "expected_step"),
     [
-        ({}, 0, {"x": 1}, None),
-        ({0: [TrajectoryEntry(1, {"x": 1}, 1, [1])]}, 0, {"x": 1}, 1),
-        ({0: [TrajectoryEntry(1, {"x": 1.0}, 1, [1])]}, 0, {"x": 1}, None),
+        ([], {"x": 1}, None),
+        ([_trajectory_entry(1, {"x": 1}, 1, [1])], {"x": 1}, 1),
+        ([_trajectory_entry(1, {"x": 1.0}, 1, [1])], {"x": 1}, None),
         (
-            {
-                0: [
-                    TrajectoryEntry(1, {"x": 1.0}, 1, [1]),
-                    TrajectoryEntry(2, {"x": 1}, 1, [1]),
-                ]
-            },
-            0,
+            [
+                _trajectory_entry(1, {"x": 1.0}, 1, [1]),
+                _trajectory_entry(2, {"x": 1}, 1, [1]),
+            ],
             {"x": 1},
             2,
         ),
-        ({0: [TrajectoryEntry(1, {"x": 1}, 1, [1])]}, 1, {"x": 1}, None),
-        ({1: [TrajectoryEntry(3, {"x": 1}, 1, [1])]}, 1, {"x": 1}, 3),
     ],
 )
 def test_get_cached_trajectory_result(
     base_tr: TestRun,
     tmp_path: Path,
-    trajectory: dict[int, list[TrajectoryEntry]],
-    current_iteration: int,
+    entries: list[TrajectoryEntry],
     action: dict[str, object],
     expected_step: int | None,
 ) -> None:
@@ -390,8 +405,7 @@ def test_get_cached_trajectory_result(
     runner.get_job_output_path.return_value = tmp_path / "scenario" / base_tr.name / "0" / "7"
 
     env = CloudAIGymEnv(test_run=base_tr, runner=runner, rewards=RewardOverrides())
-    env.test_run.current_iteration = current_iteration
-    env.trajectory = trajectory
+    env.trajectory = Trajectory(entries)
 
     actual = env.get_cached_trajectory_result(action, {})
     if actual is None:
@@ -400,8 +414,8 @@ def test_get_cached_trajectory_result(
         assert actual.step == expected_step
 
 
-def test_cached_step_appends_trajectory_row(nemorun: NeMoRunTestDefinition, tmp_path: Path) -> None:
-    """Cache hits must still append a row to trajectory.csv so the visible step list matches agent_steps."""
+def test_cached_step_appends_trajectory_record(nemorun: NeMoRunTestDefinition, tmp_path: Path) -> None:
+    """Cache hits must still append a record so the visible step list matches agent_steps."""
     tdef = nemorun.model_copy(deep=True)
     tdef.cmd_args.data.global_batch_size = 8
     tdef.agent_metrics = ["default"]
@@ -420,7 +434,7 @@ def test_cached_step_appends_trajectory_row(nemorun: NeMoRunTestDefinition, tmp_
     env = CloudAIGymEnv(test_run=test_run, runner=runner, rewards=RewardOverrides())
     cached_action = {"trainer.max_steps": 1000}
     env.test_run.current_iteration = 0
-    env.trajectory = {0: [TrajectoryEntry(step=1, action=cached_action, reward=0.42, observation=[0.84])]}
+    env.trajectory.append(step=1, action=cached_action, reward=0.42, observation=[0.84])
 
     env.test_run.step = 4
     obs, reward, done, _info = env.step(cached_action)
@@ -429,29 +443,35 @@ def test_cached_step_appends_trajectory_row(nemorun: NeMoRunTestDefinition, tmp_
     assert reward == 0.42
     assert obs == [0.84]
     assert done is False
-    rows = env.trajectory[0]
+    rows = env.trajectory
     assert len(rows) == 2
     assert rows[-1].step == 5, (
         "CloudAIGymEnv.step() advances test_run.step before recording the trajectory row; "
         "the cached row must be tagged with the advanced trial index, not the pre-step value."
     )
-    assert rows[-1].reward == 0.42
-    assert rows[-1].action == cached_action
+    result = rows[-1].get(TrialResult)
+    assert result is not None
+    assert result.reward == 0.42
+    assert result.action == cached_action
 
-    csv_path = env.trajectory_file_path
-    assert csv_path.exists()
-    contents = csv_path.read_text().strip().splitlines()
-    assert contents[0] == "step,action,reward,observation"
-    assert contents[-1].startswith("5,")
+    trajectory_path = env.trajectory_file_path
+    assert trajectory_path.exists()
+    assert trajectory_path.name == "trajectory.jsonl"
+    records = [json.loads(line) for line in trajectory_path.read_text().splitlines()]
+    assert records[-1]["step"] == 5
+    assert records[-1]["action"] == cached_action
 
 
 def _seed_cached_entry_with_env_params(
     env: CloudAIGymEnv, action: dict[str, object], env_params: dict[str, object]
 ) -> None:
-    """Seed env.trajectory with one entry carrying the given env_params."""
-    entry = TrajectoryEntry(step=1, action=action, reward=0.5, observation=[100.0], env_params=env_params)
-    env.test_run.current_iteration = 0
-    env.trajectory = {0: [entry]}
+    """Seed an environment-parameter-aware trajectory with one entry."""
+    trajectory = Trajectory(
+        components=(EnvParamsSample,),
+        identity=(EnvParamsSample,),
+    )
+    trajectory.append(step=1, action=action, reward=0.5, observation=[100.0], env_params=dict(env_params))
+    env.trajectory = trajectory
 
 
 def test_cache_miss_when_env_params_differ(base_tr: TestRun, tmp_path: Path) -> None:
@@ -506,7 +526,7 @@ def test_cache_hit_when_neither_has_env_params(base_tr: TestRun, tmp_path: Path)
 
     env = CloudAIGymEnv(test_run=base_tr, runner=runner, rewards=RewardOverrides())
     env.test_run.current_iteration = 0
-    env.trajectory = {0: [TrajectoryEntry(step=1, action={"x": 10}, reward=0.5, observation=[100.0])]}
+    env.trajectory = Trajectory([_trajectory_entry(1, {"x": 10}, 0.5, [100.0])])
     # Note: neither the cached entry nor the trial carries env_params -> existing behavior.
 
     result = env.get_cached_trajectory_result({"x": 10}, {})
@@ -519,7 +539,7 @@ def test_step_reruns_workload_when_env_params_change(tmp_path: Path) -> None:
 
     Counterpart to test_cache_miss_when_env_params_differ but exercising the
     full step() flow: increment_step -> sample env_params -> apply_params_set ->
-    cache lookup -> runner.run() -> write_trajectory. With seed 42 the sampler
+    cache lookup -> runner.run() -> trajectory append. With seed 42 the sampler
     draws ball_speed=3 then ball_speed=1 on the two consecutive trials, so the
     cache key differs and the workload must re-run both times.
     """
@@ -568,13 +588,8 @@ def test_step_reruns_workload_when_env_params_change(tmp_path: Path) -> None:
     )
 
 
-def test_env_csv_is_step_aligned_with_trajectory(tmp_path: Path) -> None:
-    """env.csv must have exactly one row per env.step() call, with steps aligned 1:1 to trajectory.csv.
-
-    This pins the corpus-friendly contract: a downstream consumer can
-    ``pd.merge(traj, env, on="step")`` without losing rows on either side,
-    independent of whether the trial hit the trajectory cache.
-    """
+def test_env_params_are_recorded_in_trajectory_output(tmp_path: Path) -> None:
+    """Every recorded trial includes its sampled environment in the trajectory output."""
     tdef = EnvVarTestDefinition(
         name="dr",
         description="dr",
@@ -610,28 +625,13 @@ def test_env_csv_is_step_aligned_with_trajectory(tmp_path: Path) -> None:
         for action in (action_a, action_b, action_a):
             env.step(action)
 
-    env_csv = env.env_params_record_path
-    traj_csv = env.trajectory_file_path
-    assert env_csv.exists(), "env.csv must be written when env_params is declared"
-
-    env_steps = [int(line.split(",", 1)[0]) for line in env_csv.read_text().strip().splitlines()[1:]]
-    traj_steps = [int(line.split(",", 1)[0]) for line in traj_csv.read_text().strip().splitlines()[1:]]
-    assert env_steps == traj_steps == [1, 2, 3], (
-        f"step columns must align 1:1 across env.csv ({env_steps}) and trajectory.csv ({traj_steps})"
-    )
-
+    records = [json.loads(line) for line in env.trajectory_file_path.read_text().splitlines()]
+    assert [record["step"] for record in records] == [1, 2, 3]
+    assert all("ball_speed" in record["env_params"] for record in records)
 
-def test_env_csv_step_alignment_holds_on_constraint_failure(tmp_path: Path) -> None:
-    """A constraint failure must not desync env.csv from trajectory.csv.
 
-    Runs three steps where the middle one fails ``constraint_check`` and the
-    other two succeed. ``env.csv`` is sunk inside ``write_trajectory`` from the
-    same ``TrajectoryEntry``, which is never reached on the early-return
-    constraint-failure path - so the failed step lands in neither file. The
-    corpus-friendly contract (``pd.merge(traj, env, on="step")`` loses no rows)
-    therefore holds via shared absence: both files record exactly the surviving
-    steps, aligned 1:1.
-    """
+def test_constraint_failure_omits_the_complete_trajectory_row(tmp_path: Path) -> None:
+    """A constraint failure records no trajectory components."""
     tdef = EnvVarTestDefinition(
         name="dr",
         description="dr",
@@ -671,31 +671,20 @@ def test_env_csv_step_alignment_holds_on_constraint_failure(tmp_path: Path) -> N
         for action in ({"paddle_width": 4}, {"paddle_width": 6}, {"paddle_width": 8}):
             env.step(action)
 
-    env_csv = env.env_params_record_path
-    traj_csv = env.trajectory_file_path
-
-    assert env_csv.exists(), "surviving steps declare env_params -> env.csv must exist"
-    env_steps = [int(line.split(",", 1)[0]) for line in env_csv.read_text().strip().splitlines()[1:]]
+    trajectory_path = env.trajectory_file_path
     traj_steps = (
-        [int(line.split(",", 1)[0]) for line in traj_csv.read_text().strip().splitlines()[1:]]
-        if traj_csv.exists()
+        [json.loads(line)["step"] for line in trajectory_path.read_text().splitlines()]
+        if trajectory_path.exists()
         else []
     )
-    assert env_steps == traj_steps == [1, 3], (
-        f"the constraint-failed step (2) must appear in neither file; env.csv ({env_steps}) "
-        f"and trajectory.csv ({traj_steps}) must stay 1:1 aligned on the surviving steps"
-    )
+    assert traj_steps == [1, 3]
 
 
-def test_step_cache_hit_with_declared_env_params_still_writes_env_csv(tmp_path: Path) -> None:
-    """End-to-end: cache HIT under observer-driven env_params still records env.csv.
+def test_step_cache_hit_with_declared_env_params_records_complete_trajectory_row(tmp_path: Path) -> None:
+    """End-to-end: a cache hit records its environment in the trajectory output.
 
-    A cache hit still calls ``write_trajectory``, which sinks the trajectory row
-    and the matching env.csv row from the same entry - keeping the two files
-    step-aligned even when the workload itself is short-circuited.
-    Asserts: (a) the workload is NOT re-run (cache short-circuit), (b)
-    env.csv gains a row, (c) trajectory.csv gains a row carrying the
-    sampled env_params.
+    A cache hit still appends a complete row even though workload execution is
+    short-circuited.
     """
     import random as _random
 
@@ -726,13 +715,17 @@ def test_step_cache_hit_with_declared_env_params_still_writes_env_csv(tmp_path:
     env = CloudAIGymEnv(test_run=test_run, runner=runner, rewards=RewardOverrides())
     assert env.params is not None, "TestDefinition.env_params declared -> EnvParams must be built"
 
-    expected_sample = {"ball_speed": _random.Random("42:ball_speed:1").choice([1, 2, 3])}
+    expected_sample = {"ball_speed": _random.Random("42:ball_speed:2").choice([1, 2, 3])}
     action = {"paddle_width": 4}
     env.test_run.current_iteration = 0
-    env.trajectory = {
-        0: [TrajectoryEntry(step=0, action=action, reward=0.42, observation=[0.84], env_params=expected_sample)]
-    }
-    env.test_run.step = 0
+    env.trajectory.append(
+        step=1,
+        action=action,
+        reward=0.42,
+        observation=[0.84],
+        env_params=expected_sample,
+    )
+    env.test_run.step = 1
 
     with patch.object(env, "get_observation", side_effect=AssertionError("cache miss path must not run")):
         obs, reward, _done, info = env.step(action)
@@ -744,14 +737,13 @@ def test_step_cache_hit_with_declared_env_params_still_writes_env_csv(tmp_path:
         "the per-trial regime behind this observation is reported on info['env_params']"
     )
 
-    env_csv = env.env_params_record_path
-    assert env_csv.exists(), "cache HIT must NOT skip the observer; env.csv must record the trial"
-    env_rows = env_csv.read_text().strip().splitlines()
-    assert env_rows[0] == "step,env"
-    assert env_rows[1].startswith("1,"), f"expected step 1 row in env.csv, got {env_rows[1]!r}"
+    trajectory_records = [json.loads(line) for line in env.trajectory_file_path.read_text().splitlines()]
+    assert trajectory_records[-1]["step"] == 2
+    assert "ball_speed" in trajectory_records[-1]["env_params"]
 
-    traj_rows = env.trajectory[0]
-    assert len(traj_rows) == 2 and traj_rows[-1].env_params == expected_sample, (
+    traj_rows = env.trajectory
+    recorded_sample = traj_rows[-1].get(EnvParamsSample)
+    assert len(traj_rows) == 2 and recorded_sample is not None and recorded_sample.env_params == expected_sample, (
         "cache-hit trajectory entry must record the per-trial env_params sample"
     )
 
@@ -826,8 +818,10 @@ def test_param_space_excludes_env_params_keys(setup_env: tuple[TestRun, BaseRunn
     )
 
 
-def test_no_env_csv_when_env_params_not_declared(nemorun: NeMoRunTestDefinition, tmp_path: Path) -> None:
-    """Workloads without [env_params.*] pay zero overhead: no observer, no env.csv."""
+def test_csv_trajectory_has_no_env_params_column_when_not_declared(
+    nemorun: NeMoRunTestDefinition, tmp_path: Path
+) -> None:
+    """Workloads without env_params retain the base trajectory schema."""
     tdef = nemorun.model_copy(deep=True)
     tdef.cmd_args.data.global_batch_size = 8
     test_run = TestRun(
@@ -842,10 +836,16 @@ def test_no_env_csv_when_env_params_not_declared(nemorun: NeMoRunTestDefinition,
     runner.scenario_root = tmp_path / "scenario"
     runner.system = MagicMock()
 
-    env = CloudAIGymEnv(test_run=test_run, runner=runner, rewards=RewardOverrides())
+    env = CloudAIGymEnv(
+        test_run=test_run,
+        runner=runner,
+        rewards=RewardOverrides(),
+        trajectory_file_type="csv",
+    )
 
     assert env.params is None, "no env_params declared -> no EnvParams object"
-    assert not env.env_params_record_path.exists()
+    env.trajectory.append(step=1, action={}, reward=1.0, observation=[1.0])
+    assert env.trajectory_file_path.read_text().splitlines()[0] == "step,action,reward,observation"
 
 
 def _dr_env(tmp_path: Path, candidates: list, *, seed: int = 42) -> CloudAIGymEnv:
diff --git a/tests/test_env_params.py b/tests/test_env_params.py
index 9eb11d2ac..14acb54ba 100644
--- a/tests/test_env_params.py
+++ b/tests/test_env_params.py
@@ -33,7 +33,6 @@
 
 import dataclasses
 import random
-from pathlib import Path
 from typing import Any, List, Union
 
 import pytest
@@ -45,7 +44,6 @@
     EnvParams,
     EnvParamSpec,
     ObsLeafDescriptor,
-    write_env_params,
 )
 from cloudai.core import TestRun
 from cloudai.models.workload import CmdArgs, TestDefinition
@@ -204,27 +202,6 @@ def test_env_params_is_immutable() -> None:
         env_params.seed = 1  # pyright: ignore[reportAttributeAccessIssue]
 
 
-# --- write_env_params: unchanged persistence contract ---
-
-
-def test_csv_sink_skips_empty_samples_and_rejects_zero_step(tmp_path: Path) -> None:
-    path = tmp_path / "env.csv"
-    write_env_params(path, 1, {})  # empty -> no-op, no file
-    assert not path.exists()
-    with pytest.raises(ValueError, match="must be a positive trial index"):
-        write_env_params(path, 0, {"ball_speed": 1})
-
-
-def test_csv_sink_writes_header_then_rows(tmp_path: Path) -> None:
-    path = tmp_path / "env.csv"
-    write_env_params(path, 1, {"ball_speed": 2})
-    write_env_params(path, 2, {"ball_speed": 3})
-    contents = path.read_text().strip().splitlines()
-    assert contents[0] == "step,env"
-    assert contents[1].startswith("1,")
-    assert contents[2].startswith("2,")
-
-
 # --- EnvParams.from_test: resolves candidate lists from cmd_args, once at env formulation ---
 
 
diff --git a/tests/test_gymnasium_adapter_contract.py b/tests/test_gymnasium_adapter_contract.py
index a995fede0..54cef9630 100644
--- a/tests/test_gymnasium_adapter_contract.py
+++ b/tests/test_gymnasium_adapter_contract.py
@@ -25,7 +25,7 @@
 contextual-bandit configs (``agent_steps=1``), RLlib calls ``reset()`` before
 *every* trial. An earlier adapter rewound ``test_run.step`` on reset and
 collapsed every trial onto step 1 — silently overwriting output dirs and
-producing duplicate-step rows in trajectory.csv / env.csv.
+producing duplicate-step trajectory records.
 
 These tests pin the negative invariant: the adapter must not mutate
 ``test_run.step``. That counter is owned by ``TestRun`` and advanced
diff --git a/tests/test_handlers.py b/tests/test_handlers.py
index bd2a76c46..7e3ac478a 100644
--- a/tests/test_handlers.py
+++ b/tests/test_handlers.py
@@ -209,16 +209,16 @@ def _job_output_path(tr: TestRun, create: bool = True):
     assert (trajectory_dir / "3").exists()
     assert caplog.text.count("Retrieved cached result from") == 1
 
-    actual_trajectory = pd.read_csv(trajectory_dir / "trajectory.csv")
+    actual_trajectory = pd.read_json(trajectory_dir / "trajectory.jsonl", lines=True)
     expected_trajectory = pd.DataFrame(
         data=[
-            [1, "{'candidate': 1}", -1.0, "[-1.0]"],
-            [2, "{'candidate': 1}", -1.0, "[-1.0]"],
-            [3, "{'candidate': 2}", -1.0, "[-1.0]"],
+            [1, {"candidate": 1}, -1.0, [-1.0]],
+            [2, {"candidate": 1}, -1.0, [-1.0]],
+            [3, {"candidate": 2}, -1.0, [-1.0]],
         ],
         columns=["step", "action", "reward", "observation"],
     )
-    pd.testing.assert_frame_equal(actual_trajectory, expected_trajectory)
+    pd.testing.assert_frame_equal(actual_trajectory, expected_trajectory, check_dtype=False)
 
     assert [tr.step for tr in reporter.trs] == [1, 3]
 
diff --git a/tests/test_reporter.py b/tests/test_reporter.py
index 95acd8ac9..b8e593d28 100644
--- a/tests/test_reporter.py
+++ b/tests/test_reporter.py
@@ -16,6 +16,7 @@
 
 import copy
 import csv
+import json
 import tarfile
 from dataclasses import asdict
 from pathlib import Path
@@ -28,7 +29,7 @@
 from cloudai.cli.handlers import generate_reports
 from cloudai.core import CommandGenStrategy, Registry, Reporter, System
 from cloudai.models.scenario import ReportConfig, TestRunDetails
-from cloudai.report_generator.dse_report import build_dse_summaries
+from cloudai.report_generator.dse_report import build_dse_summaries, load_trajectory_dataframe
 from cloudai.reporter import DSEReporter, PerTestReporter, ReportItem, StatusReporter, TarballReporter
 from cloudai.systems.slurm.slurm_metadata import (
     MetadataCUDA,
@@ -46,6 +47,18 @@
 from cloudai.workloads.nccl_test import NCCLCmdArgs, NCCLTestDefinition
 
 
+def test_load_trajectory_dataframe_prefers_json_lines(tmp_path: Path) -> None:
+    jsonl_path = tmp_path / "trajectory.jsonl"
+    jsonl_path.write_text(json.dumps({"step": 1, "action": {"x": 2}, "reward": 3.0, "observation": [4.0]}) + "\n")
+
+    loaded = load_trajectory_dataframe(tmp_path)
+
+    assert loaded is not None
+    path, dataframe = loaded
+    assert path == jsonl_path
+    assert dataframe.to_dict(orient="records") == [{"step": 1, "action": {"x": 2}, "reward": 3, "observation": [4.0]}]
+
+
 class TestLoadTestTuns:
     def test_load_test_runs_behcnmark_sorted(self, slurm_system: SlurmSystem, benchmark_tr: TestRun) -> None:
         reporter = PerTestReporter(
diff --git a/tests/test_single_sbatch_runner.py b/tests/test_single_sbatch_runner.py
index 91d3cdf27..263891552 100644
--- a/tests/test_single_sbatch_runner.py
+++ b/tests/test_single_sbatch_runner.py
@@ -561,11 +561,11 @@ def test_trajectory_saved(dse_tr: TestRun, slurm_system: SlurmSystem) -> None:
     dse_tr.output_path = slurm_system.output_path / dse_tr.name
     dse_tr.output_path.mkdir(parents=True, exist_ok=True)
 
-    trajectory_path = runner.scenario_root / dse_tr.name / f"{dse_tr.current_iteration}" / "trajectory.csv"
+    trajectory_path = runner.scenario_root / dse_tr.name / f"{dse_tr.current_iteration}" / "trajectory.jsonl"
     trajectory_path.unlink(missing_ok=True)
     runner.handle_dse()
 
     assert trajectory_path.exists()
-    df = pd.read_csv(trajectory_path)
+    df = pd.read_json(trajectory_path, lines=True)
     assert df.shape[0] == len(dse_tr.all_combinations)
     assert df["step"].tolist() == list(range(1, len(dse_tr.all_combinations) + 1))

From efb4a97981503f9da3358c391b293823e82f3d54 Mon Sep 17 00:00:00 2001
From: Alex Manley <amanley@nvidia.com>
Date: Tue, 14 Jul 2026 14:51:07 -0700
Subject: [PATCH 04/14] [Trajectory] minor tweaks to interface

---
 src/cloudai/configurator/cloudai_gym.py | 20 ++----------
 src/cloudai/configurator/trajectory.py  | 40 +++++++++++++++--------
 tests/test_cloudaigym.py                |  1 -
 tests/test_trajectory.py                | 43 ++++++++-----------------
 4 files changed, 42 insertions(+), 62 deletions(-)

diff --git a/src/cloudai/configurator/cloudai_gym.py b/src/cloudai/configurator/cloudai_gym.py
index 848824e26..07330c7dd 100644
--- a/src/cloudai/configurator/cloudai_gym.py
+++ b/src/cloudai/configurator/cloudai_gym.py
@@ -26,12 +26,9 @@
 from .base_gym import BaseGym
 from .env_params import EnvParams, ObsLeafDescriptor
 from .trajectory import (
-    CsvTrajectoryWriter,
     EnvParamsSample,
-    JsonLinesTrajectoryWriter,
     Trajectory,
     TrajectoryEntry,
-    TrajectoryWriter,
     TrialResult,
 )
 
@@ -273,21 +270,10 @@ def encode_env_params(self, env_params: dict[str, Any]) -> Dict[str, Any]:
         return self.params.encode(env_params) if self.params is not None else {}
 
     def _new_trajectory(self, file_type: Literal["csv", "jsonl"]) -> Trajectory:
-        writer: TrajectoryWriter
-        if file_type == "csv":
-            writer = CsvTrajectoryWriter(lambda: self.iteration_dir)
-        elif file_type == "jsonl":
-            writer = JsonLinesTrajectoryWriter(lambda: self.iteration_dir)
-        else:
-            raise ValueError(f"Invalid file type: {file_type}")
-
-        if self.params is None:
-            return Trajectory(writer=writer)
-
         return Trajectory(
-            writer=writer,
-            components=(EnvParamsSample,),
-            identity=(EnvParamsSample,),
+            iteration_dir=lambda: self.iteration_dir,
+            file_type=file_type,
+            components=(EnvParamsSample,) if self.params is not None else (),
         )
 
     @property
diff --git a/src/cloudai/configurator/trajectory.py b/src/cloudai/configurator/trajectory.py
index 194f9bdd8..868139baf 100644
--- a/src/cloudai/configurator/trajectory.py
+++ b/src/cloudai/configurator/trajectory.py
@@ -24,7 +24,7 @@
 import logging
 from collections.abc import Callable, Iterator, Mapping, Sequence
 from pathlib import Path
-from typing import Any, Protocol, TypeVar, cast, overload
+from typing import Any, ClassVar, Literal, Protocol, TypeVar, cast, overload
 
 ComponentT = TypeVar("ComponentT")
 
@@ -33,6 +33,8 @@
 class EnvParamsSample:
     """Environment-parameter values sampled for one trial."""
 
+    contributes_to_identity: ClassVar[bool] = True
+
     env_params: dict[str, Any]
 
 
@@ -137,8 +139,8 @@ class Trajectory(Sequence[TrajectoryEntry]):
 
     ``components`` declares optional data types every entry must contain;
     :class:`TrialResult` is always included.
-    ``identity`` is the subset that affects trial equivalence;
-    informational components such as logs and metrics can be excluded.
+    Components whose ``contributes_to_identity`` class property is true affect
+    trial equivalence; other components are stored as informational data.
 
     Steps must be appended in increasing order, but gaps are permitted because
     CloudAI does not record constraint-failed trials.
@@ -148,24 +150,21 @@ def __init__(
         self,
         entries: Sequence[TrajectoryEntry] = (),
         *,
-        writer: TrajectoryWriter | None = None,
+        iteration_dir: Path | Callable[[], Path] | None = None,
+        file_type: Literal["csv", "jsonl"] = "jsonl",
         components: Sequence[type[object]] = (),
-        identity: Sequence[type[object]] = (),
     ) -> None:
         self._component_types = (TrialResult, *components)
         self._components = frozenset(self._component_types)
-        self._identity = frozenset(identity)
+        self._identity = frozenset(
+            component_type
+            for component_type in self._component_types
+            if getattr(component_type, "contributes_to_identity", False)
+        )
         if len(self._components) != len(self._component_types):
             raise ValueError("components cannot contain duplicate types")
-        if len(self._identity) != len(identity):
-            raise ValueError("identity cannot contain duplicate types")
 
-        undeclared_identity_types = self._identity - self._components
-        if undeclared_identity_types:
-            names = ", ".join(sorted(component_type.__name__ for component_type in undeclared_identity_types))
-            raise ValueError(f"identity types must be declared in components: {names}")
-
-        self._writer = writer
+        self._writer = self._create_writer(iteration_dir, file_type)
         self._fields_by_type = self._build_fields_by_type()
 
         self._entries: list[TrajectoryEntry] = []
@@ -179,6 +178,19 @@ def __init__(
             len(self),
         )
 
+    @staticmethod
+    def _create_writer(
+        iteration_dir: Path | Callable[[], Path] | None,
+        file_type: Literal["csv", "jsonl"],
+    ) -> TrajectoryWriter | None:
+        if file_type == "csv":
+            writer_type = CsvTrajectoryWriter
+        elif file_type == "jsonl":
+            writer_type = JsonLinesTrajectoryWriter
+        else:
+            raise ValueError(f"Invalid trajectory file type: {file_type}")
+        return writer_type(iteration_dir) if iteration_dir is not None else None
+
     @overload
     def __getitem__(self, index: int) -> TrajectoryEntry: ...
 
diff --git a/tests/test_cloudaigym.py b/tests/test_cloudaigym.py
index 6b2660ba0..492746226 100644
--- a/tests/test_cloudaigym.py
+++ b/tests/test_cloudaigym.py
@@ -468,7 +468,6 @@ def _seed_cached_entry_with_env_params(
     """Seed an environment-parameter-aware trajectory with one entry."""
     trajectory = Trajectory(
         components=(EnvParamsSample,),
-        identity=(EnvParamsSample,),
     )
     trajectory.append(step=1, action=action, reward=0.5, observation=[100.0], env_params=dict(env_params))
     env.trajectory = trajectory
diff --git a/tests/test_trajectory.py b/tests/test_trajectory.py
index e78e59747..249142fab 100644
--- a/tests/test_trajectory.py
+++ b/tests/test_trajectory.py
@@ -11,7 +11,6 @@
 import pytest
 
 from cloudai.configurator.trajectory import (
-    CsvTrajectoryWriter,
     EnvParamsSample,
     JsonLinesTrajectoryWriter,
     Trajectory,
@@ -25,18 +24,6 @@ class LoggingMetrics:
     logging_metrics: Mapping[str, float]
 
 
-class RecordingWriter:
-    def __init__(self, output_path: Path, *, fail: bool = False) -> None:
-        self.output_path = output_path
-        self.fail = fail
-        self.records: list[Mapping[str, object]] = []
-
-    def append(self, record: Mapping[str, object]) -> None:
-        if self.fail:
-            raise OSError("write failed")
-        self.records.append(record)
-
-
 def _entry(step: int, action: Mapping[str, Any] | None = None) -> TrajectoryEntry:
     return TrajectoryEntry(
         step=step,
@@ -71,8 +58,12 @@ def test_trajectory_rejects_non_increasing_steps() -> None:
         trajectory.append(step=2, action={"x": 2}, reward=2.0, observation=[2])
 
 
-def test_writer_failure_does_not_append_entry_to_memory(tmp_path: Path) -> None:
-    trajectory = Trajectory(writer=RecordingWriter(tmp_path / "trajectory", fail=True))
+def test_writer_failure_does_not_append_entry_to_memory(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    def fail_write(_writer: JsonLinesTrajectoryWriter, _record: Mapping[str, object]) -> None:
+        raise OSError("write failed")
+
+    monkeypatch.setattr(JsonLinesTrajectoryWriter, "append", fail_write)
+    trajectory = Trajectory(iteration_dir=tmp_path)
 
     with pytest.raises(OSError, match="write failed"):
         trajectory.append(step=1, action={"x": 1}, reward=1.0, observation=[1])
@@ -81,18 +72,17 @@ def test_writer_failure_does_not_append_entry_to_memory(tmp_path: Path) -> None:
 
 
 def test_initial_entries_are_not_replayed_to_writer(tmp_path: Path) -> None:
-    writer = RecordingWriter(tmp_path / "trajectory")
-
-    trajectory = Trajectory([_entry(1)], writer=writer)
+    trajectory = Trajectory([_entry(1)], iteration_dir=tmp_path)
 
     assert len(trajectory) == 1
-    assert writer.records == []
+    assert not (tmp_path / "trajectory.jsonl").exists()
 
 
 def test_append_writes_component_values_to_csv(tmp_path: Path) -> None:
     path = tmp_path / "trajectory.csv"
     trajectory = Trajectory(
-        writer=CsvTrajectoryWriter(tmp_path),
+        iteration_dir=tmp_path,
+        file_type="csv",
         components=(LoggingMetrics,),
     )
 
@@ -109,7 +99,8 @@ def test_append_writes_component_values_to_csv(tmp_path: Path) -> None:
 def test_append_writes_generic_records_as_json_lines(tmp_path: Path) -> None:
     path = tmp_path / "trajectory.jsonl"
     trajectory = Trajectory(
-        writer=JsonLinesTrajectoryWriter(tmp_path),
+        iteration_dir=tmp_path,
+        file_type="jsonl",
         components=(EnvParamsSample, LoggingMetrics),
     )
 
@@ -168,15 +159,9 @@ def test_trajectory_validates_its_fixed_component_schema() -> None:
         )
 
 
-def test_identity_component_types_must_be_declared() -> None:
-    with pytest.raises(ValueError, match=r"must be declared.*EnvParamsSample"):
-        Trajectory(identity=(EnvParamsSample,))
-
-
-def test_find_uses_only_configured_identity_components() -> None:
+def test_find_uses_only_components_that_contribute_to_identity() -> None:
     trajectory = Trajectory(
         components=(EnvParamsSample, LoggingMetrics),
-        identity=(EnvParamsSample,),
     )
     first = trajectory.append(
         step=1,
@@ -207,7 +192,6 @@ def test_find_ignores_informational_component_values() -> None:
 def test_find_requires_the_configured_identity_component_types() -> None:
     trajectory = Trajectory(
         components=(EnvParamsSample,),
-        identity=(EnvParamsSample,),
     )
 
     with pytest.raises(TypeError, match="missing: env_params"):
@@ -220,7 +204,6 @@ def test_find_requires_the_configured_identity_component_types() -> None:
 def test_find_preserves_exact_value_types_inside_components() -> None:
     trajectory = Trajectory(
         components=(EnvParamsSample,),
-        identity=(EnvParamsSample,),
     )
     trajectory.append(step=1, action={"x": 1}, reward=1.0, observation=[1], env_params={"speed": 1.0})
 

From ef1189a0e9c5310be5eabcc8adf20e6cdd93162b Mon Sep 17 00:00:00 2001
From: Alex Manley <amanley@nvidia.com>
Date: Tue, 14 Jul 2026 15:02:26 -0700
Subject: [PATCH 05/14] [Trajectory] revert default to csv

---
 src/cloudai/configurator/cloudai_gym.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/cloudai/configurator/cloudai_gym.py b/src/cloudai/configurator/cloudai_gym.py
index 07330c7dd..d68416ab5 100644
--- a/src/cloudai/configurator/cloudai_gym.py
+++ b/src/cloudai/configurator/cloudai_gym.py
@@ -46,7 +46,7 @@ def __init__(
         runner: BaseRunner,
         rewards: RewardOverrides,
         *,
-        trajectory_file_type: Literal["csv", "jsonl"] = "jsonl",
+        trajectory_file_type: Literal["csv", "jsonl"] = "csv",
     ):
         """
         Initialize the Gym environment using the TestRun object.

From 55264043eb88481266aa35683e8076f96f6baf47 Mon Sep 17 00:00:00 2001
From: Alex Manley <amanley@nvidia.com>
Date: Tue, 14 Jul 2026 15:07:28 -0700
Subject: [PATCH 06/14] [Trajectory] fix tests broken by jsonl

---
 src/cloudai/configurator/trajectory.py |  2 +-
 tests/test_cloudaigym.py               | 25 +++++++++++++------------
 tests/test_handlers.py                 | 10 +++++-----
 tests/test_single_sbatch_runner.py     |  4 ++--
 tests/test_trajectory.py               |  8 ++++----
 5 files changed, 25 insertions(+), 24 deletions(-)

diff --git a/src/cloudai/configurator/trajectory.py b/src/cloudai/configurator/trajectory.py
index 868139baf..37cf6320d 100644
--- a/src/cloudai/configurator/trajectory.py
+++ b/src/cloudai/configurator/trajectory.py
@@ -151,7 +151,7 @@ def __init__(
         entries: Sequence[TrajectoryEntry] = (),
         *,
         iteration_dir: Path | Callable[[], Path] | None = None,
-        file_type: Literal["csv", "jsonl"] = "jsonl",
+        file_type: Literal["csv", "jsonl"] = "csv",
         components: Sequence[type[object]] = (),
     ) -> None:
         self._component_types = (TrialResult, *components)
diff --git a/tests/test_cloudaigym.py b/tests/test_cloudaigym.py
index 492746226..ca5683a84 100644
--- a/tests/test_cloudaigym.py
+++ b/tests/test_cloudaigym.py
@@ -14,7 +14,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import json
 from pathlib import Path
 from typing import cast
 from unittest.mock import MagicMock, PropertyMock, patch
@@ -456,10 +455,10 @@ def test_cached_step_appends_trajectory_record(nemorun: NeMoRunTestDefinition, t
 
     trajectory_path = env.trajectory_file_path
     assert trajectory_path.exists()
-    assert trajectory_path.name == "trajectory.jsonl"
-    records = [json.loads(line) for line in trajectory_path.read_text().splitlines()]
-    assert records[-1]["step"] == 5
-    assert records[-1]["action"] == cached_action
+    assert trajectory_path.name == "trajectory.csv"
+    contents = trajectory_path.read_text().strip().splitlines()
+    assert contents[0] == "step,action,reward,observation"
+    assert contents[-1].startswith("5,")
 
 
 def _seed_cached_entry_with_env_params(
@@ -624,9 +623,10 @@ def test_env_params_are_recorded_in_trajectory_output(tmp_path: Path) -> None:
         for action in (action_a, action_b, action_a):
             env.step(action)
 
-    records = [json.loads(line) for line in env.trajectory_file_path.read_text().splitlines()]
-    assert [record["step"] for record in records] == [1, 2, 3]
-    assert all("ball_speed" in record["env_params"] for record in records)
+    rows = env.trajectory_file_path.read_text().strip().splitlines()
+    assert rows[0] == "step,action,reward,observation,env_params"
+    assert [int(row.split(",", 1)[0]) for row in rows[1:]] == [1, 2, 3]
+    assert all("ball_speed" in row for row in rows[1:])
 
 
 def test_constraint_failure_omits_the_complete_trajectory_row(tmp_path: Path) -> None:
@@ -672,7 +672,7 @@ def test_constraint_failure_omits_the_complete_trajectory_row(tmp_path: Path) ->
 
     trajectory_path = env.trajectory_file_path
     traj_steps = (
-        [json.loads(line)["step"] for line in trajectory_path.read_text().splitlines()]
+        [int(line.split(",", 1)[0]) for line in trajectory_path.read_text().strip().splitlines()[1:]]
         if trajectory_path.exists()
         else []
     )
@@ -736,9 +736,10 @@ def test_step_cache_hit_with_declared_env_params_records_complete_trajectory_row
         "the per-trial regime behind this observation is reported on info['env_params']"
     )
 
-    trajectory_records = [json.loads(line) for line in env.trajectory_file_path.read_text().splitlines()]
-    assert trajectory_records[-1]["step"] == 2
-    assert "ball_speed" in trajectory_records[-1]["env_params"]
+    trajectory_rows = env.trajectory_file_path.read_text().strip().splitlines()
+    assert trajectory_rows[0] == "step,action,reward,observation,env_params"
+    assert trajectory_rows[-1].startswith("2,")
+    assert "ball_speed" in trajectory_rows[-1]
 
     traj_rows = env.trajectory
     recorded_sample = traj_rows[-1].get(EnvParamsSample)
diff --git a/tests/test_handlers.py b/tests/test_handlers.py
index 7e3ac478a..bd2a76c46 100644
--- a/tests/test_handlers.py
+++ b/tests/test_handlers.py
@@ -209,16 +209,16 @@ def _job_output_path(tr: TestRun, create: bool = True):
     assert (trajectory_dir / "3").exists()
     assert caplog.text.count("Retrieved cached result from") == 1
 
-    actual_trajectory = pd.read_json(trajectory_dir / "trajectory.jsonl", lines=True)
+    actual_trajectory = pd.read_csv(trajectory_dir / "trajectory.csv")
     expected_trajectory = pd.DataFrame(
         data=[
-            [1, {"candidate": 1}, -1.0, [-1.0]],
-            [2, {"candidate": 1}, -1.0, [-1.0]],
-            [3, {"candidate": 2}, -1.0, [-1.0]],
+            [1, "{'candidate': 1}", -1.0, "[-1.0]"],
+            [2, "{'candidate': 1}", -1.0, "[-1.0]"],
+            [3, "{'candidate': 2}", -1.0, "[-1.0]"],
         ],
         columns=["step", "action", "reward", "observation"],
     )
-    pd.testing.assert_frame_equal(actual_trajectory, expected_trajectory, check_dtype=False)
+    pd.testing.assert_frame_equal(actual_trajectory, expected_trajectory)
 
     assert [tr.step for tr in reporter.trs] == [1, 3]
 
diff --git a/tests/test_single_sbatch_runner.py b/tests/test_single_sbatch_runner.py
index 263891552..91d3cdf27 100644
--- a/tests/test_single_sbatch_runner.py
+++ b/tests/test_single_sbatch_runner.py
@@ -561,11 +561,11 @@ def test_trajectory_saved(dse_tr: TestRun, slurm_system: SlurmSystem) -> None:
     dse_tr.output_path = slurm_system.output_path / dse_tr.name
     dse_tr.output_path.mkdir(parents=True, exist_ok=True)
 
-    trajectory_path = runner.scenario_root / dse_tr.name / f"{dse_tr.current_iteration}" / "trajectory.jsonl"
+    trajectory_path = runner.scenario_root / dse_tr.name / f"{dse_tr.current_iteration}" / "trajectory.csv"
     trajectory_path.unlink(missing_ok=True)
     runner.handle_dse()
 
     assert trajectory_path.exists()
-    df = pd.read_json(trajectory_path, lines=True)
+    df = pd.read_csv(trajectory_path)
     assert df.shape[0] == len(dse_tr.all_combinations)
     assert df["step"].tolist() == list(range(1, len(dse_tr.all_combinations) + 1))
diff --git a/tests/test_trajectory.py b/tests/test_trajectory.py
index 249142fab..b030d2c90 100644
--- a/tests/test_trajectory.py
+++ b/tests/test_trajectory.py
@@ -11,8 +11,8 @@
 import pytest
 
 from cloudai.configurator.trajectory import (
+    CsvTrajectoryWriter,
     EnvParamsSample,
-    JsonLinesTrajectoryWriter,
     Trajectory,
     TrajectoryEntry,
     TrialResult,
@@ -59,10 +59,10 @@ def test_trajectory_rejects_non_increasing_steps() -> None:
 
 
 def test_writer_failure_does_not_append_entry_to_memory(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
-    def fail_write(_writer: JsonLinesTrajectoryWriter, _record: Mapping[str, object]) -> None:
+    def fail_write(_writer: CsvTrajectoryWriter, _record: Mapping[str, object]) -> None:
         raise OSError("write failed")
 
-    monkeypatch.setattr(JsonLinesTrajectoryWriter, "append", fail_write)
+    monkeypatch.setattr(CsvTrajectoryWriter, "append", fail_write)
     trajectory = Trajectory(iteration_dir=tmp_path)
 
     with pytest.raises(OSError, match="write failed"):
@@ -75,7 +75,7 @@ def test_initial_entries_are_not_replayed_to_writer(tmp_path: Path) -> None:
     trajectory = Trajectory([_entry(1)], iteration_dir=tmp_path)
 
     assert len(trajectory) == 1
-    assert not (tmp_path / "trajectory.jsonl").exists()
+    assert not (tmp_path / "trajectory.csv").exists()
 
 
 def test_append_writes_component_values_to_csv(tmp_path: Path) -> None:

From d469b00603b752a0adcb86ab9f5a1db894559eae Mon Sep 17 00:00:00 2001
From: Alex Manley <amanley@nvidia.com>
Date: Tue, 14 Jul 2026 15:10:50 -0700
Subject: [PATCH 07/14] [Trajectory] fix test copyright header

---
 tests/test_trajectory.py | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/tests/test_trajectory.py b/tests/test_trajectory.py
index b030d2c90..f46318b9d 100644
--- a/tests/test_trajectory.py
+++ b/tests/test_trajectory.py
@@ -1,6 +1,18 @@
 # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES
 # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 
 import dataclasses
 import json

From 7aefb8372b8d9efe6358ee09c8050f928927eecc Mon Sep 17 00:00:00 2001
From: Alex Manley <amanley@nvidia.com>
Date: Wed, 15 Jul 2026 10:05:19 -0700
Subject: [PATCH 08/14] [Trajectory] post-review improvements

---
 src/cloudai/configurator/trajectory.py        | 55 ++++++++++++++--
 .../systems/slurm/single_sbatch_runner.py     | 16 +++--
 tests/test_cloudaigym.py                      |  2 +-
 tests/test_reporter.py                        |  4 ++
 tests/test_single_sbatch_runner.py            | 24 +++++++
 tests/test_trajectory.py                      | 66 +++++++++++++++++++
 6 files changed, 155 insertions(+), 12 deletions(-)

diff --git a/src/cloudai/configurator/trajectory.py b/src/cloudai/configurator/trajectory.py
index 37cf6320d..f15f140d1 100644
--- a/src/cloudai/configurator/trajectory.py
+++ b/src/cloudai/configurator/trajectory.py
@@ -24,6 +24,7 @@
 import logging
 from collections.abc import Callable, Iterator, Mapping, Sequence
 from pathlib import Path
+from types import MappingProxyType
 from typing import Any, ClassVar, Literal, Protocol, TypeVar, cast, overload
 
 ComponentT = TypeVar("ComponentT")
@@ -35,7 +36,11 @@ class EnvParamsSample:
 
     contributes_to_identity: ClassVar[bool] = True
 
-    env_params: dict[str, Any]
+    env_params: Mapping[str, Any]
+
+    def __post_init__(self) -> None:
+        """Own an immutable snapshot of the sampled values."""
+        object.__setattr__(self, "env_params", _freeze(self.env_params))
 
 
 @dataclasses.dataclass(frozen=True)
@@ -46,6 +51,10 @@ class TrialResult:
     reward: float
     observation: Sequence[Any]
 
+    def __post_init__(self) -> None:
+        """Own an immutable snapshot of the action."""
+        object.__setattr__(self, "action", _freeze(self.action))
+
 
 @dataclasses.dataclass(frozen=True)
 class TrajectoryEntry:
@@ -103,17 +112,26 @@ def __init__(self, iteration_dir: Path | Callable[[], Path]) -> None:
 
     def append(self, record: Mapping[str, object]) -> None:
         fields = tuple(record)
+        path = self.output_path
+        path.parent.mkdir(parents=True, exist_ok=True)
+        write_header = not path.exists() or path.stat().st_size == 0
+        existing_fields: tuple[str, ...] = ()
+        if not write_header:
+            with path.open(newline="") as file:
+                existing_fields = tuple(next(csv.reader(file), ()))
+
         if self._fields is None:
+            if existing_fields and existing_fields != fields:
+                raise ValueError(f"trajectory file fields do not match: expected {fields}, got {existing_fields}")
             self._fields = fields
         elif fields != self._fields:
             raise ValueError(f"trajectory record fields changed: expected {self._fields}, got {fields}")
+        elif existing_fields and existing_fields != self._fields:
+            raise ValueError(f"trajectory file fields do not match: expected {self._fields}, got {existing_fields}")
 
-        path = self.output_path
-        new_file = not path.exists()
-        path.parent.mkdir(parents=True, exist_ok=True)
         with path.open("a", newline="") as file:
             writer = csv.DictWriter(file, fieldnames=self._fields)
-            if new_file:
+            if write_header:
                 writer.writeheader()
             writer.writerow(record)
         logging.debug("Wrote trajectory record to %s.", path)
@@ -237,11 +255,12 @@ def find(self, action: Mapping[str, Any], **identity_values: object) -> Trajecto
             "trajectory identity values",
         )
         identity = self._identity_for(identity_components)
+        frozen_action = _freeze(action)
         for entry in self._entries:
             result = entry.get(TrialResult)
             if result is None:
                 raise ValueError(f"trajectory entry at step {entry.step} is missing TrialResult")
-            if _values_match_exact(result.action, action) and _values_match_exact(
+            if _values_match_exact(result.action, frozen_action) and _values_match_exact(
                 self._identity_for(entry.components), identity
             ):
                 logging.debug("Found matching trajectory entry at step %s for action %s.", entry.step, action)
@@ -312,7 +331,7 @@ def _to_record(self, entry: TrajectoryEntry) -> dict[str, object]:
         for component_type in self._component_types:
             component = components_by_type[component_type]
             record.update(
-                {field.name: getattr(component, field.name) for field in self._fields_by_type[component_type]}
+                {_field.name: _thaw(getattr(component, _field.name)) for _field in self._fields_by_type[component_type]}
             )
         return record
 
@@ -345,6 +364,28 @@ def _validate_schema(
     raise TypeError(f"{context} do not match configured schema ({'; '.join(details)})")
 
 
+def _freeze(value: Any) -> Any:
+    """Recursively copy mutable containers into read-only equivalents."""
+    if isinstance(value, Mapping):
+        return MappingProxyType({key: _freeze(item) for key, item in value.items()})
+    if isinstance(value, (list, tuple)):
+        return tuple(_freeze(item) for item in value)
+    if isinstance(value, (set, frozenset)):
+        return frozenset(_freeze(item) for item in value)
+    return value
+
+
+def _thaw(value: Any) -> Any:
+    """Convert frozen containers to values supported by trajectory writers."""
+    if isinstance(value, Mapping):
+        return {key: _thaw(item) for key, item in value.items()}
+    if isinstance(value, tuple):
+        return [_thaw(item) for item in value]
+    if isinstance(value, frozenset):
+        return [_thaw(item) for item in value]
+    return value
+
+
 def _values_match_exact(left: Any, right: Any) -> bool:
     if type(left) is not type(right):
         return False
diff --git a/src/cloudai/systems/slurm/single_sbatch_runner.py b/src/cloudai/systems/slurm/single_sbatch_runner.py
index 9f73fb2af..7258d9f68 100644
--- a/src/cloudai/systems/slurm/single_sbatch_runner.py
+++ b/src/cloudai/systems/slurm/single_sbatch_runner.py
@@ -22,6 +22,7 @@
 from typing import Generator, Optional, cast
 
 from cloudai.configurator import CloudAIGymEnv
+from cloudai.configurator.env_params import EnvParams
 from cloudai.core import BaseJob, JobIdRetrievalError, Registry, System, TestRun, TestScenario
 from cloudai.util import CommandShell, format_time_limit, parse_time_limit
 
@@ -125,9 +126,11 @@ def get_single_tr_block(self, tr: TestRun) -> str:
         return srun_cmd
 
     def unroll_dse(self, tr: TestRun) -> Generator[TestRun, None, None]:
-        for idx, combination in enumerate(tr.all_combinations):
-            next_tr = tr.apply_params_set(combination)
-            next_tr.step = idx + 1
+        params = EnvParams.from_test(tr.test)
+        for idx, combination in enumerate(tr.all_combinations, start=1):
+            sampled_env_params = params.sample(idx) if params is not None else {}
+            next_tr = tr.apply_params_set(combination, env_params=sampled_env_params)
+            next_tr.step = idx
             next_tr.output_path = self.get_job_output_path(next_tr)
 
             if next_tr.test.constraint_check(next_tr, self.system):
@@ -211,18 +214,23 @@ def handle_dse(self):
             gym = CloudAIGymEnv(tr, self, rewards=agent_config.rewards)
 
             for idx, combination in enumerate(tr.all_combinations, start=1):
-                next_tr = tr.apply_params_set(combination)
+                sampled_env_params = gym.params.sample(idx) if gym.params is not None else {}
+                next_tr = tr.apply_params_set(combination, env_params=sampled_env_params)
                 next_tr.step = idx
                 next_tr.output_path = self.get_job_output_path(next_tr)
 
                 gym.test_run = next_tr
                 observation = gym.get_observation({})
                 reward = gym.compute_reward(observation)
+                trajectory_values: dict[str, object] = {}
+                if gym.params is not None:
+                    trajectory_values["env_params"] = sampled_env_params
                 gym.trajectory.append(
                     step=idx,
                     action=combination,
                     reward=reward,
                     observation=observation,
+                    **trajectory_values,
                 )
 
     def completed_test_runs(self, job: BaseJob) -> list[TestRun]:
diff --git a/tests/test_cloudaigym.py b/tests/test_cloudaigym.py
index ca5683a84..858019cd4 100644
--- a/tests/test_cloudaigym.py
+++ b/tests/test_cloudaigym.py
@@ -714,7 +714,7 @@ def test_step_cache_hit_with_declared_env_params_records_complete_trajectory_row
     env = CloudAIGymEnv(test_run=test_run, runner=runner, rewards=RewardOverrides())
     assert env.params is not None, "TestDefinition.env_params declared -> EnvParams must be built"
 
-    expected_sample = {"ball_speed": _random.Random("42:ball_speed:2").choice([1, 2, 3])}
+    expected_sample = {"ball_speed": _random.Random("42:ball_speed:2").choice([1, 2, 3])}  # noqa: S311, RUF100
     action = {"paddle_width": 4}
     env.test_run.current_iteration = 0
     env.trajectory.append(
diff --git a/tests/test_reporter.py b/tests/test_reporter.py
index b8e593d28..f6420be52 100644
--- a/tests/test_reporter.py
+++ b/tests/test_reporter.py
@@ -50,6 +50,10 @@
 def test_load_trajectory_dataframe_prefers_json_lines(tmp_path: Path) -> None:
     jsonl_path = tmp_path / "trajectory.jsonl"
     jsonl_path.write_text(json.dumps({"step": 1, "action": {"x": 2}, "reward": 3.0, "observation": [4.0]}) + "\n")
+    with (tmp_path / "trajectory.csv").open("w", newline="") as file:
+        writer = csv.DictWriter(file, fieldnames=("step", "action", "reward", "observation"))
+        writer.writeheader()
+        writer.writerow({"step": 99, "action": {"x": 99}, "reward": 99.0, "observation": [99.0]})
 
     loaded = load_trajectory_dataframe(tmp_path)
 
diff --git a/tests/test_single_sbatch_runner.py b/tests/test_single_sbatch_runner.py
index 91d3cdf27..09a63a841 100644
--- a/tests/test_single_sbatch_runner.py
+++ b/tests/test_single_sbatch_runner.py
@@ -14,6 +14,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import ast
 import copy
 import re
 from pathlib import Path
@@ -24,6 +25,7 @@
 import pytest
 import toml
 
+from cloudai.configurator.env_params import EnvParams, EnvParamSpec
 from cloudai.core import Registry, System, TestRun, TestScenario
 from cloudai.systems.slurm import SingleSbatchRunner, SlurmJob, SlurmJobMetadata, SlurmSystem
 from cloudai.workloads.nccl_test import NCCLCmdArgs, NCCLTestDefinition
@@ -569,3 +571,25 @@ def test_trajectory_saved(dse_tr: TestRun, slurm_system: SlurmSystem) -> None:
     df = pd.read_csv(trajectory_path)
     assert df.shape[0] == len(dse_tr.all_combinations)
     assert df["step"].tolist() == list(range(1, len(dse_tr.all_combinations) + 1))
+
+
+def test_dse_env_params_are_applied_to_runs_and_trajectory(dse_tr: TestRun, slurm_system: SlurmSystem) -> None:
+    dse_tr.test.cmd_args.nthreads = [1, 2]
+    dse_tr.test.env_params = {"nthreads": EnvParamSpec()}
+    dse_tr.test.agent_config = {"random_seed": 42}
+    tc = TestScenario(name="tc", test_runs=[dse_tr])
+    runner = SingleSbatchRunner(mode="run", system=slurm_system, test_scenario=tc, output_path=slurm_system.output_path)
+    trajectory_path = runner.scenario_root / dse_tr.name / f"{dse_tr.current_iteration}" / "trajectory.csv"
+    trajectory_path.unlink(missing_ok=True)
+
+    params = EnvParams.from_test(dse_tr.test)
+    assert params is not None
+    expected_samples = [params.sample(step) for step in range(1, len(dse_tr.all_combinations) + 1)]
+
+    unrolled = list(runner.unroll_dse(dse_tr))
+    assert [tr.test.cmd_args.nthreads for tr in unrolled] == [sample["nthreads"] for sample in expected_samples]
+
+    runner.handle_dse()
+
+    dataframe = pd.read_csv(trajectory_path)
+    assert dataframe["env_params"].map(ast.literal_eval).tolist() == expected_samples
diff --git a/tests/test_trajectory.py b/tests/test_trajectory.py
index f46318b9d..9c8aba417 100644
--- a/tests/test_trajectory.py
+++ b/tests/test_trajectory.py
@@ -108,6 +108,40 @@ def test_append_writes_component_values_to_csv(tmp_path: Path) -> None:
     ]
 
 
+def test_csv_writer_initializes_a_precreated_empty_file(tmp_path: Path) -> None:
+    path = tmp_path / "trajectory.csv"
+    path.touch()
+
+    Trajectory(iteration_dir=tmp_path).append(step=1, action={"x": 1}, reward=1.0, observation=[1])
+
+    assert path.read_text().splitlines() == [
+        "step,action,reward,observation",
+        "1,{'x': 1},1.0,[1]",
+    ]
+
+
+def test_csv_writer_reuses_a_matching_header_without_duplication(tmp_path: Path) -> None:
+    path = tmp_path / "trajectory.csv"
+    path.write_text("step,action,reward,observation\n")
+
+    Trajectory(iteration_dir=tmp_path).append(step=1, action={"x": 1}, reward=1.0, observation=[1])
+
+    assert path.read_text().splitlines() == [
+        "step,action,reward,observation",
+        "1,{'x': 1},1.0,[1]",
+    ]
+
+
+def test_csv_writer_rejects_an_existing_mismatched_header(tmp_path: Path) -> None:
+    path = tmp_path / "trajectory.csv"
+    path.write_text("step,reward,action,observation\n")
+
+    with pytest.raises(ValueError, match="trajectory file fields do not match"):
+        Trajectory(iteration_dir=tmp_path).append(step=1, action={"x": 1}, reward=1.0, observation=[1])
+
+    assert path.read_text() == "step,reward,action,observation\n"
+
+
 def test_append_writes_generic_records_as_json_lines(tmp_path: Path) -> None:
     path = tmp_path / "trajectory.jsonl"
     trajectory = Trajectory(
@@ -228,6 +262,38 @@ def test_find_preserves_exact_action_value_types() -> None:
     assert trajectory.find({"x": 1}) is None
 
 
+def test_identity_values_are_deeply_immutable_and_owned_by_the_entry(tmp_path: Path) -> None:
+    trajectory = Trajectory(iteration_dir=tmp_path, components=(EnvParamsSample,))
+    action = {"shape": {"layers": [1, 2]}}
+    sampled_env_params = {"regime": {"speeds": [3, 4]}}
+    entry = trajectory.append(
+        step=1,
+        action=action,
+        reward=1.0,
+        observation=[1],
+        env_params=sampled_env_params,
+    )
+    result = entry.get(TrialResult)
+    env_params = entry.get(EnvParamsSample)
+    assert result is not None and env_params is not None
+
+    action["shape"]["layers"].append(99)
+    sampled_env_params["regime"]["speeds"].append(99)
+    with pytest.raises(TypeError):
+        result.action["shape"] = {}  # type: ignore[index]
+    with pytest.raises(AttributeError):
+        result.action["shape"]["layers"].append(99)
+    with pytest.raises(TypeError):
+        env_params.env_params["regime"] = {}  # type: ignore[index]
+    with pytest.raises(AttributeError):
+        env_params.env_params["regime"]["speeds"].append(99)
+
+    original_action = {"shape": {"layers": [1, 2]}}
+    original_env_params = {"regime": {"speeds": [3, 4]}}
+    assert trajectory.find(original_action, env_params=original_env_params) is entry
+    assert "99" not in (tmp_path / "trajectory.csv").read_text()
+
+
 def test_trajectory_logs_lifecycle_and_lookup(caplog: pytest.LogCaptureFixture) -> None:
     with caplog.at_level("DEBUG"):
         trajectory = Trajectory()

From 994c688b2511c8ee75e37f0d6f5658b15a00190b Mon Sep 17 00:00:00 2001
From: Alex Manley <amanley@nvidia.com>
Date: Wed, 15 Jul 2026 10:39:56 -0700
Subject: [PATCH 09/14] [Trajectory] fix component freezing

---
 src/cloudai/configurator/trajectory.py | 25 +++++++++++---
 tests/test_trajectory.py               | 46 +++++++++++++++++++++++++-
 2 files changed, 66 insertions(+), 5 deletions(-)

diff --git a/src/cloudai/configurator/trajectory.py b/src/cloudai/configurator/trajectory.py
index f15f140d1..6b3d7b3df 100644
--- a/src/cloudai/configurator/trajectory.py
+++ b/src/cloudai/configurator/trajectory.py
@@ -18,6 +18,7 @@
 
 from __future__ import annotations
 
+import copy
 import csv
 import dataclasses
 import json
@@ -236,17 +237,19 @@ def append(self, *, step: int, **values: object) -> TrajectoryEntry:
         """Build configured components from values, then store and persist one entry."""
         components = self._construct_components(self._component_types, values, "trajectory values")
         entry = TrajectoryEntry(step=step, components=components)
-        self._store_entry(entry, persist=True)
-        return entry
+        return self._store_entry(entry, persist=True)
 
-    def _store_entry(self, entry: TrajectoryEntry, *, persist: bool) -> None:
+    def _store_entry(self, entry: TrajectoryEntry, *, persist: bool) -> TrajectoryEntry:
         self._validate_entry_components(entry)
         if self._entries and entry.step <= self._entries[-1].step:
             raise ValueError(f"trajectory steps must increase: last step is {self._entries[-1].step}, got {entry.step}")
+        if self._identity:
+            entry = TrajectoryEntry(step=entry.step, components=self._freeze_identity_components(entry.components))
         if persist and self._writer is not None:
             self._writer.append(self._to_record(entry))
         self._entries.append(entry)
         logging.debug("Appended trajectory entry for step %s (total entries: %s).", entry.step, len(self))
+        return entry
 
     def find(self, action: Mapping[str, Any], **identity_values: object) -> TrajectoryEntry | None:
         identity_components = self._construct_components(
@@ -254,7 +257,7 @@ def find(self, action: Mapping[str, Any], **identity_values: object) -> Trajecto
             identity_values,
             "trajectory identity values",
         )
-        identity = self._identity_for(identity_components)
+        identity = self._identity_for(self._freeze_identity_components(identity_components))
         frozen_action = _freeze(action)
         for entry in self._entries:
             result = entry.get(TrialResult)
@@ -278,6 +281,12 @@ def _validate_entry_components(self, entry: TrajectoryEntry) -> None:
     def _identity_for(self, components: Sequence[object]) -> dict[type[object], object]:
         return {type(component): component for component in components if type(component) in self._identity}
 
+    def _freeze_identity_components(self, components: Sequence[object]) -> tuple[object, ...]:
+        return tuple(
+            _freeze_component(component) if type(component) in self._identity else component
+            for component in components
+        )
+
     def _construct_components(
         self,
         component_types: Sequence[type[object]],
@@ -375,6 +384,14 @@ def _freeze(value: Any) -> Any:
     return value
 
 
+def _freeze_component(component: object) -> object:
+    """Copy a dataclass component and freeze all of its stored fields."""
+    snapshot = copy.copy(component)
+    for field in dataclasses.fields(cast(Any, component)):
+        object.__setattr__(snapshot, field.name, _freeze(getattr(component, field.name)))
+    return snapshot
+
+
 def _thaw(value: Any) -> Any:
     """Convert frozen containers to values supported by trajectory writers."""
     if isinstance(value, Mapping):
diff --git a/tests/test_trajectory.py b/tests/test_trajectory.py
index 9c8aba417..07672b67a 100644
--- a/tests/test_trajectory.py
+++ b/tests/test_trajectory.py
@@ -18,7 +18,7 @@
 import json
 from collections.abc import Mapping
 from pathlib import Path
-from typing import Any
+from typing import Any, ClassVar
 
 import pytest
 
@@ -36,6 +36,12 @@ class LoggingMetrics:
     logging_metrics: Mapping[str, float]
 
 
+@dataclasses.dataclass(frozen=True)
+class CacheContext:
+    contributes_to_identity: ClassVar[bool] = True
+    cache_context: Mapping[str, Any]
+
+
 def _entry(step: int, action: Mapping[str, Any] | None = None) -> TrajectoryEntry:
     return TrajectoryEntry(
         step=step,
@@ -222,6 +228,44 @@ def test_find_uses_only_components_that_contribute_to_identity() -> None:
     assert trajectory.find({"x": 1}, env_params={"speed": 2}) is None
 
 
+def test_future_identity_components_are_frozen_by_trajectory(tmp_path: Path) -> None:
+    context = {"hardware": {"gpus": [8]}}
+    trajectory = Trajectory(iteration_dir=tmp_path, components=(CacheContext,))
+
+    entry = trajectory.append(
+        step=1,
+        action={"x": 1},
+        reward=1.0,
+        observation=[1],
+        cache_context=context,
+    )
+    context["hardware"]["gpus"].append(16)
+    stored_context = entry.get(CacheContext)
+    assert stored_context is not None
+
+    with pytest.raises(TypeError):
+        stored_context.cache_context["hardware"] = {}  # type: ignore[index]
+    with pytest.raises(AttributeError):
+        stored_context.cache_context["hardware"]["gpus"].append(16)
+
+    assert trajectory.find({"x": 1}, cache_context={"hardware": {"gpus": [8]}}) is entry
+    assert trajectory.find({"x": 1}, cache_context=context) is None
+    assert "16" not in (tmp_path / "trajectory.csv").read_text()
+
+
+def test_initial_entries_snapshot_future_identity_components() -> None:
+    context = CacheContext({"hardware": {"gpus": [8]}})
+    entry = TrajectoryEntry(
+        step=1,
+        components=(TrialResult(action={"x": 1}, reward=1.0, observation=[1]), context),
+    )
+    trajectory = Trajectory([entry], components=(CacheContext,))
+
+    context.cache_context["hardware"]["gpus"].append(16)
+
+    assert trajectory.find({"x": 1}, cache_context={"hardware": {"gpus": [8]}}) is trajectory[0]
+
+
 def test_find_ignores_informational_component_values() -> None:
     trajectory = Trajectory(components=(LoggingMetrics,))
     first = trajectory.append(

From daebbb8e57dcab224ff40ef98f23008d42055117 Mon Sep 17 00:00:00 2001
From: Alex Manley <amanley@nvidia.com>
Date: Wed, 15 Jul 2026 10:48:06 -0700
Subject: [PATCH 10/14] [Trajectory] minor code cleanup and fixes

---
 src/cloudai/configurator/trajectory.py        | 32 ++++++++++---------
 .../systems/slurm/single_sbatch_runner.py     |  3 ++
 tests/test_trajectory.py                      |  2 +-
 3 files changed, 21 insertions(+), 16 deletions(-)

diff --git a/src/cloudai/configurator/trajectory.py b/src/cloudai/configurator/trajectory.py
index 6b3d7b3df..6914b3f68 100644
--- a/src/cloudai/configurator/trajectory.py
+++ b/src/cloudai/configurator/trajectory.py
@@ -39,10 +39,6 @@ class EnvParamsSample:
 
     env_params: Mapping[str, Any]
 
-    def __post_init__(self) -> None:
-        """Own an immutable snapshot of the sampled values."""
-        object.__setattr__(self, "env_params", _freeze(self.env_params))
-
 
 @dataclasses.dataclass(frozen=True)
 class TrialResult:
@@ -65,10 +61,20 @@ class TrajectoryEntry:
     components: tuple[object, ...]
 
     def __post_init__(self) -> None:
-        """Validate the trial index and component uniqueness."""
+        """Validate the entry and snapshot its identity components."""
         if self.step < 1:
             raise ValueError(f"trajectory step must be positive; got {self.step}")
-        _components_by_type(self.components)
+        _validate_components(self.components)
+        object.__setattr__(
+            self,
+            "components",
+            tuple(
+                _freeze_component(component)
+                if getattr(type(component), "contributes_to_identity", False)
+                else component
+                for component in self.components
+            ),
+        )
 
     def get(self, component_type: type[ComponentT]) -> ComponentT | None:
         """Return the component with exactly the requested type, if present."""
@@ -237,19 +243,17 @@ def append(self, *, step: int, **values: object) -> TrajectoryEntry:
         """Build configured components from values, then store and persist one entry."""
         components = self._construct_components(self._component_types, values, "trajectory values")
         entry = TrajectoryEntry(step=step, components=components)
-        return self._store_entry(entry, persist=True)
+        self._store_entry(entry, persist=True)
+        return entry
 
-    def _store_entry(self, entry: TrajectoryEntry, *, persist: bool) -> TrajectoryEntry:
+    def _store_entry(self, entry: TrajectoryEntry, *, persist: bool) -> None:
         self._validate_entry_components(entry)
         if self._entries and entry.step <= self._entries[-1].step:
             raise ValueError(f"trajectory steps must increase: last step is {self._entries[-1].step}, got {entry.step}")
-        if self._identity:
-            entry = TrajectoryEntry(step=entry.step, components=self._freeze_identity_components(entry.components))
         if persist and self._writer is not None:
             self._writer.append(self._to_record(entry))
         self._entries.append(entry)
         logging.debug("Appended trajectory entry for step %s (total entries: %s).", entry.step, len(self))
-        return entry
 
     def find(self, action: Mapping[str, Any], **identity_values: object) -> TrajectoryEntry | None:
         identity_components = self._construct_components(
@@ -345,14 +349,12 @@ def _to_record(self, entry: TrajectoryEntry) -> dict[str, object]:
         return record
 
 
-def _components_by_type(components: Sequence[object]) -> dict[type[object], object]:
+def _validate_components(components: Sequence[object]) -> None:
     non_dataclasses = [type(component).__name__ for component in components if not dataclasses.is_dataclass(component)]
     if non_dataclasses:
         raise TypeError(f"trajectory components must be dataclass instances: {', '.join(non_dataclasses)}")
-    by_type = {type(component): component for component in components}
-    if len(by_type) != len(components):
+    if len({type(component) for component in components}) != len(components):
         raise ValueError("components cannot contain duplicate component types")
-    return by_type
 
 
 def _validate_schema(
diff --git a/src/cloudai/systems/slurm/single_sbatch_runner.py b/src/cloudai/systems/slurm/single_sbatch_runner.py
index 7258d9f68..cf873465e 100644
--- a/src/cloudai/systems/slurm/single_sbatch_runner.py
+++ b/src/cloudai/systems/slurm/single_sbatch_runner.py
@@ -219,6 +219,9 @@ def handle_dse(self):
                 next_tr.step = idx
                 next_tr.output_path = self.get_job_output_path(next_tr)
 
+                if not next_tr.test.constraint_check(next_tr, self.system):
+                    continue
+
                 gym.test_run = next_tr
                 observation = gym.get_observation({})
                 reward = gym.compute_reward(observation)
diff --git a/tests/test_trajectory.py b/tests/test_trajectory.py
index 07672b67a..cff067ef5 100644
--- a/tests/test_trajectory.py
+++ b/tests/test_trajectory.py
@@ -183,7 +183,7 @@ def test_entry_contains_and_retrieves_typed_components() -> None:
 
     assert entry.components == (result, env_params, metrics)
     assert entry.get(TrialResult) is result
-    assert entry.get(EnvParamsSample) is env_params
+    assert entry.get(EnvParamsSample) is not env_params
     assert entry.get(LoggingMetrics) is metrics
 
 

From a0ad79cd5a04a433390de782c9b5a114fab29bba Mon Sep 17 00:00:00 2001
From: Alex Manley <amanley@nvidia.com>
Date: Wed, 15 Jul 2026 10:52:07 -0700
Subject: [PATCH 11/14] [Trajectory] fix linting

---
 src/cloudai/configurator/trajectory.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/src/cloudai/configurator/trajectory.py b/src/cloudai/configurator/trajectory.py
index 6914b3f68..a6cebb3da 100644
--- a/src/cloudai/configurator/trajectory.py
+++ b/src/cloudai/configurator/trajectory.py
@@ -287,8 +287,7 @@ def _identity_for(self, components: Sequence[object]) -> dict[type[object], obje
 
     def _freeze_identity_components(self, components: Sequence[object]) -> tuple[object, ...]:
         return tuple(
-            _freeze_component(component) if type(component) in self._identity else component
-            for component in components
+            _freeze_component(component) if type(component) in self._identity else component for component in components
         )
 
     def _construct_components(

From ddab127810878657cc4c9a12cb2a9515e112ebf2 Mon Sep 17 00:00:00 2001
From: Alex Manley <amanley@nvidia.com>
Date: Wed, 15 Jul 2026 11:10:40 -0700
Subject: [PATCH 12/14] [Trajectory] add agent_config flag to override
 trajectory file_type

---
 src/cloudai/cli/handlers.py            | 1 +
 src/cloudai/configurator/base_agent.py | 1 +
 src/cloudai/configurator/trajectory.py | 5 +++--
 3 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/src/cloudai/cli/handlers.py b/src/cloudai/cli/handlers.py
index 3f6488f86..0f050e618 100644
--- a/src/cloudai/cli/handlers.py
+++ b/src/cloudai/cli/handlers.py
@@ -157,6 +157,7 @@ def handle_dse_job(runner: Runner, args: argparse.Namespace) -> int:
                 test_run=test_run,
                 runner=runner.runner,
                 rewards=agent_config.rewards,
+                trajectory_file_type=agent_config.trajectory_file_type,
             )
             if agent_config.start_action == "first":
                 logging.info(f"Using deterministic first sweep for the chosen agent: {env.first_sweep}.")
diff --git a/src/cloudai/configurator/base_agent.py b/src/cloudai/configurator/base_agent.py
index d5f1162c8..8a8dec29e 100644
--- a/src/cloudai/configurator/base_agent.py
+++ b/src/cloudai/configurator/base_agent.py
@@ -49,6 +49,7 @@ class BaseAgentConfig(BaseModel):
         default_factory=RewardOverrides,
         description="Reward and observation overrides for the agent.",
     )
+    trajectory_file_type: Literal["csv", "jsonl"] = "csv"
 
 
 class BaseAgent(ABC):
diff --git a/src/cloudai/configurator/trajectory.py b/src/cloudai/configurator/trajectory.py
index a6cebb3da..a7abd7276 100644
--- a/src/cloudai/configurator/trajectory.py
+++ b/src/cloudai/configurator/trajectory.py
@@ -198,9 +198,10 @@ def __init__(
 
         component_names = ", ".join(component_type.__name__ for component_type in self._component_types)
         logging.debug(
-            "Initialized Trajectory with component types %s and %s entries.",
-            component_names,
+            "Initialized Trajectory with %s warm-start entries and persistence to local trajectory.%s. Entries contain component types: [%s].",
             len(self),
+            str(file_type),
+            component_names,
         )
 
     @staticmethod

From 496596c5f9968bf67f1454fe10700b52725c94bd Mon Sep 17 00:00:00 2001
From: Alex Manley <amanley@nvidia.com>
Date: Wed, 15 Jul 2026 11:23:49 -0700
Subject: [PATCH 13/14] [Trajectory] fix for linter/tests

---
 src/cloudai/configurator/trajectory.py | 2 +-
 tests/test_trajectory.py               | 5 ++++-
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/src/cloudai/configurator/trajectory.py b/src/cloudai/configurator/trajectory.py
index a7abd7276..b80bbaca9 100644
--- a/src/cloudai/configurator/trajectory.py
+++ b/src/cloudai/configurator/trajectory.py
@@ -198,7 +198,7 @@ def __init__(
 
         component_names = ", ".join(component_type.__name__ for component_type in self._component_types)
         logging.debug(
-            "Initialized Trajectory with %s warm-start entries and persistence to local trajectory.%s. Entries contain component types: [%s].",
+            "Initializing Trajectory: entries=%s, file_type=%s, components=[%s]. ",
             len(self),
             str(file_type),
             component_names,
diff --git a/tests/test_trajectory.py b/tests/test_trajectory.py
index cff067ef5..b043d84f9 100644
--- a/tests/test_trajectory.py
+++ b/tests/test_trajectory.py
@@ -345,7 +345,10 @@ def test_trajectory_logs_lifecycle_and_lookup(caplog: pytest.LogCaptureFixture)
         assert trajectory.find({"x": 1}) is entry
         assert trajectory.find({"x": 2}) is None
 
-    assert "Initialized Trajectory with component types TrialResult and 0 entries." in caplog.messages
+    assert (
+        "Initialized Trajectory with 0 warm-start entries and persistence to local trajectory.csv. "
+        "Entries contain component types: [TrialResult]."
+    ) in caplog.messages
     assert "Appended trajectory entry for step 1 (total entries: 1)." in caplog.messages
     assert "Found matching trajectory entry at step 1 for action {'x': 1}." in caplog.messages
     assert "No matching trajectory entry found for action {'x': 2}." in caplog.messages

From 3e9c4045392b090bfbf3cbc7a5b9f1377edc60b5 Mon Sep 17 00:00:00 2001
From: Alex Manley <amanley@nvidia.com>
Date: Wed, 15 Jul 2026 11:33:31 -0700
Subject: [PATCH 14/14] [Trajectory] final fixes, cleanup

---
 src/cloudai/configurator/trajectory.py           | 13 ++++++++++++-
 .../systems/slurm/single_sbatch_runner.py        |  4 +++-
 tests/test_trajectory.py                         | 16 ++++++++++++----
 3 files changed, 27 insertions(+), 6 deletions(-)

diff --git a/src/cloudai/configurator/trajectory.py b/src/cloudai/configurator/trajectory.py
index b80bbaca9..68fc544a4 100644
--- a/src/cloudai/configurator/trajectory.py
+++ b/src/cloudai/configurator/trajectory.py
@@ -198,7 +198,7 @@ def __init__(
 
         component_names = ", ".join(component_type.__name__ for component_type in self._component_types)
         logging.debug(
-            "Initializing Trajectory: entries=%s, file_type=%s, components=[%s]. ",
+            "Initializing Trajectory: entries=%s, file_type=%s, components=[%s].",
             len(self),
             str(file_type),
             component_names,
@@ -353,6 +353,17 @@ def _validate_components(components: Sequence[object]) -> None:
     non_dataclasses = [type(component).__name__ for component in components if not dataclasses.is_dataclass(component)]
     if non_dataclasses:
         raise TypeError(f"trajectory components must be dataclass instances: {', '.join(non_dataclasses)}")
+    mutable_identity_components = [
+        type(component).__name__
+        for component in components
+        if getattr(type(component), "contributes_to_identity", False)
+        and not cast(Any, type(component)).__dataclass_params__.frozen
+    ]
+    if mutable_identity_components:
+        raise TypeError(
+            "identity-contributing trajectory components must be frozen dataclasses: "
+            f"{', '.join(mutable_identity_components)}"
+        )
     if len({type(component) for component in components}) != len(components):
         raise ValueError("components cannot contain duplicate component types")
 
diff --git a/src/cloudai/systems/slurm/single_sbatch_runner.py b/src/cloudai/systems/slurm/single_sbatch_runner.py
index cf873465e..8bc1213a6 100644
--- a/src/cloudai/systems/slurm/single_sbatch_runner.py
+++ b/src/cloudai/systems/slurm/single_sbatch_runner.py
@@ -211,7 +211,9 @@ def handle_dse(self):
             agent_class = registry.get_agent(tr.test.agent)
             agent_config_data = tr.test.agent_config or {}
             agent_config = agent_class.get_config_class()(**agent_config_data)
-            gym = CloudAIGymEnv(tr, self, rewards=agent_config.rewards)
+            gym = CloudAIGymEnv(
+                tr, self, rewards=agent_config.rewards, trajectory_file_type=agent_config.trajectory_file_type
+            )
 
             for idx, combination in enumerate(tr.all_combinations, start=1):
                 sampled_env_params = gym.params.sample(idx) if gym.params is not None else {}
diff --git a/tests/test_trajectory.py b/tests/test_trajectory.py
index b043d84f9..2584c13f9 100644
--- a/tests/test_trajectory.py
+++ b/tests/test_trajectory.py
@@ -42,6 +42,12 @@ class CacheContext:
     cache_context: Mapping[str, Any]
 
 
+@dataclasses.dataclass
+class MutableCacheContext:
+    contributes_to_identity: ClassVar[bool] = True
+    cache_context: Mapping[str, Any]
+
+
 def _entry(step: int, action: Mapping[str, Any] | None = None) -> TrajectoryEntry:
     return TrajectoryEntry(
         step=step,
@@ -195,6 +201,11 @@ def test_entry_rejects_duplicate_component_types() -> None:
         )
 
 
+def test_entry_rejects_mutable_identity_component_dataclasses() -> None:
+    with pytest.raises(TypeError, match="must be frozen dataclasses: MutableCacheContext"):
+        TrajectoryEntry(step=1, components=(MutableCacheContext({"hardware": "H100"}),))
+
+
 def test_trajectory_validates_its_fixed_component_schema() -> None:
     trajectory = Trajectory(components=(EnvParamsSample, LoggingMetrics))
 
@@ -345,10 +356,7 @@ def test_trajectory_logs_lifecycle_and_lookup(caplog: pytest.LogCaptureFixture)
         assert trajectory.find({"x": 1}) is entry
         assert trajectory.find({"x": 2}) is None
 
-    assert (
-        "Initialized Trajectory with 0 warm-start entries and persistence to local trajectory.csv. "
-        "Entries contain component types: [TrialResult]."
-    ) in caplog.messages
+    assert "Initializing Trajectory: entries=0, file_type=csv, components=[TrialResult]." in caplog.messages
     assert "Appended trajectory entry for step 1 (total entries: 1)." in caplog.messages
     assert "Found matching trajectory entry at step 1 for action {'x': 1}." in caplog.messages
     assert "No matching trajectory entry found for action {'x': 2}." in caplog.messages