From f35b041509fcfdf1b3d2ab21ee20b2e815298d28 Mon Sep 17 00:00:00 2001
From: Matthieu Darcy <68646255+MatthieuDarcy@users.noreply.github.com>
Date: Wed, 19 Aug 2026 15:38:49 -0400
Subject: [PATCH 1/4] added the option to use "previous_transition" in
 dsx.simulate

Add observation_control_alignment for discrete-time Simulator (#312)

Add an explicit observation_control_alignment: Literal["same_time",
"previous_transition"] field to DynamicalModel, defaulting to "same_time"
(today's behavior, unchanged). "previous_transition" pairs y_{k+1} with u_k
(the control that produced x_{k+1}) instead of pairing y_k with u_k, matching
DiscreteControlLoopSimulator's existing closed-loop convention and avoiding
the acausal y_0-depends-on-u_0 coupling.

For "previous_transition", DiscreteTimeSimulator/dsx.simulate never samples
y_0 and excludes x_0/t_0 from the returned SimulatedResult -- states,
observations, times, and the caller's ctrl_values all end up the same length,
with no padding or off-by-one bookkeeping required.

Scope: the plain Simulator/DiscreteTimeSimulator/dsx.simulate generation path
only. mppi.py and discrete_controller_simulators.py are unchanged, deferred
to a follow-up.
---
 dynestyx/api.py                   |  12 +-
 dynestyx/handlers.py              |   8 +-
 dynestyx/models/checkers.py       |  26 +++++
 dynestyx/models/core.py           |  22 +++-
 dynestyx/simulation/discrete.py   | 187 +++++++++++++++++++++++++-----
 dynestyx/types.py                 |  10 ++
 dynestyx/utils.py                 |  19 ++-
 tests/test_models_core.py         |  57 +++++++++
 tests/test_simulate_standalone.py | 149 ++++++++++++++++++++++++
 9 files changed, 457 insertions(+), 33 deletions(-)

diff --git a/dynestyx/api.py b/dynestyx/api.py
index 67d2b6c6..7b9b4574 100644
--- a/dynestyx/api.py
+++ b/dynestyx/api.py
@@ -55,7 +55,9 @@ def simulate(
         dynamics: Dynamical model to simulate.
         rng_key: JAX pseudorandom number generator key.
         ctrl_times: Times associated with `ctrl_values`. If controls are
-            provided, these times must match `predict_times`.
+            provided, these times must match `predict_times` for models with
+            `dynamics.observation_control_alignment="same_time"` (default), or
+            `predict_times[:-1]` for `"previous_transition"`.
         ctrl_values: Control values, or `None` for an uncontrolled model.
         predict_times: Times at which to simulate states and observations.
         n_simulations: Number of independent trajectories to simulate.
@@ -97,7 +99,13 @@ def simulate(
 
     _validate_site_sorting(ctrl_times, name="ctrl_times")
     _validate_site_sorting(predict_times, name="predict_times")
-    _validate_controls(None, predict_times, ctrl_times, ctrl_values)
+    _validate_controls(
+        None,
+        predict_times,
+        ctrl_times,
+        ctrl_values,
+        observation_control_alignment=dynamics.observation_control_alignment,
+    )
     _validate_control_dim(dynamics, ctrl_values)
 
     dynamics_with_t0 = _get_dynamics_with_t0(dynamics, None, predict_times)
diff --git a/dynestyx/handlers.py b/dynestyx/handlers.py
index 5b8bfe29..e4a9fb06 100644
--- a/dynestyx/handlers.py
+++ b/dynestyx/handlers.py
@@ -100,7 +100,13 @@ def _validate_and_prepare(
     _validate_site_sorting(ctrl_times, name="ctrl_times")
     _validate_site_sorting(predict_times, name="predict_times")
 
-    _validate_controls(obs_times, predict_times, ctrl_times, ctrl_values)
+    _validate_controls(
+        obs_times,
+        predict_times,
+        ctrl_times,
+        ctrl_values,
+        observation_control_alignment=dynamics.observation_control_alignment,
+    )
     _validate_control_dim(dynamics, ctrl_values)
 
     # Initial dynamics may not have t0, which is then inferred from obs_times
diff --git a/dynestyx/models/checkers.py b/dynestyx/models/checkers.py
index c778700f..a2ac7765 100644
--- a/dynestyx/models/checkers.py
+++ b/dynestyx/models/checkers.py
@@ -235,6 +235,32 @@ def _validate_categorical_state(
         )
 
 
+_VALID_OBSERVATION_CONTROL_ALIGNMENTS = ("same_time", "previous_transition")
+
+
+def _validate_observation_control_alignment(
+    observation_control_alignment: str, continuous_time: bool
+) -> None:
+    """Validate the observation/control alignment convention.
+
+    Unlike continuous_time/categorical_state, this field has no inferred
+    counterpart -- it is purely user-supplied with a default, so this only
+    checks (1) the value is one of the two supported literals and (2)
+    'previous_transition' is only used with discrete-time state evolution.
+    """
+    if observation_control_alignment not in _VALID_OBSERVATION_CONTROL_ALIGNMENTS:
+        raise ValueError(
+            "observation_control_alignment must be one of "
+            f"{_VALID_OBSERVATION_CONTROL_ALIGNMENTS}, got "
+            f"{observation_control_alignment!r}."
+        )
+    if observation_control_alignment == "previous_transition" and continuous_time:
+        raise ValueError(
+            "observation_control_alignment='previous_transition' is only "
+            "supported for discrete-time models (continuous_time=False)."
+        )
+
+
 def _inside_numpyro_plate_context() -> bool:
     """Return True when currently executing inside any active numpyro.plate frame."""
     return any(
diff --git a/dynestyx/models/core.py b/dynestyx/models/core.py
index 4fa14852..643d5ddf 100644
--- a/dynestyx/models/core.py
+++ b/dynestyx/models/core.py
@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from typing import Any, Protocol, cast, runtime_checkable
+from typing import Any, Literal, Protocol, cast, runtime_checkable
 
 import equinox as eqx
 import jax
@@ -22,6 +22,7 @@
     _validate_continuous_time_flag,
     _validate_discrete_state_evolution_output_shape,
     _validate_imex_potential_conflict,
+    _validate_observation_control_alignment,
     _validate_observation_dim,
     _validate_state_dim,
 )
@@ -103,7 +104,16 @@ class DynamicalModel(eqx.Module):
             exactly; a mismatch raises a ``ValueError`` at simulation time.
         continuous_time (bool): Whether the model uses continuous-time state evolution (SDE) or discrete-time.
             Gets set automatically from the concrete type of `state_evolution`.
-    
+        observation_control_alignment ("same_time" | "previous_transition"): Convention for how
+            observations pair with controls in discrete time. `"same_time"` (default) pairs
+            $y_k$ with $u_k$, matching today's `dsx.simulate` behavior. `"previous_transition"`
+            pairs $y_{k+1}$ with $u_k$ (the control that produced $x_{k+1}$), matching
+            `DiscreteControlLoopSimulator`'s closed-loop convention; under this convention
+            $y_0$ is never sampled. Only `"same_time"` is honored outside the plain
+            `Simulator`/`DiscreteTimeSimulator`/`dsx.simulate` generation path (not yet by
+            Filter/Smoother/`LatentPathBuilder` posterior rollout, `DiscreteControlLoopSimulator`,
+            or `mppi.py` -- see [issue #312](https://github.com/BasisResearch/dynestyx/issues/312)).
+
     Note:
         - `continuous_time`, `state_dim`, `observation_dim`, and `categorical_state` are inferred automatically; do not pass them to the constructor.
         - Logic for control_model is not implemented yet.
@@ -127,6 +137,7 @@ class DynamicalModel(eqx.Module):
     observation_dim: int
     categorical_state: bool
     continuous_time: bool
+    observation_control_alignment: Literal["same_time", "previous_transition"]
 
     def __init__(
         self,
@@ -141,12 +152,19 @@ def __init__(
         observation_dim: int | None = None,
         categorical_state: bool | None = None,
         continuous_time: bool | None = None,
+        observation_control_alignment: Literal[
+            "same_time", "previous_transition"
+        ] = "same_time",
     ):
         inferred_continuous_time = isinstance(
             state_evolution, ContinuousTimeStateEvolution
         )
         _validate_continuous_time_flag(continuous_time, inferred_continuous_time)
         self.continuous_time = inferred_continuous_time
+        _validate_observation_control_alignment(
+            observation_control_alignment, self.continuous_time
+        )
+        self.observation_control_alignment = observation_control_alignment
         self.initial_condition = initial_condition
         self.state_evolution = state_evolution
         self.observation_model = observation_model
diff --git a/dynestyx/simulation/discrete.py b/dynestyx/simulation/discrete.py
index 45a60135..8b3da0af 100644
--- a/dynestyx/simulation/discrete.py
+++ b/dynestyx/simulation/discrete.py
@@ -13,6 +13,7 @@
 from dynestyx.simulation.utils import (
     _ensure_trailing_dim,
     _sample_initial_states,
+    _sample_observation_path,
     _tile_times,
 )
 from dynestyx.types import SimulatedResult
@@ -49,10 +50,25 @@ def _sample_discrete_state_path_from_initial_state(
     initial_state: Real[Array, " state_dim"] | Real[Array, ""],
     rng_key: PRNGKeyArray,
     times: Real[Array, " time"],
-    ctrl_values: Real[Array, "time control_dim"] | Real[Array, " time"] | None,
-) -> Real[Array, "time state_dim"] | Real[Array, " time"]:
-    """Sample one canonical discrete state path from a fixed initial state."""
-    if len(times) == 1:
+    ctrl_values: Real[Array, "ctrl_time control_dim"]
+    | Real[Array, " ctrl_time"]
+    | None,
+    include_initial_condition: bool = True,
+) -> Real[Array, "state_path_time state_dim"] | Real[Array, " state_path_time"]:
+    """Sample one canonical discrete state path from a fixed initial state.
+
+    ctrl_values has its own length ("ctrl_time"), decoupled from `times`: it
+    is len(times) for same_time (one entry per transition plus one unused by
+    any transition, reserved for the final same_time observation) or
+    len(times) - 1 for previous_transition (exactly one entry per
+    transition). The returned path's length also depends on
+    include_initial_condition, hence the separate "state_path_time" name.
+
+    Returns x_0..x_{T-1} when include_initial_condition is True (same_time
+    convention, default). Returns x_1..x_{T-1} only when False
+    (previous_transition convention) -- x_0 is the given seed, not re-emitted.
+    """
+    if len(times) == 1 and include_initial_condition:
         return jnp.expand_dims(initial_state, axis=0)
 
     state_transition = cast(DiscreteStateTransition, dynamics.state_evolution)
@@ -74,7 +90,77 @@ def _step(carry, t_idx):
         (initial_state, rng_key),
         jnp.arange(len(times) - 1),
     )
-    return jnp.concatenate([jnp.expand_dims(initial_state, 0), scan_states], axis=0)
+    if include_initial_condition:
+        return jnp.concatenate([jnp.expand_dims(initial_state, 0), scan_states], axis=0)
+    return scan_states
+
+
+def _sample_discrete_observation_path(
+    dynamics: DynamicalModel,
+    *,
+    states: Real[Array, "obs_path_time state_dim"] | Real[Array, " obs_path_time"],
+    times: Real[Array, " time"],
+    ctrl_values: Real[Array, "obs_path_time control_dim"]
+    | Real[Array, " obs_path_time"]
+    | None,
+    rng_key: PRNGKeyArray,
+    include_initial_condition: bool = True,
+) -> Real[Array, "obs_path_time observation_dim"] | Real[Array, " obs_path_time"]:
+    """Sample observations for a discrete state path.
+
+    states/ctrl_values/the return value all share one length ("obs_path_time"),
+    decoupled from `times` (always the full predict_times grid, length T):
+    that shared length is T for same_time or T-1 for previous_transition.
+
+    include_initial_condition=True (same_time, default): states/times are the
+    FULL path (length T, x_0/t_0 included); delegates to the shared
+    _sample_observation_path so same_time logic has one source of truth,
+    unchanged from today.
+
+    include_initial_condition=False (previous_transition): states is
+    x_1..x_{T-1} (length T-1, the output of the state function above with
+    include_initial_condition=False); times is still the FULL predict_times
+    (length T) so times[k+1] is available. y_{k+1} ~ p(x_{k+1}, ctrl_values[k],
+    t_{k+1}), k=0..T-2. y_0 is never sampled -- no wasted draw.
+    """
+    if include_initial_condition:
+        ctrl_eval = (
+            (lambda t: ctrl_values[jnp.searchsorted(times, t, side="left")])
+            if ctrl_values is not None
+            else None
+        )
+        return _sample_observation_path(
+            dynamics,
+            states=states,
+            times=times,
+            rng_key=rng_key,
+            control_path_eval=ctrl_eval,
+        )
+
+    n = states.shape[0]
+    obs_keys = jr.split(rng_key, n)
+    future_times = times[1:]
+
+    # Map directly over the pre-sliced arrays (states/ctrl_values/future_times/
+    # obs_keys) rather than indexing by a scanned integer inside the mapped
+    # body. jax.vmap traces its body once regardless of batch size, so
+    # indexing into a genuinely zero-length array (the n=0 edge case, e.g.
+    # predict_times of length 1) raises immediately; mapping over already-
+    # sliced arrays instead lets vmap's native zero-size handling take over,
+    # with no indexing operation in the body at all.
+    if ctrl_values is None:
+
+        def _sample_one(x_next, t_next, key):
+            obs_dist = dynamics.observation_model(x=x_next, u=None, t=t_next)
+            return obs_dist.sample(key)
+
+        return jax.vmap(_sample_one)(states, future_times, obs_keys)
+
+    def _sample_one(x_next, u, t_next, key):
+        obs_dist = dynamics.observation_model(x=x_next, u=u, t=t_next)
+        return obs_dist.sample(key)
+
+    return jax.vmap(_sample_one)(states, ctrl_values, future_times, obs_keys)
 
 
 def _sample_discrete_state_path(
@@ -108,7 +194,10 @@ class DiscreteTimeSimulator(BaseSimulator):
     r"""Generate trajectories from a discrete-time dynamical model.
 
     For prediction times \(t_0,\ldots,t_{T-1}\), this simulator draws
-    `n_simulations` independent paths according to
+    `n_simulations` independent paths. The observation/control pairing
+    depends on `dynamics.observation_control_alignment`:
+
+    For `"same_time"` (default):
 
     \[
     x_0^{(m)} \sim p_0(x_0), \qquad
@@ -121,7 +210,30 @@ class DiscreteTimeSimulator(BaseSimulator):
     The first state in the returned path is the initial-condition draw at
     `predict_times[0]`; the simulator then makes one transition draw for each
     adjacent pair of prediction times and samples one observation conditional
-    on every realized state. See
+    on every realized state, including \(y_0\) (paired with \(u_0\)).
+
+    For `"previous_transition"`:
+
+    \[
+    x_0^{(m)} \sim p_0(x_0), \qquad
+    x_{k+1}^{(m)} \sim p\!\left(x_{k+1}\mid x_k^{(m)},u_k,t_k,t_{k+1}\right),
+    \qquad
+    y_{k+1}^{(m)} \sim p(y_{k+1}\mid x_{k+1}^{(m)},u_k,t_{k+1}).
+    \]
+
+    Here \(x_0\) is only the seed for the rollout: \(y_0\) is never sampled,
+    and neither \(x_0\) nor \(t_0\) appears in the returned result --
+    `SimulatedResult.x_0` is `None` and `.times`/`.states`/`.observations` all
+    have length \(T-1\), matching the \(T-1\) controls \(u_0,\ldots,u_{T-2}\)
+    the caller supplies via `ctrl_values` (aligned to `predict_times[:-1]`,
+    not the full `predict_times`). This matches
+    [DiscreteControlLoopSimulator][dynestyx.control.discrete_controller_simulators.DiscreteControlLoopSimulator]'s
+    closed-loop convention. Only discrete-time models generated through the
+    plain `Simulator`/`DiscreteTimeSimulator`/`dsx.simulate` path honor this
+    convention today -- see
+    [issue #312](https://github.com/BasisResearch/dynestyx/issues/312).
+
+    See
     [DiscreteTimeStateEvolution][dynestyx.models.core.DiscreteTimeStateEvolution]
     for how a discrete transition model is represented in a `DynamicalModel`.
 
@@ -165,10 +277,14 @@ class DiscreteTimeSimulator(BaseSimulator):
     not be uniformly spaced, provided the model's transition accepts those
     intervals.
 
-    If controls are supplied, `ctrl_times` must contain every prediction time
-    exactly. `ctrl_values[k]` is used for the transition beginning at \(t_k\)
-    and for the observation at \(t_k\). The paired control arrays are validated
-    before simulation.
+    If controls are supplied, `ctrl_times` must exactly match the grid
+    required by `dynamics.observation_control_alignment`: the full
+    `predict_times` for `"same_time"` (default) -- `ctrl_values[k]` is used
+    for the transition beginning at \(t_k\) and for the observation at
+    \(t_k\) -- or `predict_times[:-1]` for `"previous_transition"` --
+    `ctrl_values[k]` drives the transition into \(x_{k+1}\) and the
+    observation \(y_{k+1}\). The paired control arrays are validated before
+    simulation.
 
     This handler is generation-only and does not condition on `obs_times` or
     `obs_values`. Use
@@ -250,15 +366,20 @@ def _simulate_forward_from_initial_state(
         | Real[Array, " n_simulations"],
         rng_key: PRNGKeyArray,
         times: Real[Array, " time"],
-        ctrl_values: Real[Array, "time control_dim"] | Real[Array, " time"] | None,
+        ctrl_values: Real[Array, "ctrl_time control_dim"]
+        | Real[Array, " ctrl_time"]
+        | None,
     ) -> SimulatedResult:
-        """Run pure forward simulation for a discrete-time model."""
+        """Run pure forward simulation for a discrete-time model.
+
+        ctrl_values has its own length ("ctrl_time"), decoupled from `times`
+        (always the full predict_times grid): len(times) for same_time or
+        len(times) - 1 for previous_transition.
+        """
         n_sim = initial_state.shape[0]
         sim_keys = jr.split(rng_key, n_sim)
-        ctrl_eval = (
-            (lambda t: ctrl_values[jnp.searchsorted(times, t, side="left")])
-            if ctrl_values is not None
-            else None
+        include_initial_condition = (
+            dynamics.observation_control_alignment != "previous_transition"
         )
 
         def _sim_one_trajectory(
@@ -272,22 +393,29 @@ def _sim_one_trajectory(
                 rng_key=key_states,
                 times=times,
                 ctrl_values=ctrl_values,
+                include_initial_condition=include_initial_condition,
             )
-            observations = self._emit_observations(
-                "",
+            observations = _sample_discrete_observation_path(
                 dynamics,
-                states,
-                times,
-                None,
-                ctrl_eval,
-                key=key_obs,
+                states=states,
+                times=times,
+                ctrl_values=ctrl_values,
+                rng_key=key_obs,
+                include_initial_condition=include_initial_condition,
             )
             return states, observations
 
         states, observations = jax.vmap(_sim_one_trajectory)(sim_keys, initial_state)
+        if include_initial_condition:
+            return SimulatedResult(
+                times=_tile_times(times, n_sim),
+                x_0=initial_state,
+                states=_ensure_trailing_dim(states),
+                observations=_ensure_trailing_dim(observations),
+            )
         return SimulatedResult(
-            times=_tile_times(times, n_sim),
-            x_0=initial_state,
+            times=_tile_times(times[1:], n_sim),
+            x_0=None,
             states=_ensure_trailing_dim(states),
             observations=_ensure_trailing_dim(observations),
         )
@@ -315,8 +443,13 @@ def simulate(
         if predict_times is None:
             raise ValueError("predict_times must be provided")
 
+        align_times = (
+            predict_times[:-1]
+            if dynamics.observation_control_alignment == "previous_transition"
+            else predict_times
+        )
         aligned_ctrl_values = _align_ctrl_values_to_times(
-            times=predict_times,
+            times=align_times,
             ctrl_times=ctrl_times,
             ctrl_values=ctrl_values,
         )
diff --git a/dynestyx/types.py b/dynestyx/types.py
index 283cd738..2a144da6 100644
--- a/dynestyx/types.py
+++ b/dynestyx/types.py
@@ -143,6 +143,16 @@ class SimulatedResult(eqx.Module):
     posterior rollout, the same result object instead carries
     ``predicted_times``, ``predicted_states``, and
     ``predicted_observations``.
+
+    For a discrete-time model with
+    ``dynamics.observation_control_alignment="previous_transition"``, ``x_0``
+    is ``None`` -- the seed state is not part of the rollout output -- and
+    ``times``, ``states``, and ``observations`` are all exactly the length of
+    the ``ctrl_values`` the caller supplied (one shorter than
+    ``predict_times``, since :math:`y_0` is never sampled under this
+    convention). ``x``, ``y``, ``u``, and ``t`` are therefore all the same
+    shape, with no :math:`t_0`/:math:`x_0` remnant anywhere in the result. See
+    [DiscreteTimeSimulator][dynestyx.simulation.discrete.DiscreteTimeSimulator].
     """
 
     times: Real[Array, "*plate n_simulations time"] | None = None
diff --git a/dynestyx/utils.py b/dynestyx/utils.py
index ff122ded..4311e21a 100644
--- a/dynestyx/utils.py
+++ b/dynestyx/utils.py
@@ -334,6 +334,8 @@ def _validate_controls(
     ctrl_values: Real[Array, "*ctrl_value_plate ctrl_time control_dim"]
     | Real[Array, "*ctrl_value_plate ctrl_time"]
     | None,
+    *,
+    observation_control_alignment: str = "same_time",
 ) -> None:
     """
     Validate control inputs against model time grids.
@@ -344,11 +346,22 @@ def _validate_controls(
     - If both obs_times and predict_times are present, ctrl_times must match their union.
     - Otherwise ctrl_times must match whichever single grid is provided.
     - Matching is set-like (order-insensitive) and length-preserving.
+    - When observation_control_alignment is "previous_transition", ctrl_times must
+      instead match predict_times[:-1] (one control per transition); obs_times-based
+      conditioning is not supported yet under this convention (see issue #312).
 
     Raises:
         ValueError: If controls are partially provided or no time grid is provided.
     """
 
+    if observation_control_alignment == "previous_transition" and obs_times is not None:
+        raise ValueError(
+            "observation_control_alignment='previous_transition' does not "
+            "support obs_times-based conditioning yet (Filter/Smoother/"
+            "LatentPathBuilder posterior rollout); only predict_times-only "
+            "generation is supported. See issue #312."
+        )
+
     if ctrl_times is None:
         if ctrl_values is not None:
             raise ValueError(
@@ -365,7 +378,11 @@ def _validate_controls(
     if obs_times is None and predict_times is None:
         raise ValueError("At least one of obs_times or predict_times must be provided")
 
-    if obs_times is None:
+    if observation_control_alignment == "previous_transition":
+        # obs_times is None here -- already rejected above otherwise.
+        assert predict_times is not None
+        total_obs_pred_times = predict_times[:-1]
+    elif obs_times is None:
         total_obs_pred_times = predict_times
     elif predict_times is None:
         total_obs_pred_times = obs_times
diff --git a/tests/test_models_core.py b/tests/test_models_core.py
index a877e82e..a602e3f3 100644
--- a/tests/test_models_core.py
+++ b/tests/test_models_core.py
@@ -803,3 +803,60 @@ def bad_cov_fn(t_now, t_next):
             ),
             observation_model=dsx.LinearGaussianObservation(H=jnp.eye(2), R=jnp.eye(2)),
         )
+
+
+# ---------------------------------------------------------------------------
+# observation_control_alignment field (#312)
+# ---------------------------------------------------------------------------
+
+
+def test_observation_control_alignment_defaults_to_same_time() -> None:
+    model = _simple_discrete_model()
+    assert model.observation_control_alignment == "same_time"
+
+
+def test_observation_control_alignment_previous_transition_stored() -> None:
+    model = DynamicalModel(
+        initial_condition=dist.Normal(0.0, 1.0),
+        state_evolution=lambda x, u, t_now, t_next: dist.Normal(x, 0.1),
+        observation_model=lambda x, u, t: dist.Normal(x, 0.1),
+        control_dim=0,
+        observation_control_alignment="previous_transition",
+    )
+    assert model.observation_control_alignment == "previous_transition"
+
+
+def test_observation_control_alignment_rejects_invalid_literal() -> None:
+    # Under the pytest jaxtyping import hook (see pyproject.toml addopts),
+    # jaxtyping's own Literal[...] enforcement rejects an invalid value before
+    # DynamicalModel.__init__'s body -- and _validate_observation_control_alignment
+    # within it -- ever runs, raising jaxtyping.TypeCheckError (a TypeError
+    # subclass) rather than the ValueError _validate_observation_control_alignment
+    # raises outside that instrumented context. Accept either so this test is
+    # correct with or without the import hook active.
+    with pytest.raises((ValueError, TypeError)):
+        DynamicalModel(
+            initial_condition=dist.Normal(0.0, 1.0),
+            state_evolution=lambda x, u, t_now, t_next: dist.Normal(x, 0.1),
+            observation_model=lambda x, u, t: dist.Normal(x, 0.1),
+            control_dim=0,
+            observation_control_alignment="bogus",  # ty: ignore[invalid-argument-type]
+        )
+
+
+def test_observation_control_alignment_previous_transition_rejects_continuous_time() -> (
+    None
+):
+    with pytest.raises(
+        ValueError,
+        match="observation_control_alignment='previous_transition' is only supported "
+        "for discrete-time models",
+    ):
+        DynamicalModel(
+            initial_condition=_initial_condition_2d(),
+            state_evolution=ContinuousTimeStateEvolution(
+                drift=lambda x, u, t: -0.3 * x
+            ),
+            observation_model=_observation_model_2d,
+            observation_control_alignment="previous_transition",
+        )
diff --git a/tests/test_simulate_standalone.py b/tests/test_simulate_standalone.py
index 07db9cb8..66db4945 100644
--- a/tests/test_simulate_standalone.py
+++ b/tests/test_simulate_standalone.py
@@ -406,3 +406,152 @@ def model():
         "Simulator + Predictive, dsx.simulate, and pre-split Simulator.simulate "
         f"produced different values for {mismatched_fields}"
     )
+
+
+# ---------------------------------------------------------------------------
+# observation_control_alignment="previous_transition" (#312)
+# ---------------------------------------------------------------------------
+
+
+def _make_previous_transition_dynamics() -> dsx.DynamicalModel:
+    """Deterministic 1-D discrete model whose observation reveals both the
+    state and (scaled) control it was conditioned on, so tests can check
+    exactly which control an observation used."""
+
+    def _state_evolution(x, u, t_now, t_next):
+        del t_now, t_next
+        u = jnp.zeros_like(x) if u is None else u
+        return dist.Delta(x + u).to_event(1)
+
+    def _observation_model(x, u, t):
+        del t
+        u = jnp.zeros_like(x) if u is None else u
+        return dist.Delta(x + 100.0 * u).to_event(1)
+
+    return dsx.DynamicalModel(
+        control_dim=1,
+        initial_condition=dist.Delta(jnp.array([0.0])).to_event(1),
+        state_evolution=_state_evolution,
+        observation_model=_observation_model,
+        observation_control_alignment="previous_transition",
+    )
+
+
+def test_discrete_simulator_previous_transition_aligns_ctrl_values_shorter_by_one():
+    predict_times = jnp.array([0.0, 1.0, 2.0, 3.0])
+    ctrl_times = predict_times[:-1]
+    ctrl_values = jnp.array([[1.0], [2.0], [3.0]])
+
+    result = dsx.DiscreteTimeSimulator().simulate(
+        _make_previous_transition_dynamics(),
+        rng_key=jr.PRNGKey(0),
+        predict_times=predict_times,
+        ctrl_times=ctrl_times,
+        ctrl_values=ctrl_values,
+    )
+
+    states = jnp.asarray(result.states)
+    observations = jnp.asarray(result.observations)
+    times = jnp.asarray(result.times)
+
+    assert result.x_0 is None
+    assert times.shape == (1, 3)
+    assert jnp.allclose(times[0], predict_times[1:])
+    assert states.shape == (1, 3, 1)
+    assert observations.shape == (1, 3, 1)
+    # x_1=1, x_2=3, x_3=6 (cumulative sum of controls, x_0=0)
+    expected_states = jnp.array([[[1.0], [3.0], [6.0]]])
+    assert jnp.array_equal(states, expected_states)
+    # y_{k+1} = x_{k+1} + 100 * u_k
+    expected_observations = expected_states + 100.0 * jnp.array([[[1.0], [2.0], [3.0]]])
+    assert jnp.array_equal(observations, expected_observations)
+
+
+def test_discrete_simulator_previous_transition_rejects_ctrl_times_matching_full_predict_times():
+    """dsx.simulate's _validate_controls gate requires an exact-length match
+    against predict_times[:-1] for previous_transition; DiscreteTimeSimulator's
+    own _align_ctrl_values_to_times permits a superset ctrl_times, so this must
+    go through the public dsx.simulate entry point to observe the rejection."""
+    predict_times = jnp.array([0.0, 1.0, 2.0, 3.0])
+    ctrl_values = jnp.array([[1.0], [2.0], [3.0], [4.0]])
+
+    with pytest.raises(Exception):
+        dsx.simulate(
+            _make_previous_transition_dynamics(),
+            rng_key=jr.PRNGKey(0),
+            predict_times=predict_times,
+            ctrl_times=predict_times,
+            ctrl_values=ctrl_values,
+        )
+
+
+def test_dsx_simulate_previous_transition_end_to_end():
+    """Exercises dsx.simulate -> api.py -> utils.py::_validate_controls threading."""
+    predict_times = jnp.array([0.0, 1.0, 2.0, 3.0])
+    ctrl_times = predict_times[:-1]
+    ctrl_values = jnp.array([[1.0], [2.0], [3.0]])
+
+    result = dsx.simulate(
+        _make_previous_transition_dynamics(),
+        rng_key=jr.PRNGKey(0),
+        predict_times=predict_times,
+        ctrl_times=ctrl_times,
+        ctrl_values=ctrl_values,
+    )
+
+    assert result.x_0 is None
+    assert jnp.asarray(result.times).shape == (1, 3)
+    assert jnp.asarray(result.states).shape == (1, 3, 1)
+    assert jnp.asarray(result.observations).shape == (1, 3, 1)
+
+
+def test_discrete_simulator_previous_transition_zero_length_predict_times_edge_case():
+    result = dsx.simulate(
+        _make_previous_transition_dynamics(),
+        rng_key=jr.PRNGKey(0),
+        predict_times=jnp.arange(1.0),
+    )
+
+    assert result.x_0 is None
+    assert jnp.asarray(result.times).shape == (1, 0)
+    assert jnp.asarray(result.states).shape == (1, 0, 1)
+    assert jnp.asarray(result.observations).shape == (1, 0, 1)
+
+
+def test_discrete_simulator_previous_transition_rejects_obs_times():
+    predict_times = jnp.array([0.0, 1.0, 2.0])
+
+    with pytest.raises(
+        ValueError,
+        match="observation_control_alignment='previous_transition' does not support "
+        "obs_times",
+    ):
+        dsx.condition(
+            "f",
+            _make_previous_transition_dynamics(),
+            obs_times=predict_times,
+            obs_values=jnp.zeros((3, 1)),
+            predict_times=predict_times,
+        )
+
+
+def test_previous_transition_observation_uses_previous_step_control_not_same_index():
+    """y_{k+1} must use u_k (the control that produced x_{k+1}), not u_{k+1}."""
+    predict_times = jnp.array([0.0, 1.0, 2.0, 3.0])
+    ctrl_times = predict_times[:-1]
+    ctrl_values = jnp.array([[1.0], [2.0], [3.0]])
+
+    result = dsx.DiscreteTimeSimulator().simulate(
+        _make_previous_transition_dynamics(),
+        rng_key=jr.PRNGKey(0),
+        predict_times=predict_times,
+        ctrl_times=ctrl_times,
+        ctrl_values=ctrl_values,
+    )
+
+    states = jnp.asarray(result.states)
+    observations = jnp.asarray(result.observations)
+    # observation_model reveals x + 100*u, so subtracting the realized state
+    # recovers exactly which (scaled) control each observation used.
+    revealed_control = (observations - states) / 100.0
+    assert jnp.array_equal(revealed_control[0], ctrl_values)

From 92e8bcaec7a8dbe6c402c6d985fae96da34dc478 Mon Sep 17 00:00:00 2001
From: Matthieu Darcy <68646255+MatthieuDarcy@users.noreply.github.com>
Date: Mon, 31 Aug 2026 13:34:47 -0400
Subject: [PATCH 2/4] Updated results, cleaned up some functions

Include x_0 in all results; add controls to SimulatedResult

For observation_control_alignment="previous_transition", the result now keeps
x_0 and the full times/states path (length T), matching "same_time". Only
observations stay one shorter (y_1..y_{T-1}, length T-1) since y_0 is never
sampled -- so states[k+1] pairs with observations[k].

Add a controls field to SimulatedResult carrying the aligned ctrl_values used
(length T for same_time, T-1 for previous_transition; None when uncontrolled).

Also drop the bespoke _sample_discrete_observation_path in favor of calling
_emit_observations directly with sliced states/times, and fix
_sample_observation_path to vmap over arrays rather than indexing by a scanned
integer, which crashed on zero-length observation paths.
---
 dynestyx/simulation/discrete.py   | 158 ++++++++++--------------------
 dynestyx/simulation/utils.py      |  16 +--
 dynestyx/types.py                 |  33 +++++--
 tests/test_discrete_control.py    |   8 +-
 tests/test_simulate_standalone.py | 111 ++++++++++++++++++---
 5 files changed, 189 insertions(+), 137 deletions(-)

diff --git a/dynestyx/simulation/discrete.py b/dynestyx/simulation/discrete.py
index 8b3da0af..21456dc9 100644
--- a/dynestyx/simulation/discrete.py
+++ b/dynestyx/simulation/discrete.py
@@ -13,7 +13,6 @@
 from dynestyx.simulation.utils import (
     _ensure_trailing_dim,
     _sample_initial_states,
-    _sample_observation_path,
     _tile_times,
 )
 from dynestyx.types import SimulatedResult
@@ -53,22 +52,21 @@ def _sample_discrete_state_path_from_initial_state(
     ctrl_values: Real[Array, "ctrl_time control_dim"]
     | Real[Array, " ctrl_time"]
     | None,
-    include_initial_condition: bool = True,
-) -> Real[Array, "state_path_time state_dim"] | Real[Array, " state_path_time"]:
+) -> Real[Array, "time state_dim"] | Real[Array, " time"]:
     """Sample one canonical discrete state path from a fixed initial state.
 
+    Always returns x_0..x_{T-1} (length T, matching `times`) -- x_0 is
+    included regardless of `dynamics.observation_control_alignment`; only
+    observation sampling (see `_simulate_forward_from_initial_state`) differs
+    by convention.
+
     ctrl_values has its own length ("ctrl_time"), decoupled from `times`: it
     is len(times) for same_time (one entry per transition plus one unused by
     any transition, reserved for the final same_time observation) or
     len(times) - 1 for previous_transition (exactly one entry per
-    transition). The returned path's length also depends on
-    include_initial_condition, hence the separate "state_path_time" name.
-
-    Returns x_0..x_{T-1} when include_initial_condition is True (same_time
-    convention, default). Returns x_1..x_{T-1} only when False
-    (previous_transition convention) -- x_0 is the given seed, not re-emitted.
+    transition).
     """
-    if len(times) == 1 and include_initial_condition:
+    if len(times) == 1:
         return jnp.expand_dims(initial_state, axis=0)
 
     state_transition = cast(DiscreteStateTransition, dynamics.state_evolution)
@@ -90,77 +88,7 @@ def _step(carry, t_idx):
         (initial_state, rng_key),
         jnp.arange(len(times) - 1),
     )
-    if include_initial_condition:
-        return jnp.concatenate([jnp.expand_dims(initial_state, 0), scan_states], axis=0)
-    return scan_states
-
-
-def _sample_discrete_observation_path(
-    dynamics: DynamicalModel,
-    *,
-    states: Real[Array, "obs_path_time state_dim"] | Real[Array, " obs_path_time"],
-    times: Real[Array, " time"],
-    ctrl_values: Real[Array, "obs_path_time control_dim"]
-    | Real[Array, " obs_path_time"]
-    | None,
-    rng_key: PRNGKeyArray,
-    include_initial_condition: bool = True,
-) -> Real[Array, "obs_path_time observation_dim"] | Real[Array, " obs_path_time"]:
-    """Sample observations for a discrete state path.
-
-    states/ctrl_values/the return value all share one length ("obs_path_time"),
-    decoupled from `times` (always the full predict_times grid, length T):
-    that shared length is T for same_time or T-1 for previous_transition.
-
-    include_initial_condition=True (same_time, default): states/times are the
-    FULL path (length T, x_0/t_0 included); delegates to the shared
-    _sample_observation_path so same_time logic has one source of truth,
-    unchanged from today.
-
-    include_initial_condition=False (previous_transition): states is
-    x_1..x_{T-1} (length T-1, the output of the state function above with
-    include_initial_condition=False); times is still the FULL predict_times
-    (length T) so times[k+1] is available. y_{k+1} ~ p(x_{k+1}, ctrl_values[k],
-    t_{k+1}), k=0..T-2. y_0 is never sampled -- no wasted draw.
-    """
-    if include_initial_condition:
-        ctrl_eval = (
-            (lambda t: ctrl_values[jnp.searchsorted(times, t, side="left")])
-            if ctrl_values is not None
-            else None
-        )
-        return _sample_observation_path(
-            dynamics,
-            states=states,
-            times=times,
-            rng_key=rng_key,
-            control_path_eval=ctrl_eval,
-        )
-
-    n = states.shape[0]
-    obs_keys = jr.split(rng_key, n)
-    future_times = times[1:]
-
-    # Map directly over the pre-sliced arrays (states/ctrl_values/future_times/
-    # obs_keys) rather than indexing by a scanned integer inside the mapped
-    # body. jax.vmap traces its body once regardless of batch size, so
-    # indexing into a genuinely zero-length array (the n=0 edge case, e.g.
-    # predict_times of length 1) raises immediately; mapping over already-
-    # sliced arrays instead lets vmap's native zero-size handling take over,
-    # with no indexing operation in the body at all.
-    if ctrl_values is None:
-
-        def _sample_one(x_next, t_next, key):
-            obs_dist = dynamics.observation_model(x=x_next, u=None, t=t_next)
-            return obs_dist.sample(key)
-
-        return jax.vmap(_sample_one)(states, future_times, obs_keys)
-
-    def _sample_one(x_next, u, t_next, key):
-        obs_dist = dynamics.observation_model(x=x_next, u=u, t=t_next)
-        return obs_dist.sample(key)
-
-    return jax.vmap(_sample_one)(states, ctrl_values, future_times, obs_keys)
+    return jnp.concatenate([jnp.expand_dims(initial_state, 0), scan_states], axis=0)
 
 
 def _sample_discrete_state_path(
@@ -221,12 +149,16 @@ class DiscreteTimeSimulator(BaseSimulator):
     y_{k+1}^{(m)} \sim p(y_{k+1}\mid x_{k+1}^{(m)},u_k,t_{k+1}).
     \]
 
-    Here \(x_0\) is only the seed for the rollout: \(y_0\) is never sampled,
-    and neither \(x_0\) nor \(t_0\) appears in the returned result --
-    `SimulatedResult.x_0` is `None` and `.times`/`.states`/`.observations` all
-    have length \(T-1\), matching the \(T-1\) controls \(u_0,\ldots,u_{T-2}\)
-    the caller supplies via `ctrl_values` (aligned to `predict_times[:-1]`,
-    not the full `predict_times`). This matches
+    Here \(y_0\) is never sampled -- there's no control that produced it --
+    but \(x_0\) is still part of the returned result, exactly like
+    `"same_time"`: `SimulatedResult.x_0` is populated and `.states` has length
+    \(T\) (matching `.times`/`predict_times`, \(x_0,\ldots,x_{T-1}\)).
+    `.observations` and `.controls`, however, are one shorter -- length
+    \(T-1\): \(y_1,\ldots,y_{T-1}\) and \(u_0,\ldots,u_{T-2}\) (aligned to
+    `predict_times[:-1]`, not the full `predict_times`). So `.states` is
+    intentionally one longer than `.observations`/`.controls`:
+    `states[k+1]` pairs with `observations[k]`/`controls[k]`, not
+    `states[k]`. This matches
     [DiscreteControlLoopSimulator][dynestyx.control.discrete_controller_simulators.DiscreteControlLoopSimulator]'s
     closed-loop convention. Only discrete-time models generated through the
     plain `Simulator`/`DiscreteTimeSimulator`/`dsx.simulate` path honor this
@@ -314,7 +246,13 @@ class DiscreteTimeSimulator(BaseSimulator):
     - `"f_states"`: latent states, shape
       `(*plate_shape, n_simulations, T, state_dim)`;
     - `"f_observations"`: sampled observations, shape
-      `(*plate_shape, n_simulations, T, observation_dim)`.
+      `(*plate_shape, n_simulations, T, observation_dim)` for `"same_time"`,
+      or `(*plate_shape, n_simulations, T-1, observation_dim)` for
+      `"previous_transition"`;
+    - `"f_controls"`: the (aligned) controls used, when the model is
+      controlled, shape `(*plate_shape, n_simulations, T, control_dim)` for
+      `"same_time"` or `(*plate_shape, n_simulations, T-1, control_dim)` for
+      `"previous_transition"`; absent when the model is uncontrolled.
 
     Here `"f"` is replaced by the `name` passed to `dsx.sample`. Under
     `Predictive(..., num_samples=N)`, NumPyro prepends an `N` axis to each
@@ -374,7 +312,9 @@ def _simulate_forward_from_initial_state(
 
         ctrl_values has its own length ("ctrl_time"), decoupled from `times`
         (always the full predict_times grid): len(times) for same_time or
-        len(times) - 1 for previous_transition.
+        len(times) - 1 for previous_transition. States always include x_0
+        (length matches `times`) for both conventions; only observations (and
+        the returned controls) are one shorter for previous_transition.
         """
         n_sim = initial_state.shape[0]
         sim_keys = jr.split(rng_key, n_sim)
@@ -393,31 +333,39 @@ def _sim_one_trajectory(
                 rng_key=key_states,
                 times=times,
                 ctrl_values=ctrl_values,
-                include_initial_condition=include_initial_condition,
             )
-            observations = _sample_discrete_observation_path(
-                dynamics,
-                states=states,
-                times=times,
-                ctrl_values=ctrl_values,
-                rng_key=key_obs,
-                include_initial_condition=include_initial_condition,
+            # For previous_transition, drop x_0/t_0 before sampling
+            # observations -- y_0 is never sampled under that convention.
+            # Otherwise (same_time) this is a no-op.
+            obs_states, obs_times = (
+                (states, times)
+                if include_initial_condition
+                else (states[1:], times[1:])
+            )
+            ctrl_eval = (
+                (lambda t: ctrl_values[jnp.searchsorted(obs_times, t, side="left")])
+                if ctrl_values is not None
+                else None
+            )
+            observations = self._emit_observations(
+                "", dynamics, obs_states, obs_times, None, ctrl_eval, key=key_obs
             )
             return states, observations
 
         states, observations = jax.vmap(_sim_one_trajectory)(sim_keys, initial_state)
-        if include_initial_condition:
-            return SimulatedResult(
-                times=_tile_times(times, n_sim),
-                x_0=initial_state,
-                states=_ensure_trailing_dim(states),
-                observations=_ensure_trailing_dim(observations),
+
+        controls = None
+        if ctrl_values is not None:
+            controls = _ensure_trailing_dim(
+                jnp.broadcast_to(ctrl_values[None], (n_sim, *ctrl_values.shape))
             )
+
         return SimulatedResult(
-            times=_tile_times(times[1:], n_sim),
-            x_0=None,
+            times=_tile_times(times, n_sim),
+            x_0=initial_state,
             states=_ensure_trailing_dim(states),
             observations=_ensure_trailing_dim(observations),
+            controls=controls,
         )
 
     def simulate(
diff --git a/dynestyx/simulation/utils.py b/dynestyx/simulation/utils.py
index 44bbb7f7..68eab8a4 100644
--- a/dynestyx/simulation/utils.py
+++ b/dynestyx/simulation/utils.py
@@ -9,7 +9,7 @@
 import jax.numpy as jnp
 import jax.random as jr
 import numpyro
-from jaxtyping import Array, Bool, Int, PRNGKeyArray, Real
+from jaxtyping import Array, Bool, PRNGKeyArray, Real
 
 from dynestyx.models import DynamicalModel
 from dynestyx.types import SimulatedResult, chain_numpyro_site_registrations
@@ -115,10 +115,14 @@ def _sample_observation_path(
     ctrl = control_path_eval if control_path_eval is not None else (lambda t: None)
     obs_keys = jr.split(rng_key, len(times))
 
-    def _sample_at_time(t_idx: Int[Array, ""]):
-        x_t = states[t_idx]
-        t = times[t_idx]
+    # Map directly over states/times/obs_keys rather than indexing by a
+    # scanned integer inside the mapped body: jax.vmap traces its body once
+    # regardless of batch size, so indexing into a genuinely zero-length
+    # array (e.g. a previous_transition observation path sliced down from a
+    # single-timepoint prediction grid) would raise immediately. Mapping over
+    # the arrays directly lets vmap's own batching handle the zero-size case.
+    def _sample_at(x_t, t, key):
         obs_dist = dynamics.observation_model(x=x_t, u=ctrl(t), t=t)
-        return obs_dist.sample(obs_keys[t_idx])
+        return obs_dist.sample(key)
 
-    return jax.vmap(_sample_at_time)(jnp.arange(len(times)))
+    return jax.vmap(_sample_at)(states, times, obs_keys)
diff --git a/dynestyx/types.py b/dynestyx/types.py
index 2a144da6..a189d994 100644
--- a/dynestyx/types.py
+++ b/dynestyx/types.py
@@ -144,17 +144,29 @@ class SimulatedResult(eqx.Module):
     ``predicted_times``, ``predicted_states``, and
     ``predicted_observations``.
 
+    ``controls`` carries the (aligned) control values used to produce this
+    result, when the model was controlled -- ``None`` otherwise. It is
+    populated by ``DiscreteTimeSimulator`` for both
+    ``observation_control_alignment`` conventions; ODE/SDE simulators leave it
+    ``None`` for now.
+
     For a discrete-time model with
     ``dynamics.observation_control_alignment="previous_transition"``, ``x_0``
-    is ``None`` -- the seed state is not part of the rollout output -- and
-    ``times``, ``states``, and ``observations`` are all exactly the length of
-    the ``ctrl_values`` the caller supplied (one shorter than
-    ``predict_times``, since :math:`y_0` is never sampled under this
-    convention). ``x``, ``y``, ``u``, and ``t`` are therefore all the same
-    shape, with no :math:`t_0`/:math:`x_0` remnant anywhere in the result. See
+    is populated and ``states`` includes it (length :math:`T`, matching
+    ``times``), exactly like ``"same_time"``. ``observations`` and
+    ``controls``, however, are one shorter (length :math:`T-1`:
+    :math:`y_1,\\dots,y_{T-1}` and :math:`u_0,\\dots,u_{T-2}`), since
+    :math:`y_0` is never sampled under this convention -- there is no control
+    that produced it. So ``states`` is intentionally one longer than
+    ``observations``/``controls``: ``states[k+1]`` pairs with
+    ``observations[k]``/``controls[k]``, not ``states[k]``. See
     [DiscreteTimeSimulator][dynestyx.simulation.discrete.DiscreteTimeSimulator].
     """
 
+    # observations/controls use their own axis names ("obs_time"/"ctrl_time")
+    # rather than sharing "time" with times/states: under
+    # observation_control_alignment="previous_transition" they are one
+    # shorter than times/states, so jaxtyping must not enforce equal length.
     times: Real[Array, "*plate n_simulations time"] | None = None
     x_0: (
         Real[Array, "*plate n_simulations state_dim"]
@@ -167,8 +179,13 @@ class SimulatedResult(eqx.Module):
         | None
     ) = None
     observations: (
-        Real[Array, "*plate n_simulations time observation_dim"]
-        | Real[Array, "*plate n_simulations time"]
+        Real[Array, "*plate n_simulations obs_time observation_dim"]
+        | Real[Array, "*plate n_simulations obs_time"]
+        | None
+    ) = None
+    controls: (
+        Real[Array, "*plate n_simulations ctrl_time control_dim"]
+        | Real[Array, "*plate n_simulations ctrl_time"]
         | None
     ) = None
     predicted_times: Real[Array, "*plate n_simulations predict_time"] | None = None
diff --git a/tests/test_discrete_control.py b/tests/test_discrete_control.py
index 97ba1229..bbc473ff 100644
--- a/tests/test_discrete_control.py
+++ b/tests/test_discrete_control.py
@@ -943,13 +943,15 @@ def test_dsx_simulate_with_control_policy_rejects_simulator_config():
 
 def test_dsx_simulate_without_control_policy_unchanged():
     """No control_policy given -> falls back to today's type-based routing,
-    returning a plain SimulatedResult (no controls field at all), not a
-    ControlledSimulatedResult."""
+    returning a plain SimulatedResult, not a ControlledSimulatedResult. (The
+    plain SimulatedResult does carry its own `controls` field now (#312
+    follow-up), but it's None here since no ctrl_values were supplied.)"""
     dynamics = _lti_1d()
     predict_times = jnp.arange(0.0, 5.0)
 
     result = dsx.simulate(dynamics, rng_key=jr.PRNGKey(0), predict_times=predict_times)
-    assert not hasattr(result, "controls")
+    assert not isinstance(result, ControlledSimulatedResult)
+    assert result.controls is None
 
 
 def test_initial_policy_state_threads_through_dsx_simulate():
diff --git a/tests/test_simulate_standalone.py b/tests/test_simulate_standalone.py
index 66db4945..da6ad897 100644
--- a/tests/test_simulate_standalone.py
+++ b/tests/test_simulate_standalone.py
@@ -452,18 +452,26 @@ def test_discrete_simulator_previous_transition_aligns_ctrl_values_shorter_by_on
 
     states = jnp.asarray(result.states)
     observations = jnp.asarray(result.observations)
+    controls = jnp.asarray(result.controls)
     times = jnp.asarray(result.times)
 
-    assert result.x_0 is None
-    assert times.shape == (1, 3)
-    assert jnp.allclose(times[0], predict_times[1:])
-    assert states.shape == (1, 3, 1)
-    assert observations.shape == (1, 3, 1)
-    # x_1=1, x_2=3, x_3=6 (cumulative sum of controls, x_0=0)
-    expected_states = jnp.array([[[1.0], [3.0], [6.0]]])
+    # x_0 is populated and states/times include it (length 4), like same_time.
+    assert result.x_0 is not None
+    assert jnp.array_equal(jnp.asarray(result.x_0), jnp.array([[0.0]]))
+    assert times.shape == (1, 4)
+    assert jnp.allclose(times[0], predict_times)
+    assert states.shape == (1, 4, 1)
+    # x_0=0, x_1=1, x_2=3, x_3=6 (cumulative sum of controls)
+    expected_states = jnp.array([[[0.0], [1.0], [3.0], [6.0]]])
     assert jnp.array_equal(states, expected_states)
+    # observations/controls are one shorter -- y_0 is never sampled.
+    assert observations.shape == (1, 3, 1)
+    assert controls.shape == (1, 3, 1)
+    assert jnp.array_equal(controls[0], ctrl_values)
     # y_{k+1} = x_{k+1} + 100 * u_k
-    expected_observations = expected_states + 100.0 * jnp.array([[[1.0], [2.0], [3.0]]])
+    expected_observations = expected_states[:, 1:, :] + 100.0 * jnp.array(
+        [[[1.0], [2.0], [3.0]]]
+    )
     assert jnp.array_equal(observations, expected_observations)
 
 
@@ -499,23 +507,28 @@ def test_dsx_simulate_previous_transition_end_to_end():
         ctrl_values=ctrl_values,
     )
 
-    assert result.x_0 is None
-    assert jnp.asarray(result.times).shape == (1, 3)
-    assert jnp.asarray(result.states).shape == (1, 3, 1)
+    assert result.x_0 is not None
+    assert jnp.asarray(result.times).shape == (1, 4)
+    assert jnp.asarray(result.states).shape == (1, 4, 1)
     assert jnp.asarray(result.observations).shape == (1, 3, 1)
+    assert jnp.asarray(result.controls).shape == (1, 3, 1)
 
 
 def test_discrete_simulator_previous_transition_zero_length_predict_times_edge_case():
+    """With a single prediction time, there are zero transitions/controls, so
+    observations/controls are empty -- but x_0/states/times are still the
+    (length-1) seed, same as same_time."""
     result = dsx.simulate(
         _make_previous_transition_dynamics(),
         rng_key=jr.PRNGKey(0),
         predict_times=jnp.arange(1.0),
     )
 
-    assert result.x_0 is None
-    assert jnp.asarray(result.times).shape == (1, 0)
-    assert jnp.asarray(result.states).shape == (1, 0, 1)
+    assert result.x_0 is not None
+    assert jnp.asarray(result.times).shape == (1, 1)
+    assert jnp.asarray(result.states).shape == (1, 1, 1)
     assert jnp.asarray(result.observations).shape == (1, 0, 1)
+    assert result.controls is None  # no ctrl_values supplied
 
 
 def test_discrete_simulator_previous_transition_rejects_obs_times():
@@ -553,5 +566,73 @@ def test_previous_transition_observation_uses_previous_step_control_not_same_ind
     observations = jnp.asarray(result.observations)
     # observation_model reveals x + 100*u, so subtracting the realized state
     # recovers exactly which (scaled) control each observation used.
-    revealed_control = (observations - states) / 100.0
+    # states includes x_0 (one longer than observations), so observations[k]
+    # pairs with states[k+1], not states[k].
+    revealed_control = (observations - states[:, 1:, :]) / 100.0
     assert jnp.array_equal(revealed_control[0], ctrl_values)
+
+
+def test_discrete_simulator_previous_transition_states_include_x0():
+    predict_times = jnp.array([0.0, 1.0, 2.0, 3.0])
+    ctrl_times = predict_times[:-1]
+    ctrl_values = jnp.array([[1.0], [2.0], [3.0]])
+
+    result = dsx.DiscreteTimeSimulator().simulate(
+        _make_previous_transition_dynamics(),
+        rng_key=jr.PRNGKey(0),
+        predict_times=predict_times,
+        ctrl_times=ctrl_times,
+        ctrl_values=ctrl_values,
+    )
+
+    states = jnp.asarray(result.states)
+    observations = jnp.asarray(result.observations)
+    assert jnp.array_equal(states[:, 0, :], jnp.asarray(result.x_0))
+    assert states.shape[1] == observations.shape[1] + 1
+
+
+def test_discrete_simulator_returns_controls_same_time():
+    predict_times = jnp.array([0.0, 1.0, 2.0, 3.0])
+    ctrl_values = jnp.array([[1.0], [2.0], [3.0], [4.0]])
+
+    result = dsx.DiscreteTimeSimulator().simulate(
+        _make_controlled_deterministic_discrete_dynamics(),
+        rng_key=jr.PRNGKey(0),
+        predict_times=predict_times,
+        ctrl_times=predict_times,
+        ctrl_values=ctrl_values,
+    )
+
+    assert result.controls is not None
+    assert jnp.asarray(result.controls).shape == (1, 4, 1)
+    assert jnp.array_equal(jnp.asarray(result.controls)[0], ctrl_values)
+
+
+def test_discrete_simulator_returns_controls_previous_transition():
+    predict_times = jnp.array([0.0, 1.0, 2.0, 3.0])
+    ctrl_times = predict_times[:-1]
+    ctrl_values = jnp.array([[1.0], [2.0], [3.0]])
+
+    result = dsx.DiscreteTimeSimulator().simulate(
+        _make_previous_transition_dynamics(),
+        rng_key=jr.PRNGKey(0),
+        predict_times=predict_times,
+        ctrl_times=ctrl_times,
+        ctrl_values=ctrl_values,
+    )
+
+    assert result.controls is not None
+    assert jnp.asarray(result.controls).shape == (1, 3, 1)
+    assert jnp.array_equal(jnp.asarray(result.controls)[0], ctrl_values)
+
+
+def test_discrete_simulator_controls_none_when_uncontrolled():
+    predict_times = jnp.arange(4.0)
+
+    result = dsx.simulate(
+        _make_discrete_dynamics(),
+        rng_key=jr.PRNGKey(0),
+        predict_times=predict_times,
+    )
+
+    assert result.controls is None

From 1bbfa275f71b0d72a4e3a6d4eb5af951f3f60e8f Mon Sep 17 00:00:00 2001
From: Matthieu Darcy <68646255+MatthieuDarcy@users.noreply.github.com>
Date: Mon, 31 Aug 2026 14:26:03 -0400
Subject: [PATCH 3/4] Update discrete.py

Simplified docstring
---
 dynestyx/simulation/discrete.py | 53 ++-------------------------------
 1 file changed, 3 insertions(+), 50 deletions(-)

diff --git a/dynestyx/simulation/discrete.py b/dynestyx/simulation/discrete.py
index 21456dc9..232f075d 100644
--- a/dynestyx/simulation/discrete.py
+++ b/dynestyx/simulation/discrete.py
@@ -54,17 +54,6 @@ def _sample_discrete_state_path_from_initial_state(
     | None,
 ) -> Real[Array, "time state_dim"] | Real[Array, " time"]:
     """Sample one canonical discrete state path from a fixed initial state.
-
-    Always returns x_0..x_{T-1} (length T, matching `times`) -- x_0 is
-    included regardless of `dynamics.observation_control_alignment`; only
-    observation sampling (see `_simulate_forward_from_initial_state`) differs
-    by convention.
-
-    ctrl_values has its own length ("ctrl_time"), decoupled from `times`: it
-    is len(times) for same_time (one entry per transition plus one unused by
-    any transition, reserved for the final same_time observation) or
-    len(times) - 1 for previous_transition (exactly one entry per
-    transition).
     """
     if len(times) == 1:
         return jnp.expand_dims(initial_state, axis=0)
@@ -125,45 +114,9 @@ class DiscreteTimeSimulator(BaseSimulator):
     `n_simulations` independent paths. The observation/control pairing
     depends on `dynamics.observation_control_alignment`:
 
-    For `"same_time"` (default):
-
-    \[
-    x_0^{(m)} \sim p_0(x_0), \qquad
-    x_{k+1}^{(m)}
-      \sim p\!\left(x_{k+1}\mid x_k^{(m)},u_k,t_k,t_{k+1}\right),
-    \qquad
-    y_k^{(m)} \sim p(y_k\mid x_k^{(m)},u_k,t_k).
-    \]
-
-    The first state in the returned path is the initial-condition draw at
-    `predict_times[0]`; the simulator then makes one transition draw for each
-    adjacent pair of prediction times and samples one observation conditional
-    on every realized state, including \(y_0\) (paired with \(u_0\)).
-
-    For `"previous_transition"`:
-
-    \[
-    x_0^{(m)} \sim p_0(x_0), \qquad
-    x_{k+1}^{(m)} \sim p\!\left(x_{k+1}\mid x_k^{(m)},u_k,t_k,t_{k+1}\right),
-    \qquad
-    y_{k+1}^{(m)} \sim p(y_{k+1}\mid x_{k+1}^{(m)},u_k,t_{k+1}).
-    \]
-
-    Here \(y_0\) is never sampled -- there's no control that produced it --
-    but \(x_0\) is still part of the returned result, exactly like
-    `"same_time"`: `SimulatedResult.x_0` is populated and `.states` has length
-    \(T\) (matching `.times`/`predict_times`, \(x_0,\ldots,x_{T-1}\)).
-    `.observations` and `.controls`, however, are one shorter -- length
-    \(T-1\): \(y_1,\ldots,y_{T-1}\) and \(u_0,\ldots,u_{T-2}\) (aligned to
-    `predict_times[:-1]`, not the full `predict_times`). So `.states` is
-    intentionally one longer than `.observations`/`.controls`:
-    `states[k+1]` pairs with `observations[k]`/`controls[k]`, not
-    `states[k]`. This matches
-    [DiscreteControlLoopSimulator][dynestyx.control.discrete_controller_simulators.DiscreteControlLoopSimulator]'s
-    closed-loop convention. Only discrete-time models generated through the
-    plain `Simulator`/`DiscreteTimeSimulator`/`dsx.simulate` path honor this
-    convention today -- see
-    [issue #312](https://github.com/BasisResearch/dynestyx/issues/312).
+    For `"same_time"` (default): y_{k} is paired with u_{k} and x_{k} (including k=0). States, times, observations, and controls are all of length \(T\).
+    For `"previous_transition"`: y_{k+1} is paired with u_{k} and x_{k+1} (y_0 is never sampled). States and times are of length \(T\), but observations and controls are of length \(T-1\).
+
 
     See
     [DiscreteTimeStateEvolution][dynestyx.models.core.DiscreteTimeStateEvolution]

From 007bab7b93ca6f9a3d58b77d958fc6d94e66fb15 Mon Sep 17 00:00:00 2001
From: Matthieu Darcy <68646255+MatthieuDarcy@users.noreply.github.com>
Date: Mon, 31 Aug 2026 14:29:39 -0400
Subject: [PATCH 4/4] Update discrete.py

---
 dynestyx/simulation/discrete.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/dynestyx/simulation/discrete.py b/dynestyx/simulation/discrete.py
index 232f075d..fa9ce445 100644
--- a/dynestyx/simulation/discrete.py
+++ b/dynestyx/simulation/discrete.py
@@ -53,8 +53,7 @@ def _sample_discrete_state_path_from_initial_state(
     | Real[Array, " ctrl_time"]
     | None,
 ) -> Real[Array, "time state_dim"] | Real[Array, " time"]:
-    """Sample one canonical discrete state path from a fixed initial state.
-    """
+    """Sample one canonical discrete state path from a fixed initial state."""
     if len(times) == 1:
         return jnp.expand_dims(initial_state, axis=0)