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/18] [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/18] [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/18] [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/18] [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/18] [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/18] [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/18] [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/18] [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/18] [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/18] [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/18] [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/18] [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/18] [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/18] [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

From 424b2645743928ac70951fc8bda61029634ade6a Mon Sep 17 00:00:00 2001
From: Alex Manley <amanley@nvidia.com>
Date: Tue, 21 Jul 2026 14:26:09 -0700
Subject: [PATCH 15/18] [Trajectory] update to simplify and use pandas

---
 src/cloudai/cli/handlers.py                   |   1 -
 src/cloudai/configurator/__init__.py          |  16 +-
 src/cloudai/configurator/base_agent.py        |   1 -
 src/cloudai/configurator/cloudai_gym.py       |  83 +--
 src/cloudai/configurator/trajectory.py        | 477 ++++--------------
 src/cloudai/report_generator/dse_report.py    |  33 +-
 .../systems/slurm/single_sbatch_runner.py     |  11 +-
 tests/test_cloudaigym.py                      |  96 ++--
 tests/test_handlers.py                        |   8 +-
 tests/test_reporter.py                        |  42 +-
 tests/test_single_sbatch_runner.py            |   3 +-
 tests/test_trajectory.py                      | 396 +++++----------
 12 files changed, 368 insertions(+), 799 deletions(-)

diff --git a/src/cloudai/cli/handlers.py b/src/cloudai/cli/handlers.py
index 0f050e618..3f6488f86 100644
--- a/src/cloudai/cli/handlers.py
+++ b/src/cloudai/cli/handlers.py
@@ -157,7 +157,6 @@ 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/__init__.py b/src/cloudai/configurator/__init__.py
index 9d00017cd..3965bc93b 100644
--- a/src/cloudai/configurator/__init__.py
+++ b/src/cloudai/configurator/__init__.py
@@ -19,27 +19,13 @@
 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,
-)
+from .trajectory import Trajectory
 
 __all__ = [
     "BaseAgent",
     "BaseGym",
     "CloudAIGymEnv",
-    "CsvTrajectoryWriter",
-    "EnvParamsSample",
     "GridSearchAgent",
     "GymnasiumAdapter",
-    "JsonLinesTrajectoryWriter",
     "Trajectory",
-    "TrajectoryEntry",
-    "TrajectoryWriter",
-    "TrialResult",
 ]
diff --git a/src/cloudai/configurator/base_agent.py b/src/cloudai/configurator/base_agent.py
index 8a8dec29e..d5f1162c8 100644
--- a/src/cloudai/configurator/base_agent.py
+++ b/src/cloudai/configurator/base_agent.py
@@ -49,7 +49,6 @@ 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/cloudai_gym.py b/src/cloudai/configurator/cloudai_gym.py
index d68416ab5..85c60ee3c 100644
--- a/src/cloudai/configurator/cloudai_gym.py
+++ b/src/cloudai/configurator/cloudai_gym.py
@@ -17,7 +17,7 @@
 import copy
 import logging
 from pathlib import Path
-from typing import Any, Dict, Literal, Optional, Tuple
+from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, cast
 
 from cloudai.core import METRIC_ERROR, BaseRunner, Registry, TestRun
 from cloudai.util.lazy_imports import lazy
@@ -25,12 +25,10 @@
 from .base_agent import RewardOverrides
 from .base_gym import BaseGym
 from .env_params import EnvParams, ObsLeafDescriptor
-from .trajectory import (
-    EnvParamsSample,
-    Trajectory,
-    TrajectoryEntry,
-    TrialResult,
-)
+from .trajectory import Trajectory
+
+if TYPE_CHECKING:
+    import pandas as pd
 
 
 class CloudAIGymEnv(BaseGym):
@@ -45,8 +43,6 @@ def __init__(
         test_run: TestRun,
         runner: BaseRunner,
         rewards: RewardOverrides,
-        *,
-        trajectory_file_type: Literal["csv", "jsonl"] = "csv",
     ):
         """
         Initialize the Gym environment using the TestRun object.
@@ -55,7 +51,6 @@ def __init__(
             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
@@ -64,7 +59,7 @@ def __init__(
         self.max_steps = test_run.test.agent_steps
         self.reward_function = Registry().get_reward_function(test_run.test.agent_reward_function)
         self.params: EnvParams | None = EnvParams.from_test(test_run.test)
-        self.trajectory = self._new_trajectory(trajectory_file_type)
+        self.trajectory = self._new_trajectory()
         super().__init__()
 
     @property
@@ -147,16 +142,15 @@ 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_trial_result.reward,
-                cached_result.step,
+                cached_result["reward"],
+                cached_result["step"],
             )
-            observation = list(cached_trial_result.observation)
-            reward = cached_trial_result.reward
+            observation = {
+                metric: cached_result[f"observation.{metric}"] for metric in self.test_run.test.agent_metrics
+            }
+            reward = cast(float, cached_result["reward"])
         else:
             if not self.test_run.test.constraint_check(self.test_run, self.runner.system):
                 logging.info("Constraint check failed. Skipping step.")
@@ -182,21 +176,18 @@ def step(self, action: Any) -> Tuple[list, float, bool, dict]:
                 self.test_run.step = new_tr.step
                 self.test_run.output_path = new_tr.output_path
 
-            observation = self.get_observation(action)
+            observation = self.get_observation()
             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),
+            action=action,
             reward=reward,
             observation=observation,
-            **optional_values,
+            env_params=sampled_env_params,
         )
 
-        return observation, reward, False, info
+        return list(observation.values()), reward, False, info
 
     def render(self, mode: str = "human"):
         """
@@ -217,38 +208,35 @@ def seed(self, seed: Optional[int] = None):
         if seed is not None:
             lazy.np.random.seed(seed)
 
-    def compute_reward(self, observation: list) -> float:
+    def compute_reward(self, observation: dict[str, Any]) -> float:
         """
         Compute a reward based on the TestRun result.
 
         Args:
-            observation (list): The observation list containing the average value.
+            observation: Metric values keyed by configured metric name.
 
         Returns:
             float: Reward value.
         """
-        return self.reward_function(observation)
+        return self.reward_function(list(observation.values()))
 
-    def get_observation(self, action: Any) -> list:
+    def get_observation(self) -> dict[str, Any]:
         """
         Get the observation from the TestRun object.
 
-        Args:
-            action (Any): Action taken by the agent.
-
         Returns:
-            list: The observation.
+            Metric values keyed by configured metric name.
         """
         all_metrics = self.test_run.test.agent_metrics
         if not all_metrics:
             raise ValueError("No agent metrics defined for the test run")
 
-        observation = []
+        observation: dict[str, Any] = {}
         for metric in all_metrics:
             v = self.test_run.get_metric_value(self.runner.system, metric)
             if v is METRIC_ERROR:
                 v = self.rewards.metric_failure
-            observation.append(v)
+            observation[metric] = v
         return observation
 
     def structured_observation_descriptors(self) -> Optional[Dict[str, ObsLeafDescriptor]]:
@@ -269,12 +257,8 @@ 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:
-        return Trajectory(
-            iteration_dir=lambda: self.iteration_dir,
-            file_type=file_type,
-            components=(EnvParamsSample,) if self.params is not None else (),
-        )
+    def _new_trajectory(self) -> Trajectory:
+        return Trajectory(iteration_dir=lambda: self.iteration_dir)
 
     @property
     def iteration_dir(self) -> Path:
@@ -288,17 +272,6 @@ def trajectory_file_path(self) -> Path:
             raise RuntimeError("trajectory persistence is not configured")
         return path
 
-    def get_cached_trajectory_result(self, action: Any, env_params: dict[str, Any]) -> TrajectoryEntry | None:
-        """
-        Return a cached entry only when the full trial identity matches.
-
-        Trial identity is ``(action, env_params)``: env-randomized parameters
-        change the workload's behaviour, so a trial repeating the same action
-        under a different ``env_params`` sample must miss and re-run. Empty
-        env_params on both sides is the back-compat path for workloads that
-        do not declare any ``[env_params.*]`` block. The sample is passed in (a
-        per-trial local owned by ``step``), exactly like ``action``.
-        """
-        if not env_params:
-            return self.trajectory.find(action)
-        return self.trajectory.find(action, env_params=dict(env_params))
+    def get_cached_trajectory_result(self, action: Any, env_params: dict[str, Any]) -> pd.Series | None:
+        """Return the first trajectory row matching the action and environment parameters."""
+        return self.trajectory.find(action=action, env_params=env_params)
diff --git a/src/cloudai/configurator/trajectory.py b/src/cloudai/configurator/trajectory.py
index 68fc544a4..cf6b503f3 100644
--- a/src/cloudai/configurator/trajectory.py
+++ b/src/cloudai/configurator/trajectory.py
@@ -14,422 +14,167 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Ordered trajectory steps composed from typed dataclass components."""
+"""Pandas-native trajectory storage with flat, namespaced columns."""
 
 from __future__ import annotations
 
-import copy
 import csv
-import dataclasses
-import json
 import logging
-from collections.abc import Callable, Iterator, Mapping, Sequence
+from collections.abc import Callable, Mapping, Sequence
+from numbers import Integral
 from pathlib import Path
-from types import MappingProxyType
-from typing import Any, ClassVar, Literal, Protocol, TypeVar, cast, overload
+from typing import TYPE_CHECKING, Any
 
-ComponentT = TypeVar("ComponentT")
+from cloudai.util.lazy_imports import lazy
 
+if TYPE_CHECKING:
+    import pandas as pd
 
-@dataclasses.dataclass(frozen=True)
-class EnvParamsSample:
-    """Environment-parameter values sampled for one trial."""
 
-    contributes_to_identity: ClassVar[bool] = True
-
-    env_params: Mapping[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]
-
-    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:
-    """One immutable step containing a fixed set of typed data components."""
-
-    step: int
-    components: tuple[object, ...]
-
-    def __post_init__(self) -> None:
-        """Validate the entry and snapshot its identity components."""
-        if self.step < 1:
-            raise ValueError(f"trajectory step must be positive; got {self.step}")
-        _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."""
-        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."""
+class Trajectory:
+    """An ordered DataFrame of DSE steps persisted 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)
-        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}")
-
-        with path.open("a", newline="") as file:
-            writer = csv.DictWriter(file, fieldnames=self._fields)
-            if write_header:
-                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.
-    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.
-    """
-
     def __init__(
         self,
-        entries: Sequence[TrajectoryEntry] = (),
+        dataframe: pd.DataFrame | None = None,
         *,
         iteration_dir: Path | Callable[[], Path] | None = None,
-        file_type: Literal["csv", "jsonl"] = "csv",
-        components: Sequence[type[object]] = (),
     ) -> None:
-        self._component_types = (TrialResult, *components)
-        self._components = frozenset(self._component_types)
-        self._identity = frozenset(
-            component_type
-            for component_type in self._component_types
-            if getattr(component_type, "contributes_to_identity", False)
+        self._iteration_dir = iteration_dir
+        self._dataframe = lazy.pd.DataFrame() if dataframe is None else dataframe.copy(deep=True)
+        self._validate_dataframe()
+        self._dataframe = self._dataframe.astype(object)
+        self._fields: tuple[str, ...] | None = (
+            tuple(self._dataframe.columns) if len(self._dataframe.columns) > 0 else None
         )
-        if len(self._components) != len(self._component_types):
-            raise ValueError("components cannot contain duplicate types")
-
-        self._writer = self._create_writer(iteration_dir, file_type)
-        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(
-            "Initializing Trajectory: entries=%s, file_type=%s, components=[%s].",
+            "Initializing Trajectory: entries=%s, columns=%s.",
             len(self),
-            str(file_type),
-            component_names,
+            list(self._dataframe.columns),
         )
 
-    @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: ...
-
-    @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)
+        """Return the number of trajectory rows."""
+        return len(self._dataframe)
 
-    def __iter__(self) -> Iterator[TrajectoryEntry]:
-        """Iterate over entries in step order."""
-        return iter(self._entries)
+    @property
+    def dataframe(self) -> pd.DataFrame:
+        """Return a copy of the trajectory DataFrame for analysis."""
+        return self._dataframe.copy(deep=True)
 
     @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(self._freeze_identity_components(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, 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)
-                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 _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]],
-        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
+        """Return the current trajectory CSV path when persistence is configured."""
+        if self._iteration_dir is None:
+            return None
+        iteration_dir = self._iteration_dir() if callable(self._iteration_dir) else self._iteration_dir
+        return iteration_dir / self.file_name
 
-    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: _thaw(getattr(component, _field.name)) for _field in self._fields_by_type[component_type]}
+    def append(self, *, step: int, **values: object) -> pd.Series:
+        """Flatten, persist, and store one trajectory row."""
+        self._validate_step(step)
+        if len(self) and step <= self._dataframe.iloc[-1]["step"]:
+            raise ValueError(
+                f"trajectory steps must increase: last step is {self._dataframe.iloc[-1]['step']}, got {step}"
             )
-        return record
 
+        record: dict[str, object] = {"step": step}
+        for domain, value in values.items():
+            _flatten_value(record, domain, value)
 
-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")
-
+        fields = tuple(record)
+        if self._fields is not None and fields != self._fields:
+            raise ValueError(f"trajectory record fields changed: expected {self._fields}, got {fields}")
 
-def _validate_schema(
-    actual: frozenset[type[object]],
-    expected: frozenset[type[object]],
-    context: str,
-) -> None:
-    if actual == expected:
-        return
+        row = lazy.pd.Series(record, dtype=object)
+        self._persist(row, fields)
+
+        row_frame = row.to_frame().T.astype(object)
+        self._dataframe = lazy.pd.concat([self._dataframe, row_frame], ignore_index=True).astype(object)
+        self._fields = fields
+        logging.debug("Appended trajectory row for step %s (total rows: %s).", step, len(self))
+        return row.copy(deep=True)
+
+    def find(self, **values: object) -> pd.Series | None:
+        """Return a copy of the first row matching all supplied domain values."""
+        criteria: dict[str, object] = {}
+        for domain, value in values.items():
+            _flatten_value(criteria, domain, value)
+
+        if any(field not in self._dataframe.columns for field in criteria):
+            return None
+
+        for _, row in self._dataframe.iterrows():
+            if all(_values_match_exact(row[field], value) for field, value in criteria.items()):
+                logging.debug("Found matching trajectory row at step %s for %s.", row["step"], values)
+                return row.copy(deep=True)
+        logging.debug("No matching trajectory row found for %s.", values)
+        return None
 
-    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 _validate_dataframe(self) -> None:
+        if not self._dataframe.columns.is_unique:
+            raise ValueError("trajectory dataframe columns must be unique")
+        non_string_columns = [column for column in self._dataframe.columns if not isinstance(column, str)]
+        if non_string_columns:
+            raise TypeError(f"trajectory dataframe columns must be strings: {non_string_columns}")
+        if self._dataframe.empty and len(self._dataframe.columns) == 0:
+            return
+        if "step" not in self._dataframe.columns:
+            raise ValueError("trajectory dataframe must contain a step column")
+
+        previous_step: int | None = None
+        for step in self._dataframe["step"]:
+            self._validate_step(step)
+            if previous_step is not None and step <= previous_step:
+                raise ValueError(f"trajectory steps must increase: last step is {previous_step}, got {step}")
+            previous_step = int(step)
 
+    @staticmethod
+    def _validate_step(step: object) -> None:
+        if isinstance(step, bool) or not isinstance(step, Integral) or step < 1:
+            raise ValueError(f"trajectory step must be a positive integer; got {step}")
 
-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 _persist(self, row: pd.Series, fields: tuple[str, ...]) -> None:
+        path = self.output_path
+        if path is None:
+            return
 
+        path.parent.mkdir(parents=True, exist_ok=True)
+        write_header = not path.exists() or path.stat().st_size == 0
+        if not write_header:
+            with path.open(newline="") as file:
+                existing_fields = tuple(next(csv.reader(file), ()))
+            if existing_fields != fields:
+                raise ValueError(f"trajectory file fields do not match: expected {fields}, got {existing_fields}")
 
-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
+        row.to_frame().T.to_csv(path, mode="a", header=write_header, index=False)
+        logging.debug("Wrote trajectory row to %s.", path)
 
 
-def _thaw(value: Any) -> Any:
-    """Convert frozen containers to values supported by trajectory writers."""
+def _flatten_value(record: dict[str, object], key: str, value: object) -> None:
+    """Flatten mappings into dot-separated columns while preserving leaf values."""
     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
+        for child_key, child_value in value.items():
+            if not isinstance(child_key, str):
+                raise TypeError(f"trajectory mapping keys must be strings: {child_key}")
+            _flatten_value(record, f"{key}.{child_key}", child_value)
+        return
+    if key in record:
+        raise ValueError(f"trajectory values produce duplicate column: {key}")
+    record[key] = value
 
 
 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)):
+    if isinstance(left, Sequence) and not isinstance(left, (str, bytes)):
         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
+    return bool(left == right)
diff --git a/src/cloudai/report_generator/dse_report.py b/src/cloudai/report_generator/dse_report.py
index 3089d09cf..f9fdea5ed 100644
--- a/src/cloudai/report_generator/dse_report.py
+++ b/src/cloudai/report_generator/dse_report.py
@@ -132,10 +132,6 @@ def _safe_literal_eval(raw: Any, default: Any) -> Any:
 
 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)
@@ -143,6 +139,15 @@ def load_trajectory_dataframe(iteration_dir: Path) -> tuple[Path, Any] | None:
     return None
 
 
+def _prefixed_values(row: dict[str, Any], prefix: str) -> dict[str, Any]:
+    """Return non-missing values from flat namespaced columns."""
+    return {
+        key.removeprefix(prefix): value
+        for key, value in row.items()
+        if key.startswith(prefix) and not lazy.pd.isna(value)
+    }
+
+
 def _format_scalar(value: Any) -> str:
     if isinstance(value, float):
         return f"{value:.4f}".rstrip("0").rstrip(".")
@@ -268,12 +273,20 @@ def _build_trajectory_steps(
     steps: list[TrajectoryStep] = []
     for row in df.to_dict(orient="records"):
         step_no = int(row["step"])
-        action = _safe_literal_eval(row.get("action"), {})
-        if not isinstance(action, dict):
-            action = {}
-        observation = _safe_literal_eval(row.get("observation"), [])
-        if not isinstance(observation, list):
-            observation = [observation]
+        action = _prefixed_values(row, "action.")
+        if not action:
+            action = _safe_literal_eval(row.get("action"), {})
+            if not isinstance(action, dict):
+                action = {}
+
+        flat_observation = _prefixed_values(row, "observation.")
+        if flat_observation:
+            metric_names = test_case.test.agent_metrics or list(flat_observation)
+            observation = [flat_observation[metric] for metric in metric_names if metric in flat_observation]
+        else:
+            observation = _safe_literal_eval(row.get("observation"), [])
+            if not isinstance(observation, list):
+                observation = [observation]
         step_run = runs_by_step.get(step_no)
         steps.append(
             TrajectoryStep(
diff --git a/src/cloudai/systems/slurm/single_sbatch_runner.py b/src/cloudai/systems/slurm/single_sbatch_runner.py
index 8bc1213a6..decb1c0ea 100644
--- a/src/cloudai/systems/slurm/single_sbatch_runner.py
+++ b/src/cloudai/systems/slurm/single_sbatch_runner.py
@@ -211,9 +211,7 @@ 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, trajectory_file_type=agent_config.trajectory_file_type
-            )
+            gym = CloudAIGymEnv(tr, self, rewards=agent_config.rewards)
 
             for idx, combination in enumerate(tr.all_combinations, start=1):
                 sampled_env_params = gym.params.sample(idx) if gym.params is not None else {}
@@ -225,17 +223,14 @@ def handle_dse(self):
                     continue
 
                 gym.test_run = next_tr
-                observation = gym.get_observation({})
+                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,
+                    env_params=sampled_env_params,
                 )
 
     def completed_test_runs(self, job: BaseJob) -> list[TestRun]:
diff --git a/tests/test_cloudaigym.py b/tests/test_cloudaigym.py
index 858019cd4..e6749c294 100644
--- a/tests/test_cloudaigym.py
+++ b/tests/test_cloudaigym.py
@@ -18,15 +18,13 @@
 from typing import cast
 from unittest.mock import MagicMock, PropertyMock, patch
 
+import pandas as pd
 import pytest
 
 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
@@ -44,17 +42,18 @@
 from tests.test_env_params import EnvVarCmdArgs, EnvVarTestDefinition
 
 
-def _trajectory_entry(
+def _trajectory_row(
     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),
-    )
+) -> dict[str, object]:
+    return {
+        "step": step,
+        **{f"action.{key}": value for key, value in action.items()},
+        "reward": reward,
+        "observation.default": observation[0],
+    }
 
 
 @pytest.fixture
@@ -150,7 +149,8 @@ def test_compute_reward(reward_function, test_cases, base_tr: TestRun):
     env = CloudAIGymEnv(test_run=base_tr, runner=MagicMock(), rewards=RewardOverrides())
 
     for input_value, expected_reward in test_cases:
-        reward = env.compute_reward(input_value)
+        observation = {f"metric_{index}": value for index, value in enumerate(input_value)}
+        reward = env.compute_reward(observation)
         assert reward == expected_reward
 
 
@@ -376,12 +376,12 @@ def test_apply_params_set__preserves_installables_state(setup_env: tuple[TestRun
     ("entries", "action", "expected_step"),
     [
         ([], {"x": 1}, None),
-        ([_trajectory_entry(1, {"x": 1}, 1, [1])], {"x": 1}, 1),
-        ([_trajectory_entry(1, {"x": 1.0}, 1, [1])], {"x": 1}, None),
+        ([_trajectory_row(1, {"x": 1}, 1, [1])], {"x": 1}, 1),
+        ([_trajectory_row(1, {"x": 1.0}, 1, [1])], {"x": 1}, None),
         (
             [
-                _trajectory_entry(1, {"x": 1.0}, 1, [1]),
-                _trajectory_entry(2, {"x": 1}, 1, [1]),
+                _trajectory_row(1, {"x": 1.0}, 1, [1]),
+                _trajectory_row(2, {"x": 1}, 1, [1]),
             ],
             {"x": 1},
             2,
@@ -391,7 +391,7 @@ def test_apply_params_set__preserves_installables_state(setup_env: tuple[TestRun
 def test_get_cached_trajectory_result(
     base_tr: TestRun,
     tmp_path: Path,
-    entries: list[TrajectoryEntry],
+    entries: list[dict[str, object]],
     action: dict[str, object],
     expected_step: int | None,
 ) -> None:
@@ -404,13 +404,13 @@ 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.trajectory = Trajectory(entries)
+    env.trajectory = Trajectory(dataframe=pd.DataFrame(entries, dtype=object))
 
     actual = env.get_cached_trajectory_result(action, {})
     if actual is None:
         assert expected_step is None
     else:
-        assert actual.step == expected_step
+        assert actual["step"] == expected_step
 
 
 def test_cached_step_appends_trajectory_record(nemorun: NeMoRunTestDefinition, tmp_path: Path) -> None:
@@ -433,7 +433,7 @@ def test_cached_step_appends_trajectory_record(nemorun: NeMoRunTestDefinition, t
     env = CloudAIGymEnv(test_run=test_run, runner=runner, rewards=RewardOverrides())
     cached_action = {"trainer.max_steps": 1000}
     env.test_run.current_iteration = 0
-    env.trajectory.append(step=1, action=cached_action, reward=0.42, observation=[0.84])
+    env.trajectory.append(step=1, action=cached_action, reward=0.42, observation={"default": 0.84})
 
     env.test_run.step = 4
     obs, reward, done, _info = env.step(cached_action)
@@ -442,22 +442,20 @@ def test_cached_step_appends_trajectory_record(nemorun: NeMoRunTestDefinition, t
     assert reward == 0.42
     assert obs == [0.84]
     assert done is False
-    rows = env.trajectory
+    rows = env.trajectory.dataframe
     assert len(rows) == 2
-    assert rows[-1].step == 5, (
+    assert rows.iloc[-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."
     )
-    result = rows[-1].get(TrialResult)
-    assert result is not None
-    assert result.reward == 0.42
-    assert result.action == cached_action
+    assert rows.iloc[-1]["reward"] == 0.42
+    assert rows.iloc[-1]["action.trainer.max_steps"] == cached_action["trainer.max_steps"]
 
     trajectory_path = env.trajectory_file_path
     assert trajectory_path.exists()
     assert trajectory_path.name == "trajectory.csv"
     contents = trajectory_path.read_text().strip().splitlines()
-    assert contents[0] == "step,action,reward,observation"
+    assert contents[0] == "step,action.trainer.max_steps,reward,observation.default"
     assert contents[-1].startswith("5,")
 
 
@@ -465,10 +463,14 @@ def _seed_cached_entry_with_env_params(
     env: CloudAIGymEnv, action: dict[str, object], env_params: dict[str, object]
 ) -> None:
     """Seed an environment-parameter-aware trajectory with one entry."""
-    trajectory = Trajectory(
-        components=(EnvParamsSample,),
+    trajectory = Trajectory()
+    trajectory.append(
+        step=1,
+        action=action,
+        reward=0.5,
+        observation={"default": 100.0},
+        env_params=dict(env_params),
     )
-    trajectory.append(step=1, action=action, reward=0.5, observation=[100.0], env_params=dict(env_params))
     env.trajectory = trajectory
 
 
@@ -524,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 = Trajectory([_trajectory_entry(1, {"x": 10}, 0.5, [100.0])])
+    env.trajectory = Trajectory(dataframe=pd.DataFrame([_trajectory_row(1, {"x": 10}, 0.5, [100.0])], dtype=object))
     # Note: neither the cached entry nor the trial carries env_params -> existing behavior.
 
     result = env.get_cached_trajectory_result({"x": 10}, {})
@@ -570,9 +572,9 @@ def test_step_reruns_workload_when_env_params_change(tmp_path: Path) -> None:
 
     env = CloudAIGymEnv(test_run=test_run, runner=runner, rewards=RewardOverrides())
     action = {"paddle_width": 4}
-    fake_obs = iter([[100.0], [50.0]])
+    fake_obs = iter([{"default": 100.0}, {"default": 50.0}])
 
-    with patch.object(env, "get_observation", side_effect=lambda _action: next(fake_obs)):
+    with patch.object(env, "get_observation", side_effect=lambda: next(fake_obs)):
         env.test_run.step = 0
         *_, info1 = env.step(action)  # samples ball_speed=3
         *_, info2 = env.step(action)  # samples ball_speed=1
@@ -616,17 +618,17 @@ def test_env_params_are_recorded_in_trajectory_output(tmp_path: Path) -> None:
 
     env = CloudAIGymEnv(test_run=test_run, runner=runner, rewards=RewardOverrides())
     action_a, action_b = {"paddle_width": 4}, {"paddle_width": 8}
-    fake_obs = iter([[100.0], [50.0], [25.0]])
+    fake_obs = iter([{"default": 100.0}, {"default": 50.0}, {"default": 25.0}])
 
-    with patch.object(env, "get_observation", side_effect=lambda _action: next(fake_obs)):
+    with patch.object(env, "get_observation", side_effect=lambda: next(fake_obs)):
         env.test_run.step = 0
         for action in (action_a, action_b, action_a):
             env.step(action)
 
     rows = env.trajectory_file_path.read_text().strip().splitlines()
-    assert rows[0] == "step,action,reward,observation,env_params"
+    assert rows[0] == "step,action.paddle_width,reward,observation.default,env_params.ball_speed"
     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:])
+    assert pd.read_csv(env.trajectory_file_path)["env_params.ball_speed"].notna().all()
 
 
 def test_constraint_failure_omits_the_complete_trajectory_row(tmp_path: Path) -> None:
@@ -661,10 +663,10 @@ def test_constraint_failure_omits_the_complete_trajectory_row(tmp_path: Path) ->
 
     # Step 2 fails the constraint; steps 1 and 3 survive. get_observation is only
     # reached on the surviving steps, so it yields exactly two values.
-    fake_obs = iter([[100.0], [25.0]])
+    fake_obs = iter([{"default": 100.0}, {"default": 25.0}])
     with (
         patch.object(EnvVarTestDefinition, "constraint_check", side_effect=[True, False, True]),
-        patch.object(env, "get_observation", side_effect=lambda _action: next(fake_obs)),
+        patch.object(env, "get_observation", side_effect=lambda: next(fake_obs)),
     ):
         env.test_run.step = 0
         for action in ({"paddle_width": 4}, {"paddle_width": 6}, {"paddle_width": 8}):
@@ -721,7 +723,7 @@ def test_step_cache_hit_with_declared_env_params_records_complete_trajectory_row
         step=1,
         action=action,
         reward=0.42,
-        observation=[0.84],
+        observation={"default": 0.84},
         env_params=expected_sample,
     )
     env.test_run.step = 1
@@ -737,13 +739,12 @@ def test_step_cache_hit_with_declared_env_params_records_complete_trajectory_row
     )
 
     trajectory_rows = env.trajectory_file_path.read_text().strip().splitlines()
-    assert trajectory_rows[0] == "step,action,reward,observation,env_params"
+    assert trajectory_rows[0] == "step,action.paddle_width,reward,observation.default,env_params.ball_speed"
     assert trajectory_rows[-1].startswith("2,")
-    assert "ball_speed" in trajectory_rows[-1]
+    assert pd.read_csv(env.trajectory_file_path).iloc[-1]["env_params.ball_speed"] == expected_sample["ball_speed"]
 
-    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, (
+    traj_rows = env.trajectory.dataframe
+    assert len(traj_rows) == 2 and traj_rows.iloc[-1]["env_params.ball_speed"] == expected_sample["ball_speed"], (
         "cache-hit trajectory entry must record the per-trial env_params sample"
     )
 
@@ -786,7 +787,7 @@ def test_step_overlays_env_params_onto_cmd_args(tmp_path: Path) -> None:
     env = CloudAIGymEnv(test_run=test_run, runner=runner, rewards=RewardOverrides())
     expected = _random.Random("42:ball_speed:1").choice([1, 2])
 
-    with patch.object(env, "get_observation", side_effect=lambda _action: [1.0]):
+    with patch.object(env, "get_observation", side_effect=lambda: {"default": 1.0}):
         env.test_run.step = 0
         env.step({"paddle_width": 4})
 
@@ -840,12 +841,11 @@ def test_csv_trajectory_has_no_env_params_column_when_not_declared(
         test_run=test_run,
         runner=runner,
         rewards=RewardOverrides(),
-        trajectory_file_type="csv",
     )
 
     assert env.params is None, "no env_params declared -> no EnvParams object"
-    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"
+    env.trajectory.append(step=1, action={}, reward=1.0, observation={"default": 1.0})
+    assert env.trajectory_file_path.read_text().splitlines()[0] == "step,reward,observation.default"
 
 
 def _dr_env(tmp_path: Path, candidates: list, *, seed: int = 42) -> CloudAIGymEnv:
diff --git a/tests/test_handlers.py b/tests/test_handlers.py
index bd2a76c46..f4e96ac7f 100644
--- a/tests/test_handlers.py
+++ b/tests/test_handlers.py
@@ -212,11 +212,11 @@ def _job_output_path(tr: TestRun, create: bool = 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, 1, -1.0, -1.0],
+            [2, 1, -1.0, -1.0],
+            [3, 2, -1.0, -1.0],
         ],
-        columns=["step", "action", "reward", "observation"],
+        columns=["step", "action.candidate", "reward", "observation.default"],
     )
     pd.testing.assert_frame_equal(actual_trajectory, expected_trajectory)
 
diff --git a/tests/test_reporter.py b/tests/test_reporter.py
index f6420be52..4c5c6d971 100644
--- a/tests/test_reporter.py
+++ b/tests/test_reporter.py
@@ -16,7 +16,6 @@
 
 import copy
 import csv
-import json
 import tarfile
 from dataclasses import asdict
 from pathlib import Path
@@ -47,20 +46,21 @@
 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")
-    with (tmp_path / "trajectory.csv").open("w", newline="") as file:
-        writer = csv.DictWriter(file, fieldnames=("step", "action", "reward", "observation"))
+def test_load_trajectory_dataframe_reads_flat_csv(tmp_path: Path) -> None:
+    csv_path = tmp_path / "trajectory.csv"
+    with csv_path.open("w", newline="") as file:
+        writer = csv.DictWriter(file, fieldnames=("step", "action.x", "reward", "observation.default"))
         writer.writeheader()
-        writer.writerow({"step": 99, "action": {"x": 99}, "reward": 99.0, "observation": [99.0]})
+        writer.writerow({"step": 1, "action.x": 2, "reward": 3.0, "observation.default": 4.0})
 
     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]}]
+    assert path == csv_path
+    assert dataframe.to_dict(orient="records") == [
+        {"step": 1, "action.x": 2, "reward": 3.0, "observation.default": 4.0}
+    ]
 
 
 class TestLoadTestTuns:
@@ -405,16 +405,33 @@ def _create_dse_iteration(
     results_root: Path,
     slurm_metadata: SlurmSystemMetadata,
     steps: list[dict[str, Any]],
+    *,
+    flat_trajectory: bool = False,
 ) -> None:
     iteration_dir = results_root / case.name / str(iteration)
     iteration_dir.mkdir(parents=True, exist_ok=True)
 
     with (iteration_dir / "trajectory.csv").open("w", newline="") as f:
         writer = csv.writer(f)
-        writer.writerow(["step", "action", "reward", "observation"])
+        action_columns = [f"action.{key}" for key in steps[0]["action"]]
+        writer.writerow(
+            ["step", *action_columns, "reward", "observation.default"]
+            if flat_trajectory
+            else ["step", "action", "reward", "observation"]
+        )
         for step in steps:
             step_no = step["step"]
-            writer.writerow([step_no, step["action"], step["reward"], step["observation"]])
+            row = (
+                [
+                    step_no,
+                    *(step["action"][column.removeprefix("action.")] for column in action_columns),
+                    step["reward"],
+                    step["observation"][0],
+                ]
+                if flat_trajectory
+                else [step_no, step["action"], step["reward"], step["observation"]]
+            )
+            writer.writerow(row)
 
             step_dir = iteration_dir / str(step_no)
             step_dir.mkdir(parents=True, exist_ok=True)
@@ -432,9 +449,11 @@ def _create_dse_iteration(
                 toml.dump(TestRunDetails.from_test_run(step_tr, "", "").model_dump(mode="json"), dump_file)
 
 
+@pytest.mark.parametrize("flat_trajectory", [False, True])
 def test_dse_reporter(
     slurm_system: SlurmSystem,
     slurm_metadata: SlurmSystemMetadata,
+    flat_trajectory: bool,
 ) -> None:
     slurm_metadata.system.gpu_arch_type = "NVIDIA H100 80GB HBM3"
 
@@ -482,6 +501,7 @@ def test_dse_reporter(
         results_root=slurm_system.output_path,
         slurm_metadata=slurm_metadata,
         steps=steps,
+        flat_trajectory=flat_trajectory,
     )
 
     scenario = TestScenario(
diff --git a/tests/test_single_sbatch_runner.py b/tests/test_single_sbatch_runner.py
index 09a63a841..a3563bf7f 100644
--- a/tests/test_single_sbatch_runner.py
+++ b/tests/test_single_sbatch_runner.py
@@ -14,7 +14,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import ast
 import copy
 import re
 from pathlib import Path
@@ -592,4 +591,4 @@ def test_dse_env_params_are_applied_to_runs_and_trajectory(dse_tr: TestRun, slur
     runner.handle_dse()
 
     dataframe = pd.read_csv(trajectory_path)
-    assert dataframe["env_params"].map(ast.literal_eval).tolist() == expected_samples
+    assert dataframe["env_params.nthreads"].tolist() == [sample["nthreads"] for sample in expected_samples]
diff --git a/tests/test_trajectory.py b/tests/test_trajectory.py
index 2584c13f9..3a0accbeb 100644
--- a/tests/test_trajectory.py
+++ b/tests/test_trajectory.py
@@ -14,349 +14,189 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import dataclasses
-import json
-from collections.abc import Mapping
 from pathlib import Path
-from typing import Any, ClassVar
 
+import pandas as pd
 import pytest
 
-from cloudai.configurator.trajectory import (
-    CsvTrajectoryWriter,
-    EnvParamsSample,
-    Trajectory,
-    TrajectoryEntry,
-    TrialResult,
-)
-
-
-@dataclasses.dataclass(frozen=True)
-class LoggingMetrics:
-    logging_metrics: Mapping[str, float]
-
-
-@dataclasses.dataclass(frozen=True)
-class CacheContext:
-    contributes_to_identity: ClassVar[bool] = True
-    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,
-        components=(TrialResult(action=action or {"x": step}, reward=float(step), observation=[step]),),
-    )
+from cloudai.configurator import Trajectory
 
 
-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, monkeypatch: pytest.MonkeyPatch) -> None:
-    def fail_write(_writer: CsvTrajectoryWriter, _record: Mapping[str, object]) -> None:
-        raise OSError("write failed")
-
-    monkeypatch.setattr(CsvTrajectoryWriter, "append", fail_write)
+def test_append_flattens_domains_into_dataframe_columns(tmp_path: Path) -> None:
     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])
-
-    assert len(trajectory) == 0
-
-
-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.csv").exists()
-
-
-def test_append_writes_component_values_to_csv(tmp_path: Path) -> None:
-    path = tmp_path / "trajectory.csv"
-    trajectory = Trajectory(
-        iteration_dir=tmp_path,
-        file_type="csv",
-        components=(LoggingMetrics,),
+    row = trajectory.append(
+        step=1,
+        action={"model": {"layers": 8}},
+        reward=0.95,
+        observation={"throughput": 120.0},
+        env_params={"network_speed": 100},
+        logging={"gpu_power_watts": 610.0},
     )
 
-    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}",
-    ]
-
+    expected = {
+        "step": 1,
+        "action.model.layers": 8,
+        "reward": 0.95,
+        "observation.throughput": 120.0,
+        "env_params.network_speed": 100,
+        "logging.gpu_power_watts": 610.0,
+    }
+    assert row.to_dict() == expected
+    assert trajectory.dataframe.to_dict(orient="records") == [expected]
+    assert all(pd.api.types.is_object_dtype(dtype) for dtype in trajectory.dataframe.dtypes)
+    pd.testing.assert_frame_equal(pd.read_csv(tmp_path / "trajectory.csv"), pd.DataFrame([expected]))
 
-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])
+def test_dataframe_and_returned_rows_are_copies() -> None:
+    trajectory = Trajectory()
+    row = trajectory.append(step=1, action={"x": 1}, reward=1.0)
 
-    assert path.read_text().splitlines() == [
-        "step,action,reward,observation",
-        "1,{'x': 1},1.0,[1]",
-    ]
+    row["reward"] = 99.0
+    dataframe = trajectory.dataframe
+    dataframe.loc[0, "reward"] = 88.0
 
+    assert trajectory.dataframe.loc[0, "reward"] == 1.0
 
-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])
+def test_find_matches_flattened_subset_and_preserves_exact_types() -> None:
+    trajectory = Trajectory()
+    first = trajectory.append(step=1, action={"x": 1.0}, reward=1.0, env_params={"speed": 2})
+    trajectory.append(step=2, action={"x": 1}, reward=2.0, env_params={"speed": 2})
 
-    assert path.read_text().splitlines() == [
-        "step,action,reward,observation",
-        "1,{'x': 1},1.0,[1]",
-    ]
+    match = trajectory.find(action={"x": 1.0}, env_params={"speed": 2})
 
+    assert match is not None
+    assert match.to_dict() == first.to_dict()
+    assert trajectory.find(action={"x": True}, env_params={"speed": 2}) is None
 
-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])
+def test_find_returns_none_for_unknown_columns() -> None:
+    trajectory = Trajectory()
+    trajectory.append(step=1, action={"x": 1}, reward=1.0)
 
-    assert path.read_text() == "step,reward,action,observation\n"
+    assert trajectory.find(env_params={"speed": 2}) is None
 
 
-def test_append_writes_generic_records_as_json_lines(tmp_path: Path) -> None:
-    path = tmp_path / "trajectory.jsonl"
-    trajectory = Trajectory(
-        iteration_dir=tmp_path,
-        file_type="jsonl",
-        components=(EnvParamsSample, LoggingMetrics),
-    )
+def test_trajectory_rejects_invalid_or_non_increasing_steps() -> None:
+    trajectory = Trajectory()
 
-    trajectory.append(
-        step=1,
-        action={"x": 1},
-        reward=1.0,
-        observation=[1],
-        env_params={"speed": 8},
-        logging_metrics={"gpu_power_watts": 610.0},
-    )
+    with pytest.raises(ValueError, match="positive integer"):
+        trajectory.append(step=0, action={"x": 1})
 
-    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},
-    }
+    trajectory.append(step=2, action={"x": 1})
+    with pytest.raises(ValueError, match="steps must increase"):
+        trajectory.append(step=2, action={"x": 2})
 
 
-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))
+def test_first_row_establishes_fixed_schema() -> None:
+    trajectory = Trajectory()
+    trajectory.append(step=1, action={"x": 1}, reward=1.0)
 
-    assert entry.components == (result, env_params, metrics)
-    assert entry.get(TrialResult) is result
-    assert entry.get(EnvParamsSample) is not env_params
-    assert entry.get(LoggingMetrics) is metrics
+    with pytest.raises(ValueError, match="record fields changed"):
+        trajectory.append(step=2, action={"x": 2}, reward=2.0, logging={"power": 600.0})
 
 
-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_flattening_rejects_duplicate_columns() -> None:
+    trajectory = Trajectory()
 
+    with pytest.raises(ValueError, match=r"duplicate column: action\.x"):
+        trajectory.append(step=1, action={"x": 1}, **{"action.x": 2})
 
-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_flattening_rejects_non_string_mapping_keys() -> None:
+    trajectory = Trajectory()
 
-def test_trajectory_validates_its_fixed_component_schema() -> None:
-    trajectory = Trajectory(components=(EnvParamsSample, LoggingMetrics))
+    with pytest.raises(TypeError, match="mapping keys must be strings"):
+        trajectory.append(step=1, action={1: "x"})
 
-    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_warm_start_dataframe_is_copied_and_not_replayed(tmp_path: Path) -> None:
+    dataframe = pd.DataFrame([{"step": 1, "action.x": 1, "reward": 1.0}], dtype=object)
+    trajectory = Trajectory(dataframe=dataframe, iteration_dir=tmp_path)
+    dataframe.loc[0, "reward"] = 99.0
 
+    assert len(trajectory) == 1
+    assert trajectory.dataframe.loc[0, "reward"] == 1.0
+    assert not (tmp_path / "trajectory.csv").exists()
 
-def test_find_uses_only_components_that_contribute_to_identity() -> None:
-    trajectory = Trajectory(
-        components=(EnvParamsSample, LoggingMetrics),
-    )
-    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
+@pytest.mark.parametrize(
+    ("dataframe", "message"),
+    [
+        (pd.DataFrame([{"action.x": 1}]), "must contain a step column"),
+        (pd.DataFrame([{"step": 0}]), "positive integer"),
+        (pd.DataFrame([{"step": 2}, {"step": 1}]), "steps must increase"),
+    ],
+)
+def test_warm_start_dataframe_validation(dataframe: pd.DataFrame, message: str) -> None:
+    with pytest.raises(ValueError, match=message):
+        Trajectory(dataframe=dataframe)
 
 
-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,))
+def test_warm_start_dataframe_rejects_duplicate_columns() -> None:
+    dataframe = pd.DataFrame([[1, 2]], columns=["step", "step"])
 
-    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(ValueError, match="columns must be unique"):
+        Trajectory(dataframe=dataframe)
 
-    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_csv_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)
 
-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,))
+    assert path.read_text().splitlines() == ["step,action.x,reward", "1,1,1.0"]
 
-    context.cache_context["hardware"]["gpus"].append(16)
 
-    assert trajectory.find({"x": 1}, cache_context={"hardware": {"gpus": [8]}}) is trajectory[0]
+def test_csv_reuses_a_matching_header_without_duplication(tmp_path: Path) -> None:
+    path = tmp_path / "trajectory.csv"
+    path.write_text("step,action.x,reward\n")
 
+    Trajectory(iteration_dir=tmp_path).append(step=1, action={"x": 1}, reward=1.0)
 
-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 path.read_text().splitlines() == ["step,action.x,reward", "1,1,1.0"]
 
-    assert trajectory.find({"x": 1}) is first
 
+def test_csv_rejects_an_existing_mismatched_header(tmp_path: Path) -> None:
+    path = tmp_path / "trajectory.csv"
+    path.write_text("step,reward,action.x\n")
 
-def test_find_requires_the_configured_identity_component_types() -> None:
-    trajectory = Trajectory(
-        components=(EnvParamsSample,),
-    )
+    with pytest.raises(ValueError, match="file fields do not match"):
+        Trajectory(iteration_dir=tmp_path).append(step=1, action={"x": 1}, reward=1.0)
 
-    with pytest.raises(TypeError, match="missing: env_params"):
-        trajectory.find({"x": 1})
+    assert path.read_text() == "step,reward,action.x\n"
 
-    with pytest.raises(TypeError, match="unexpected: logging_metrics"):
-        trajectory.find({"x": 1}, logging_metrics={})
 
+def test_persistence_failure_does_not_append_to_memory(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    trajectory = Trajectory(iteration_dir=tmp_path)
 
-def test_find_preserves_exact_value_types_inside_components() -> None:
-    trajectory = Trajectory(
-        components=(EnvParamsSample,),
-    )
-    trajectory.append(step=1, action={"x": 1}, reward=1.0, observation=[1], env_params={"speed": 1.0})
+    def fail_write(*_args: object, **_kwargs: object) -> None:
+        raise OSError("write failed")
 
-    assert trajectory.find({"x": 1}, env_params={"speed": 1}) is None
+    monkeypatch.setattr(pd.DataFrame, "to_csv", fail_write)
 
+    with pytest.raises(OSError, match="write failed"):
+        trajectory.append(step=1, action={"x": 1}, reward=1.0)
 
-def test_find_preserves_exact_action_value_types() -> None:
-    trajectory = Trajectory([_entry(1, {"x": 1.0})])
+    assert len(trajectory) == 0
 
-    assert trajectory.find({"x": 1}) is None
 
+def test_callable_iteration_directory_is_resolved_per_append(tmp_path: Path) -> None:
+    current = tmp_path / "first"
+    trajectory = Trajectory(iteration_dir=lambda: current)
+    trajectory.append(step=1, action={"x": 1})
 
-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()
+    assert trajectory.output_path == current / "trajectory.csv"
 
 
 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 "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
+        trajectory.append(step=1, action={"x": 1}, reward=1.0)
+        assert trajectory.find(action={"x": 1}) is not None
+        assert trajectory.find(action={"x": 2}) is None
+
+    assert "Initializing Trajectory: entries=0, columns=[]." in caplog.messages
+    assert "Appended trajectory row for step 1 (total rows: 1)." in caplog.messages

From a3da4661c56f34d93f29c8b85d7ae1d6a2d2ac23 Mon Sep 17 00:00:00 2001
From: Alex Manley <amanley@nvidia.com>
Date: Wed, 22 Jul 2026 12:27:39 -0700
Subject: [PATCH 16/18] [Trajectory] fix failing tests

---
 src/cloudai/configurator/cloudai_gym.py |  2 +-
 tests/test_trajectory.py                | 12 ++++++------
 2 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/src/cloudai/configurator/cloudai_gym.py b/src/cloudai/configurator/cloudai_gym.py
index 85c60ee3c..920b193cd 100644
--- a/src/cloudai/configurator/cloudai_gym.py
+++ b/src/cloudai/configurator/cloudai_gym.py
@@ -272,6 +272,6 @@ def trajectory_file_path(self) -> Path:
             raise RuntimeError("trajectory persistence is not configured")
         return path
 
-    def get_cached_trajectory_result(self, action: Any, env_params: dict[str, Any]) -> pd.Series | None:
+    def get_cached_trajectory_result(self, action: Any, env_params: dict[str, Any]) -> "pd.Series" | None:
         """Return the first trajectory row matching the action and environment parameters."""
         return self.trajectory.find(action=action, env_params=env_params)
diff --git a/tests/test_trajectory.py b/tests/test_trajectory.py
index 3a0accbeb..e9e2955bd 100644
--- a/tests/test_trajectory.py
+++ b/tests/test_trajectory.py
@@ -122,16 +122,16 @@ def test_warm_start_dataframe_is_copied_and_not_replayed(tmp_path: Path) -> None
 
 
 @pytest.mark.parametrize(
-    ("dataframe", "message"),
+    ("records", "message"),
     [
-        (pd.DataFrame([{"action.x": 1}]), "must contain a step column"),
-        (pd.DataFrame([{"step": 0}]), "positive integer"),
-        (pd.DataFrame([{"step": 2}, {"step": 1}]), "steps must increase"),
+        ([{"action.x": 1}], "must contain a step column"),
+        ([{"step": 0}], "positive integer"),
+        ([{"step": 2}, {"step": 1}], "steps must increase"),
     ],
 )
-def test_warm_start_dataframe_validation(dataframe: pd.DataFrame, message: str) -> None:
+def test_warm_start_dataframe_validation(records: list[dict[str, object]], message: str) -> None:
     with pytest.raises(ValueError, match=message):
-        Trajectory(dataframe=dataframe)
+        Trajectory(dataframe=pd.DataFrame(records))
 
 
 def test_warm_start_dataframe_rejects_duplicate_columns() -> None:

From e15dab5fd577abccaedc08b49dc477de22e03afd Mon Sep 17 00:00:00 2001
From: Alex Manley <amanley@nvidia.com>
Date: Wed, 22 Jul 2026 12:33:46 -0700
Subject: [PATCH 17/18] [Trajectory] fix pyright check failing

---
 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 920b193cd..bfea7bb39 100644
--- a/src/cloudai/configurator/cloudai_gym.py
+++ b/src/cloudai/configurator/cloudai_gym.py
@@ -272,6 +272,6 @@ def trajectory_file_path(self) -> Path:
             raise RuntimeError("trajectory persistence is not configured")
         return path
 
-    def get_cached_trajectory_result(self, action: Any, env_params: dict[str, Any]) -> "pd.Series" | None:
+    def get_cached_trajectory_result(self, action: Any, env_params: dict[str, Any]) -> "pd.Series | None":
         """Return the first trajectory row matching the action and environment parameters."""
         return self.trajectory.find(action=action, env_params=env_params)

From 7c9fc4a0663493796ef5549b9be9197d6e5f862e Mon Sep 17 00:00:00 2001
From: Alex Manley <amanley@nvidia.com>
Date: Wed, 22 Jul 2026 14:15:23 -0700
Subject: [PATCH 18/18] [Trajectory] ensure RLDS format compatibility

---
 src/cloudai/configurator/cloudai_gym.py    |  10 +-
 src/cloudai/configurator/trajectory.py     | 131 +++++++----
 src/cloudai/report_generator/dse_report.py |  42 +---
 tests/test_cloudaigym.py                   |  51 +++--
 tests/test_handlers.py                     |   8 +-
 tests/test_reporter.py                     |  43 +---
 tests/test_single_sbatch_runner.py         |   8 +-
 tests/test_trajectory.py                   | 243 ++++++++++++++++-----
 8 files changed, 345 insertions(+), 191 deletions(-)

diff --git a/src/cloudai/configurator/cloudai_gym.py b/src/cloudai/configurator/cloudai_gym.py
index bfea7bb39..cdafd5ce5 100644
--- a/src/cloudai/configurator/cloudai_gym.py
+++ b/src/cloudai/configurator/cloudai_gym.py
@@ -59,7 +59,7 @@ def __init__(
         self.max_steps = test_run.test.agent_steps
         self.reward_function = Registry().get_reward_function(test_run.test.agent_reward_function)
         self.params: EnvParams | None = EnvParams.from_test(test_run.test)
-        self.trajectory = self._new_trajectory()
+        self.trajectory = Trajectory(iteration_dir=self.iteration_dir)
         super().__init__()
 
     @property
@@ -257,9 +257,6 @@ 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) -> Trajectory:
-        return Trajectory(iteration_dir=lambda: self.iteration_dir)
-
     @property
     def iteration_dir(self) -> Path:
         """Per-iteration output directory containing the trajectory output."""
@@ -267,10 +264,7 @@ def iteration_dir(self) -> Path:
 
     @property
     def trajectory_file_path(self) -> Path:
-        path = self.trajectory.output_path
-        if path is None:
-            raise RuntimeError("trajectory persistence is not configured")
-        return path
+        return self.trajectory.output_path
 
     def get_cached_trajectory_result(self, action: Any, env_params: dict[str, Any]) -> "pd.Series | None":
         """Return the first trajectory row matching the action and environment parameters."""
diff --git a/src/cloudai/configurator/trajectory.py b/src/cloudai/configurator/trajectory.py
index cf6b503f3..85543b5d9 100644
--- a/src/cloudai/configurator/trajectory.py
+++ b/src/cloudai/configurator/trajectory.py
@@ -20,7 +20,8 @@
 
 import csv
 import logging
-from collections.abc import Callable, Mapping, Sequence
+from collections.abc import Mapping, Sequence
+from copy import deepcopy
 from numbers import Integral
 from pathlib import Path
 from typing import TYPE_CHECKING, Any
@@ -32,23 +33,23 @@
 
 
 class Trajectory:
-    """An ordered DataFrame of DSE steps persisted to ``trajectory.csv``."""
+    """An ordered DataFrame of DSE steps persisted as core and metadata CSVs."""
 
     file_name = "trajectory.csv"
+    metadata_file_name = "metadata.csv"
+    _core_fields = ("step", "action", "reward", "observation")
+    _core_domains = frozenset(_core_fields)
 
     def __init__(
         self,
-        dataframe: pd.DataFrame | None = None,
         *,
-        iteration_dir: Path | Callable[[], Path] | None = None,
+        iteration_dir: Path,
+        dataframe: pd.DataFrame | None = None,
     ) -> None:
         self._iteration_dir = iteration_dir
-        self._dataframe = lazy.pd.DataFrame() if dataframe is None else dataframe.copy(deep=True)
+        self._dataframe = lazy.pd.DataFrame() if dataframe is None else _copy_dataframe(dataframe)
         self._validate_dataframe()
         self._dataframe = self._dataframe.astype(object)
-        self._fields: tuple[str, ...] | None = (
-            tuple(self._dataframe.columns) if len(self._dataframe.columns) > 0 else None
-        )
         logging.debug(
             "Initializing Trajectory: entries=%s, columns=%s.",
             len(self),
@@ -62,17 +63,27 @@ def __len__(self) -> int:
     @property
     def dataframe(self) -> pd.DataFrame:
         """Return a copy of the trajectory DataFrame for analysis."""
-        return self._dataframe.copy(deep=True)
+        return _copy_dataframe(self._dataframe)
 
     @property
-    def output_path(self) -> Path | None:
-        """Return the current trajectory CSV path when persistence is configured."""
-        if self._iteration_dir is None:
-            return None
-        iteration_dir = self._iteration_dir() if callable(self._iteration_dir) else self._iteration_dir
-        return iteration_dir / self.file_name
+    def output_path(self) -> Path:
+        """Return the trajectory CSV path."""
+        return self._iteration_dir / self.file_name
+
+    @property
+    def metadata_path(self) -> Path:
+        """Return the trajectory metadata CSV path."""
+        return self._iteration_dir / self.metadata_file_name
 
-    def append(self, *, step: int, **values: object) -> pd.Series:
+    def append(
+        self,
+        *,
+        step: int,
+        action: object,
+        reward: object,
+        observation: object,
+        **values: object,
+    ) -> pd.Series:
         """Flatten, persist, and store one trajectory row."""
         self._validate_step(step)
         if len(self) and step <= self._dataframe.iloc[-1]["step"]:
@@ -81,21 +92,22 @@ def append(self, *, step: int, **values: object) -> pd.Series:
             )
 
         record: dict[str, object] = {"step": step}
-        for domain, value in values.items():
+        domains = {"action": action, "reward": reward, "observation": observation, **values}
+        for domain, value in domains.items():
             _flatten_value(record, domain, value)
 
         fields = tuple(record)
-        if self._fields is not None and fields != self._fields:
-            raise ValueError(f"trajectory record fields changed: expected {self._fields}, got {fields}")
+        expected_fields = tuple(self._dataframe.columns)
+        if expected_fields and fields != expected_fields:
+            raise ValueError(f"trajectory record fields changed: expected {expected_fields}, got {fields}")
 
         row = lazy.pd.Series(record, dtype=object)
-        self._persist(row, fields)
+        self._persist(row, action=action, reward=reward, observation=observation)
 
         row_frame = row.to_frame().T.astype(object)
         self._dataframe = lazy.pd.concat([self._dataframe, row_frame], ignore_index=True).astype(object)
-        self._fields = fields
         logging.debug("Appended trajectory row for step %s (total rows: %s).", step, len(self))
-        return row.copy(deep=True)
+        return _copy_series(row)
 
     def find(self, **values: object) -> pd.Series | None:
         """Return a copy of the first row matching all supplied domain values."""
@@ -109,7 +121,7 @@ def find(self, **values: object) -> pd.Series | None:
         for _, row in self._dataframe.iterrows():
             if all(_values_match_exact(row[field], value) for field, value in criteria.items()):
                 logging.debug("Found matching trajectory row at step %s for %s.", row["step"], values)
-                return row.copy(deep=True)
+                return _copy_series(row)
         logging.debug("No matching trajectory row found for %s.", values)
         return None
 
@@ -136,21 +148,49 @@ def _validate_step(step: object) -> None:
         if isinstance(step, bool) or not isinstance(step, Integral) or step < 1:
             raise ValueError(f"trajectory step must be a positive integer; got {step}")
 
-    def _persist(self, row: pd.Series, fields: tuple[str, ...]) -> None:
-        path = self.output_path
-        if path is None:
-            return
+    def _persist(self, row: pd.Series, *, action: object, reward: object, observation: object) -> None:
+        core_row = lazy.pd.Series(
+            {
+                "step": row["step"],
+                "action": action,
+                "reward": reward,
+                "observation": list(observation.values()) if isinstance(observation, Mapping) else observation,
+            },
+            dtype=object,
+        )
+        metadata_fields = tuple(
+            field for field in row.index if field == "step" or field.split(".", maxsplit=1)[0] not in self._core_domains
+        )
+        metadata_row = row[list(metadata_fields)] if len(metadata_fields) > 1 else None
+
+        self.output_path.parent.mkdir(parents=True, exist_ok=True)
+        core_header = _validate_csv_header(self.output_path, self._core_fields)
+        metadata_header = (
+            _validate_csv_header(self.metadata_path, metadata_fields) if metadata_row is not None else False
+        )
+
+        core_row.to_frame().T.to_csv(self.output_path, mode="a", header=core_header, index=False)
+        logging.debug("Wrote trajectory row to %s.", self.output_path)
+        if metadata_row is not None:
+            metadata_row.to_frame().T.to_csv(
+                self.metadata_path,
+                mode="a",
+                header=metadata_header,
+                index=False,
+            )
+            logging.debug("Wrote trajectory metadata row to %s.", self.metadata_path)
 
-        path.parent.mkdir(parents=True, exist_ok=True)
-        write_header = not path.exists() or path.stat().st_size == 0
-        if not write_header:
-            with path.open(newline="") as file:
-                existing_fields = tuple(next(csv.reader(file), ()))
-            if existing_fields != fields:
-                raise ValueError(f"trajectory file fields do not match: expected {fields}, got {existing_fields}")
 
-        row.to_frame().T.to_csv(path, mode="a", header=write_header, index=False)
-        logging.debug("Wrote trajectory row to %s.", path)
+def _validate_csv_header(path: Path, fields: tuple[str, ...]) -> bool:
+    """Validate an existing CSV header and return whether a header must be written."""
+    write_header = not path.exists() or path.stat().st_size == 0
+    if write_header:
+        return True
+    with path.open(newline="") as file:
+        existing_fields = tuple(next(csv.reader(file), ()))
+    if existing_fields != fields:
+        raise ValueError(f"trajectory file fields do not match: expected {fields}, got {existing_fields}")
+    return False
 
 
 def _flatten_value(record: dict[str, object], key: str, value: object) -> None:
@@ -163,7 +203,24 @@ def _flatten_value(record: dict[str, object], key: str, value: object) -> None:
         return
     if key in record:
         raise ValueError(f"trajectory values produce duplicate column: {key}")
-    record[key] = value
+    record[key] = deepcopy(value)
+
+
+def _copy_dataframe(dataframe: pd.DataFrame) -> pd.DataFrame:
+    """Copy a DataFrame and recursively snapshot object cells."""
+    copied = dataframe.copy(deep=True).astype(object)
+    for row_index in range(dataframe.shape[0]):
+        for column_index in range(dataframe.shape[1]):
+            copied.iat[row_index, column_index] = deepcopy(dataframe.iat[row_index, column_index])
+    return copied
+
+
+def _copy_series(series: pd.Series) -> pd.Series:
+    """Copy a Series and recursively snapshot object cells."""
+    copied = series.copy(deep=True).astype(object)
+    for index in range(len(series)):
+        copied.iat[index] = deepcopy(series.iat[index])
+    return copied
 
 
 def _values_match_exact(left: Any, right: Any) -> bool:
diff --git a/src/cloudai/report_generator/dse_report.py b/src/cloudai/report_generator/dse_report.py
index f9fdea5ed..ddd0b8165 100644
--- a/src/cloudai/report_generator/dse_report.py
+++ b/src/cloudai/report_generator/dse_report.py
@@ -122,8 +122,6 @@ 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)
@@ -131,21 +129,11 @@ def _safe_literal_eval(raw: Any, default: Any) -> Any:
 
 
 def load_trajectory_dataframe(iteration_dir: Path) -> tuple[Path, Any] | None:
-    """Load trajectory output."""
-    csv_path = iteration_dir / "trajectory.csv"
-    if csv_path.is_file():
-        return csv_path, lazy.pd.read_csv(csv_path)
-
-    return None
-
-
-def _prefixed_values(row: dict[str, Any], prefix: str) -> dict[str, Any]:
-    """Return non-missing values from flat namespaced columns."""
-    return {
-        key.removeprefix(prefix): value
-        for key, value in row.items()
-        if key.startswith(prefix) and not lazy.pd.isna(value)
-    }
+    """Load a trajectory CSV as a DataFrame."""
+    trajectory_file = iteration_dir / "trajectory.csv"
+    if not trajectory_file.is_file():
+        return None
+    return trajectory_file, lazy.pd.read_csv(trajectory_file)
 
 
 def _format_scalar(value: Any) -> str:
@@ -273,20 +261,12 @@ def _build_trajectory_steps(
     steps: list[TrajectoryStep] = []
     for row in df.to_dict(orient="records"):
         step_no = int(row["step"])
-        action = _prefixed_values(row, "action.")
-        if not action:
-            action = _safe_literal_eval(row.get("action"), {})
-            if not isinstance(action, dict):
-                action = {}
-
-        flat_observation = _prefixed_values(row, "observation.")
-        if flat_observation:
-            metric_names = test_case.test.agent_metrics or list(flat_observation)
-            observation = [flat_observation[metric] for metric in metric_names if metric in flat_observation]
-        else:
-            observation = _safe_literal_eval(row.get("observation"), [])
-            if not isinstance(observation, list):
-                observation = [observation]
+        action = _safe_literal_eval(row.get("action"), {})
+        if not isinstance(action, dict):
+            action = {}
+        observation = _safe_literal_eval(row.get("observation"), [])
+        if not isinstance(observation, list):
+            observation = [observation]
         step_run = runs_by_step.get(step_no)
         steps.append(
             TrajectoryStep(
diff --git a/tests/test_cloudaigym.py b/tests/test_cloudaigym.py
index e6749c294..c1d8c3153 100644
--- a/tests/test_cloudaigym.py
+++ b/tests/test_cloudaigym.py
@@ -190,7 +190,9 @@ def test_tr_output_path(setup_env: tuple[TestRun, BaseRunner]):
         pytest.param(RewardOverrides(constraint_failure=-2.5), -2.5, id="custom_penalty"),
     ],
 )
-def test_constraint_failure(nemorun: NeMoRunTestDefinition, rewards: RewardOverrides, expected_reward: float):
+def test_constraint_failure(
+    nemorun: NeMoRunTestDefinition, tmp_path: Path, rewards: RewardOverrides, expected_reward: float
+):
     tdef = nemorun.model_copy(deep=True)
     tdef.cmd_args.data.global_batch_size = 8
     tdef.agent_metrics = ["default"]
@@ -202,6 +204,7 @@ def test_constraint_failure(nemorun: NeMoRunTestDefinition, rewards: RewardOverr
         reports={NeMoRunReportGenerationStrategy},
     )
     runner = MagicMock(spec=BaseRunner)
+    runner.scenario_root = tmp_path / "scenario"
     runner.system = MagicMock()
     env = CloudAIGymEnv(test_run=test_run, runner=runner, rewards=rewards)
 
@@ -404,7 +407,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.trajectory = Trajectory(dataframe=pd.DataFrame(entries, dtype=object))
+    env.trajectory = Trajectory(iteration_dir=env.iteration_dir, dataframe=pd.DataFrame(entries, dtype=object))
 
     actual = env.get_cached_trajectory_result(action, {})
     if actual is None:
@@ -455,7 +458,7 @@ def test_cached_step_appends_trajectory_record(nemorun: NeMoRunTestDefinition, t
     assert trajectory_path.exists()
     assert trajectory_path.name == "trajectory.csv"
     contents = trajectory_path.read_text().strip().splitlines()
-    assert contents[0] == "step,action.trainer.max_steps,reward,observation.default"
+    assert contents[0] == "step,action,reward,observation"
     assert contents[-1].startswith("5,")
 
 
@@ -463,7 +466,7 @@ def _seed_cached_entry_with_env_params(
     env: CloudAIGymEnv, action: dict[str, object], env_params: dict[str, object]
 ) -> None:
     """Seed an environment-parameter-aware trajectory with one entry."""
-    trajectory = Trajectory()
+    trajectory = Trajectory(iteration_dir=env.iteration_dir)
     trajectory.append(
         step=1,
         action=action,
@@ -526,7 +529,10 @@ 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 = Trajectory(dataframe=pd.DataFrame([_trajectory_row(1, {"x": 10}, 0.5, [100.0])], dtype=object))
+    env.trajectory = Trajectory(
+        iteration_dir=env.iteration_dir,
+        dataframe=pd.DataFrame([_trajectory_row(1, {"x": 10}, 0.5, [100.0])], dtype=object),
+    )
     # Note: neither the cached entry nor the trial carries env_params -> existing behavior.
 
     result = env.get_cached_trajectory_result({"x": 10}, {})
@@ -588,8 +594,8 @@ def test_step_reruns_workload_when_env_params_change(tmp_path: Path) -> None:
     )
 
 
-def test_env_params_are_recorded_in_trajectory_output(tmp_path: Path) -> None:
-    """Every recorded trial includes its sampled environment in the trajectory output."""
+def test_env_params_are_recorded_in_trajectory_metadata(tmp_path: Path) -> None:
+    """Every recorded trial includes its sampled environment in trajectory metadata."""
     tdef = EnvVarTestDefinition(
         name="dr",
         description="dr",
@@ -625,10 +631,12 @@ 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)
 
-    rows = env.trajectory_file_path.read_text().strip().splitlines()
-    assert rows[0] == "step,action.paddle_width,reward,observation.default,env_params.ball_speed"
-    assert [int(row.split(",", 1)[0]) for row in rows[1:]] == [1, 2, 3]
-    assert pd.read_csv(env.trajectory_file_path)["env_params.ball_speed"].notna().all()
+    trajectory = pd.read_csv(env.trajectory_file_path)
+    metadata = pd.read_csv(env.trajectory.metadata_path)
+    assert list(trajectory.columns) == ["step", "action", "reward", "observation"]
+    assert list(metadata.columns) == ["step", "env_params.ball_speed"]
+    assert trajectory["step"].tolist() == metadata["step"].tolist() == [1, 2, 3]
+    assert metadata["env_params.ball_speed"].notna().all()
 
 
 def test_constraint_failure_omits_the_complete_trajectory_row(tmp_path: Path) -> None:
@@ -673,12 +681,18 @@ def test_constraint_failure_omits_the_complete_trajectory_row(tmp_path: Path) ->
             env.step(action)
 
     trajectory_path = env.trajectory_file_path
+    metadata_path = env.trajectory.metadata_path
     traj_steps = (
         [int(line.split(",", 1)[0]) for line in trajectory_path.read_text().strip().splitlines()[1:]]
         if trajectory_path.exists()
         else []
     )
-    assert traj_steps == [1, 3]
+    metadata_steps = (
+        [int(line.split(",", 1)[0]) for line in metadata_path.read_text().strip().splitlines()[1:]]
+        if metadata_path.exists()
+        else []
+    )
+    assert traj_steps == metadata_steps == [1, 3]
 
 
 def test_step_cache_hit_with_declared_env_params_records_complete_trajectory_row(tmp_path: Path) -> None:
@@ -738,10 +752,11 @@ 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_rows = env.trajectory_file_path.read_text().strip().splitlines()
-    assert trajectory_rows[0] == "step,action.paddle_width,reward,observation.default,env_params.ball_speed"
-    assert trajectory_rows[-1].startswith("2,")
-    assert pd.read_csv(env.trajectory_file_path).iloc[-1]["env_params.ball_speed"] == expected_sample["ball_speed"]
+    trajectory = pd.read_csv(env.trajectory_file_path)
+    metadata = pd.read_csv(env.trajectory.metadata_path)
+    assert list(trajectory.columns) == ["step", "action", "reward", "observation"]
+    assert trajectory.iloc[-1]["step"] == 2
+    assert metadata.iloc[-1]["env_params.ball_speed"] == expected_sample["ball_speed"]
 
     traj_rows = env.trajectory.dataframe
     assert len(traj_rows) == 2 and traj_rows.iloc[-1]["env_params.ball_speed"] == expected_sample["ball_speed"], (
@@ -845,7 +860,8 @@ def test_csv_trajectory_has_no_env_params_column_when_not_declared(
 
     assert env.params is None, "no env_params declared -> no EnvParams object"
     env.trajectory.append(step=1, action={}, reward=1.0, observation={"default": 1.0})
-    assert env.trajectory_file_path.read_text().splitlines()[0] == "step,reward,observation.default"
+    assert env.trajectory_file_path.read_text().splitlines()[0] == "step,action,reward,observation"
+    assert not env.trajectory.metadata_path.exists()
 
 
 def _dr_env(tmp_path: Path, candidates: list, *, seed: int = 42) -> CloudAIGymEnv:
@@ -883,6 +899,7 @@ def test_metrics_only_env_opts_out(self, nemorun: NeMoRunTestDefinition, tmp_pat
         tdef.cmd_args.data.global_batch_size = 8
         test_run = TestRun(name="plain_tr", test=tdef, num_nodes=1, nodes=[])
         runner = MagicMock(spec=BaseRunner)
+        runner.scenario_root = tmp_path / "scenario"
         runner.system = MagicMock()
         env = CloudAIGymEnv(test_run=test_run, runner=runner, rewards=RewardOverrides())
 
diff --git a/tests/test_handlers.py b/tests/test_handlers.py
index f4e96ac7f..bd2a76c46 100644
--- a/tests/test_handlers.py
+++ b/tests/test_handlers.py
@@ -212,11 +212,11 @@ def _job_output_path(tr: TestRun, create: bool = True):
     actual_trajectory = pd.read_csv(trajectory_dir / "trajectory.csv")
     expected_trajectory = pd.DataFrame(
         data=[
-            [1, 1, -1.0, -1.0],
-            [2, 1, -1.0, -1.0],
-            [3, 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.candidate", "reward", "observation.default"],
+        columns=["step", "action", "reward", "observation"],
     )
     pd.testing.assert_frame_equal(actual_trajectory, expected_trajectory)
 
diff --git a/tests/test_reporter.py b/tests/test_reporter.py
index 4c5c6d971..95acd8ac9 100644
--- a/tests/test_reporter.py
+++ b/tests/test_reporter.py
@@ -28,7 +28,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, load_trajectory_dataframe
+from cloudai.report_generator.dse_report import build_dse_summaries
 from cloudai.reporter import DSEReporter, PerTestReporter, ReportItem, StatusReporter, TarballReporter
 from cloudai.systems.slurm.slurm_metadata import (
     MetadataCUDA,
@@ -46,23 +46,6 @@
 from cloudai.workloads.nccl_test import NCCLCmdArgs, NCCLTestDefinition
 
 
-def test_load_trajectory_dataframe_reads_flat_csv(tmp_path: Path) -> None:
-    csv_path = tmp_path / "trajectory.csv"
-    with csv_path.open("w", newline="") as file:
-        writer = csv.DictWriter(file, fieldnames=("step", "action.x", "reward", "observation.default"))
-        writer.writeheader()
-        writer.writerow({"step": 1, "action.x": 2, "reward": 3.0, "observation.default": 4.0})
-
-    loaded = load_trajectory_dataframe(tmp_path)
-
-    assert loaded is not None
-    path, dataframe = loaded
-    assert path == csv_path
-    assert dataframe.to_dict(orient="records") == [
-        {"step": 1, "action.x": 2, "reward": 3.0, "observation.default": 4.0}
-    ]
-
-
 class TestLoadTestTuns:
     def test_load_test_runs_behcnmark_sorted(self, slurm_system: SlurmSystem, benchmark_tr: TestRun) -> None:
         reporter = PerTestReporter(
@@ -405,33 +388,16 @@ def _create_dse_iteration(
     results_root: Path,
     slurm_metadata: SlurmSystemMetadata,
     steps: list[dict[str, Any]],
-    *,
-    flat_trajectory: bool = False,
 ) -> None:
     iteration_dir = results_root / case.name / str(iteration)
     iteration_dir.mkdir(parents=True, exist_ok=True)
 
     with (iteration_dir / "trajectory.csv").open("w", newline="") as f:
         writer = csv.writer(f)
-        action_columns = [f"action.{key}" for key in steps[0]["action"]]
-        writer.writerow(
-            ["step", *action_columns, "reward", "observation.default"]
-            if flat_trajectory
-            else ["step", "action", "reward", "observation"]
-        )
+        writer.writerow(["step", "action", "reward", "observation"])
         for step in steps:
             step_no = step["step"]
-            row = (
-                [
-                    step_no,
-                    *(step["action"][column.removeprefix("action.")] for column in action_columns),
-                    step["reward"],
-                    step["observation"][0],
-                ]
-                if flat_trajectory
-                else [step_no, step["action"], step["reward"], step["observation"]]
-            )
-            writer.writerow(row)
+            writer.writerow([step_no, step["action"], step["reward"], step["observation"]])
 
             step_dir = iteration_dir / str(step_no)
             step_dir.mkdir(parents=True, exist_ok=True)
@@ -449,11 +415,9 @@ def _create_dse_iteration(
                 toml.dump(TestRunDetails.from_test_run(step_tr, "", "").model_dump(mode="json"), dump_file)
 
 
-@pytest.mark.parametrize("flat_trajectory", [False, True])
 def test_dse_reporter(
     slurm_system: SlurmSystem,
     slurm_metadata: SlurmSystemMetadata,
-    flat_trajectory: bool,
 ) -> None:
     slurm_metadata.system.gpu_arch_type = "NVIDIA H100 80GB HBM3"
 
@@ -501,7 +465,6 @@ def test_dse_reporter(
         results_root=slurm_system.output_path,
         slurm_metadata=slurm_metadata,
         steps=steps,
-        flat_trajectory=flat_trajectory,
     )
 
     scenario = TestScenario(
diff --git a/tests/test_single_sbatch_runner.py b/tests/test_single_sbatch_runner.py
index a3563bf7f..6b3f029b3 100644
--- a/tests/test_single_sbatch_runner.py
+++ b/tests/test_single_sbatch_runner.py
@@ -579,7 +579,9 @@ def test_dse_env_params_are_applied_to_runs_and_trajectory(dse_tr: TestRun, slur
     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"
+    metadata_path = trajectory_path.with_name("metadata.csv")
     trajectory_path.unlink(missing_ok=True)
+    metadata_path.unlink(missing_ok=True)
 
     params = EnvParams.from_test(dse_tr.test)
     assert params is not None
@@ -590,5 +592,7 @@ def test_dse_env_params_are_applied_to_runs_and_trajectory(dse_tr: TestRun, slur
 
     runner.handle_dse()
 
-    dataframe = pd.read_csv(trajectory_path)
-    assert dataframe["env_params.nthreads"].tolist() == [sample["nthreads"] for sample in expected_samples]
+    trajectory = pd.read_csv(trajectory_path)
+    metadata = pd.read_csv(metadata_path)
+    assert trajectory["step"].tolist() == metadata["step"].tolist()
+    assert metadata["env_params.nthreads"].tolist() == [sample["nthreads"] for sample in expected_samples]
diff --git a/tests/test_trajectory.py b/tests/test_trajectory.py
index e9e2955bd..8b0a3654b 100644
--- a/tests/test_trajectory.py
+++ b/tests/test_trajectory.py
@@ -14,6 +14,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from ast import literal_eval
 from pathlib import Path
 
 import pandas as pd
@@ -45,24 +46,64 @@ def test_append_flattens_domains_into_dataframe_columns(tmp_path: Path) -> None:
     assert row.to_dict() == expected
     assert trajectory.dataframe.to_dict(orient="records") == [expected]
     assert all(pd.api.types.is_object_dtype(dtype) for dtype in trajectory.dataframe.dtypes)
-    pd.testing.assert_frame_equal(pd.read_csv(tmp_path / "trajectory.csv"), pd.DataFrame([expected]))
 
+    core = pd.read_csv(trajectory.output_path)
+    assert list(core.columns) == ["step", "action", "reward", "observation"]
+    action = core.loc[0, "action"]
+    observation = core.loc[0, "observation"]
+    assert isinstance(action, str)
+    assert isinstance(observation, str)
+    assert literal_eval(action) == {"model": {"layers": 8}}
+    assert core.loc[0, "reward"] == 0.95
+    assert literal_eval(observation) == [120.0]
+
+    metadata = pd.read_csv(trajectory.metadata_path)
+    assert metadata.to_dict(orient="records") == [
+        {
+            "step": 1,
+            "env_params.network_speed": 100,
+            "logging.gpu_power_watts": 610.0,
+        }
+    ]
+
+
+def test_dataframe_and_returned_rows_are_copies(tmp_path: Path) -> None:
+    trajectory = Trajectory(iteration_dir=tmp_path)
+    samples = [1, 2]
+    row = trajectory.append(
+        step=1,
+        action={"x": 1},
+        reward=1.0,
+        observation={"metric": 2.0},
+        logging={"samples": samples},
+    )
 
-def test_dataframe_and_returned_rows_are_copies() -> None:
-    trajectory = Trajectory()
-    row = trajectory.append(step=1, action={"x": 1}, reward=1.0)
-
+    samples.append(3)
     row["reward"] = 99.0
+    row_samples = row["logging.samples"]
+    assert isinstance(row_samples, list)
+    row_samples.append(4)
     dataframe = trajectory.dataframe
     dataframe.loc[0, "reward"] = 88.0
+    dataframe_samples = dataframe.loc[0, "logging.samples"]
+    assert isinstance(dataframe_samples, list)
+    dataframe_samples.append(5)
+    found = trajectory.find(action={"x": 1})
+    assert found is not None
+    found_samples = found["logging.samples"]
+    assert isinstance(found_samples, list)
+    found_samples.append(6)
 
     assert trajectory.dataframe.loc[0, "reward"] == 1.0
+    assert trajectory.dataframe.loc[0, "logging.samples"] == [1, 2]
 
 
-def test_find_matches_flattened_subset_and_preserves_exact_types() -> None:
-    trajectory = Trajectory()
-    first = trajectory.append(step=1, action={"x": 1.0}, reward=1.0, env_params={"speed": 2})
-    trajectory.append(step=2, action={"x": 1}, reward=2.0, env_params={"speed": 2})
+def test_find_matches_flattened_subset_and_preserves_exact_types(tmp_path: Path) -> None:
+    trajectory = Trajectory(iteration_dir=tmp_path)
+    first = trajectory.append(
+        step=1, action={"x": 1.0}, reward=1.0, observation={"metric": 1.0}, env_params={"speed": 2}
+    )
+    trajectory.append(step=2, action={"x": 1}, reward=2.0, observation={"metric": 2.0}, env_params={"speed": 2})
 
     match = trajectory.find(action={"x": 1.0}, env_params={"speed": 2})
 
@@ -71,53 +112,90 @@ def test_find_matches_flattened_subset_and_preserves_exact_types() -> None:
     assert trajectory.find(action={"x": True}, env_params={"speed": 2}) is None
 
 
-def test_find_returns_none_for_unknown_columns() -> None:
-    trajectory = Trajectory()
-    trajectory.append(step=1, action={"x": 1}, reward=1.0)
+def test_find_returns_none_for_unknown_columns(tmp_path: Path) -> None:
+    trajectory = Trajectory(iteration_dir=tmp_path)
+    trajectory.append(step=1, action={"x": 1}, reward=1.0, observation={"metric": 1.0})
 
     assert trajectory.find(env_params={"speed": 2}) is None
 
 
-def test_trajectory_rejects_invalid_or_non_increasing_steps() -> None:
-    trajectory = Trajectory()
+def test_trajectory_rejects_invalid_or_non_increasing_steps(tmp_path: Path) -> None:
+    trajectory = Trajectory(iteration_dir=tmp_path)
 
     with pytest.raises(ValueError, match="positive integer"):
-        trajectory.append(step=0, action={"x": 1})
+        trajectory.append(step=0, action={"x": 1}, reward=1.0, observation={"metric": 1.0})
 
-    trajectory.append(step=2, action={"x": 1})
+    trajectory.append(step=2, action={"x": 1}, reward=1.0, observation={"metric": 1.0})
     with pytest.raises(ValueError, match="steps must increase"):
-        trajectory.append(step=2, action={"x": 2})
+        trajectory.append(step=2, action={"x": 2}, reward=1.0, observation={"metric": 1.0})
 
 
-def test_first_row_establishes_fixed_schema() -> None:
-    trajectory = Trajectory()
-    trajectory.append(step=1, action={"x": 1}, reward=1.0)
+def test_first_row_establishes_fixed_schema(tmp_path: Path) -> None:
+    trajectory = Trajectory(iteration_dir=tmp_path)
+    trajectory.append(step=1, action={"x": 1}, reward=1.0, observation={"metric": 1.0})
 
     with pytest.raises(ValueError, match="record fields changed"):
-        trajectory.append(step=2, action={"x": 2}, reward=2.0, logging={"power": 600.0})
+        trajectory.append(
+            step=2,
+            action={"x": 2},
+            reward=2.0,
+            observation={"metric": 2.0},
+            logging={"power": 600.0},
+        )
 
 
-def test_flattening_rejects_duplicate_columns() -> None:
-    trajectory = Trajectory()
+def test_flattening_rejects_duplicate_columns(tmp_path: Path) -> None:
+    trajectory = Trajectory(iteration_dir=tmp_path)
 
     with pytest.raises(ValueError, match=r"duplicate column: action\.x"):
-        trajectory.append(step=1, action={"x": 1}, **{"action.x": 2})
+        trajectory.append(
+            step=1,
+            action={"x": 1},
+            reward=1.0,
+            observation={"metric": 1.0},
+            **{"action.x": 2},
+        )
 
 
-def test_flattening_rejects_non_string_mapping_keys() -> None:
-    trajectory = Trajectory()
+def test_flattening_rejects_non_string_mapping_keys(tmp_path: Path) -> None:
+    trajectory = Trajectory(iteration_dir=tmp_path)
 
     with pytest.raises(TypeError, match="mapping keys must be strings"):
-        trajectory.append(step=1, action={1: "x"})
+        trajectory.append(step=1, action={1: "x"}, reward=1.0, observation={"metric": 1.0})
+
+
+@pytest.mark.parametrize("missing", ["action", "reward", "observation"])
+def test_append_requires_core_values(tmp_path: Path, missing: str) -> None:
+    values = {"action": {"x": 1}, "reward": 1.0, "observation": {"metric": 2.0}}
+    del values[missing]
+
+    with pytest.raises(TypeError, match=missing):
+        Trajectory(iteration_dir=tmp_path).append(step=1, **values)
+
+    assert not (tmp_path / "trajectory.csv").exists()
 
 
 def test_warm_start_dataframe_is_copied_and_not_replayed(tmp_path: Path) -> None:
-    dataframe = pd.DataFrame([{"step": 1, "action.x": 1, "reward": 1.0}], dtype=object)
+    samples = [1, 2]
+    dataframe = pd.DataFrame(
+        [
+            {
+                "step": 1,
+                "action.x": 1,
+                "reward": 1.0,
+                "observation.metric": 2.0,
+                "logging.samples": samples,
+            }
+        ],
+        dtype=object,
+    )
     trajectory = Trajectory(dataframe=dataframe, iteration_dir=tmp_path)
     dataframe.loc[0, "reward"] = 99.0
+    samples.append(3)
 
     assert len(trajectory) == 1
     assert trajectory.dataframe.loc[0, "reward"] == 1.0
+    assert trajectory.dataframe.loc[0, "logging.samples"] == [1, 2]
     assert not (tmp_path / "trajectory.csv").exists()
 
 
@@ -129,44 +207,113 @@ def test_warm_start_dataframe_is_copied_and_not_replayed(tmp_path: Path) -> None
         ([{"step": 2}, {"step": 1}], "steps must increase"),
     ],
 )
-def test_warm_start_dataframe_validation(records: list[dict[str, object]], message: str) -> None:
+def test_warm_start_dataframe_validation(tmp_path: Path, records: list[dict[str, object]], message: str) -> None:
     with pytest.raises(ValueError, match=message):
-        Trajectory(dataframe=pd.DataFrame(records))
+        Trajectory(iteration_dir=tmp_path, dataframe=pd.DataFrame(records))
 
 
-def test_warm_start_dataframe_rejects_duplicate_columns() -> None:
+def test_warm_start_dataframe_rejects_duplicate_columns(tmp_path: Path) -> None:
     dataframe = pd.DataFrame([[1, 2]], columns=["step", "step"])
 
     with pytest.raises(ValueError, match="columns must be unique"):
-        Trajectory(dataframe=dataframe)
+        Trajectory(iteration_dir=tmp_path, dataframe=dataframe)
 
 
 def test_csv_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)
+    Trajectory(iteration_dir=tmp_path).append(step=1, action={"x": 1}, reward=1.0, observation={"metric": 2.0})
 
-    assert path.read_text().splitlines() == ["step,action.x,reward", "1,1,1.0"]
+    assert path.read_text().splitlines() == [
+        "step,action,reward,observation",
+        "1,{'x': 1},1.0,[2.0]",
+    ]
 
 
 def test_csv_reuses_a_matching_header_without_duplication(tmp_path: Path) -> None:
     path = tmp_path / "trajectory.csv"
-    path.write_text("step,action.x,reward\n")
+    path.write_text("step,action,reward,observation\n")
 
-    Trajectory(iteration_dir=tmp_path).append(step=1, action={"x": 1}, reward=1.0)
+    Trajectory(iteration_dir=tmp_path).append(step=1, action={"x": 1}, reward=1.0, observation={"metric": 2.0})
 
-    assert path.read_text().splitlines() == ["step,action.x,reward", "1,1,1.0"]
+    assert path.read_text().splitlines() == [
+        "step,action,reward,observation",
+        "1,{'x': 1},1.0,[2.0]",
+    ]
 
 
 def test_csv_rejects_an_existing_mismatched_header(tmp_path: Path) -> None:
     path = tmp_path / "trajectory.csv"
-    path.write_text("step,reward,action.x\n")
+    path.write_text("step,reward,action,observation\n")
 
     with pytest.raises(ValueError, match="file fields do not match"):
-        Trajectory(iteration_dir=tmp_path).append(step=1, action={"x": 1}, reward=1.0)
+        Trajectory(iteration_dir=tmp_path).append(step=1, action={"x": 1}, reward=1.0, observation={"metric": 2.0})
 
-    assert path.read_text() == "step,reward,action.x\n"
+    assert path.read_text() == "step,reward,action,observation\n"
+
+
+def test_core_only_trajectory_does_not_create_metadata_file(tmp_path: Path) -> None:
+    trajectory = Trajectory(iteration_dir=tmp_path)
+
+    trajectory.append(step=1, action={"x": 1}, reward=1.0, observation={"metric": 2.0})
+
+    assert not trajectory.metadata_path.exists()
+
+
+def test_metadata_reuses_matching_header_without_duplication(tmp_path: Path) -> None:
+    path = tmp_path / "metadata.csv"
+    path.write_text("step,env_params.speed,logging.power\n")
+    trajectory = Trajectory(iteration_dir=tmp_path)
+
+    trajectory.append(
+        step=1,
+        action={"x": 1},
+        reward=1.0,
+        observation={"metric": 2.0},
+        env_params={"speed": 100},
+        logging={"power": 600.0},
+    )
+
+    assert path.read_text().splitlines() == [
+        "step,env_params.speed,logging.power",
+        "1,100,600.0",
+    ]
+
+
+def test_metadata_initializes_a_precreated_empty_file(tmp_path: Path) -> None:
+    path = tmp_path / "metadata.csv"
+    path.touch()
+    trajectory = Trajectory(iteration_dir=tmp_path)
+
+    trajectory.append(
+        step=1,
+        action={"x": 1},
+        reward=1.0,
+        observation={"metric": 2.0},
+        env_params={"speed": 100},
+    )
+
+    assert path.read_text().splitlines() == ["step,env_params.speed", "1,100"]
+
+
+def test_metadata_rejects_an_existing_mismatched_header_before_core_write(tmp_path: Path) -> None:
+    metadata_path = tmp_path / "metadata.csv"
+    metadata_path.write_text("step,logging.power,env_params.speed\n")
+    trajectory = Trajectory(iteration_dir=tmp_path)
+
+    with pytest.raises(ValueError, match="file fields do not match"):
+        trajectory.append(
+            step=1,
+            action={"x": 1},
+            reward=1.0,
+            observation={"metric": 2.0},
+            env_params={"speed": 100},
+            logging={"power": 600.0},
+        )
+
+    assert not trajectory.output_path.exists()
+    assert len(trajectory) == 0
 
 
 def test_persistence_failure_does_not_append_to_memory(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
@@ -178,23 +325,15 @@ def fail_write(*_args: object, **_kwargs: object) -> None:
     monkeypatch.setattr(pd.DataFrame, "to_csv", fail_write)
 
     with pytest.raises(OSError, match="write failed"):
-        trajectory.append(step=1, action={"x": 1}, reward=1.0)
+        trajectory.append(step=1, action={"x": 1}, reward=1.0, observation={"metric": 2.0})
 
     assert len(trajectory) == 0
 
 
-def test_callable_iteration_directory_is_resolved_per_append(tmp_path: Path) -> None:
-    current = tmp_path / "first"
-    trajectory = Trajectory(iteration_dir=lambda: current)
-    trajectory.append(step=1, action={"x": 1})
-
-    assert trajectory.output_path == current / "trajectory.csv"
-
-
-def test_trajectory_logs_lifecycle_and_lookup(caplog: pytest.LogCaptureFixture) -> None:
+def test_trajectory_logs_lifecycle_and_lookup(tmp_path: Path, caplog: pytest.LogCaptureFixture) -> None:
     with caplog.at_level("DEBUG"):
-        trajectory = Trajectory()
-        trajectory.append(step=1, action={"x": 1}, reward=1.0)
+        trajectory = Trajectory(iteration_dir=tmp_path)
+        trajectory.append(step=1, action={"x": 1}, reward=1.0, observation={"metric": 2.0})
         assert trajectory.find(action={"x": 1}) is not None
         assert trajectory.find(action={"x": 2}) is None