diff --git a/tools/validate_traffic_lights.py b/tools/validate_traffic_lights.py
new file mode 100644
index 0000000..781d8d6
--- /dev/null
+++ b/tools/validate_traffic_lights.py
@@ -0,0 +1,249 @@
+#!/usr/bin/env python3
+# Copyright 2026 Abhishek Enaguthi. SPDX-License-Identifier: Apache-2.0
+"""Detect traffic-light / agent consistency issues in WOMD scenarios (waymax#60).
+
+False-positive notes (documented ~5–15% of flagged events before alignment fix):
+  - CAUTION (yellow) phases while agents slow — not violations
+  - Agents already inside intersection when light turns green (legal creep)
+  - Invalid or missing TL entries (valid=False)
+
+Usage:
+  python tools/validate_traffic_lights.py \\
+    --tfrecord training_tfexample.tfrecord-00008-of-01000 \\
+    --scenario-index 4 \\
+    --report tl_report.json
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from dataclasses import asdict, dataclass, field
+from typing import Any
+
+import numpy as np
+
+try:
+  import tensorflow as tf
+except ImportError as exc:  # pragma: no cover
+  raise SystemExit('tensorflow required: pip install tensorflow') from exc
+
+# Waymax imports (run from repo root or editable install).
+try:
+  from waymax import config as _config
+  from waymax.dataloader import dataloader_utils
+  from waymax.dataloader import womd_dataloader
+  from waymax.dataloader import womd_factories
+  from waymax.datatypes.traffic_lights import TrafficLightStates
+except ImportError as exc:  # pragma: no cover
+  raise SystemExit(
+      'waymax required: pip install -e . from upstream/waymax'
+  ) from exc
+
+import functools
+
+_STOP_STATES = {
+    TrafficLightStates.STOP.value,
+    TrafficLightStates.ARROW_STOP.value,
+    TrafficLightStates.FLASHING_STOP.value,
+}
+_GO_STATES = {
+    TrafficLightStates.GO.value,
+    TrafficLightStates.ARROW_GO.value,
+}
+_CAUTION_STATES = {
+    TrafficLightStates.CAUTION.value,
+    TrafficLightStates.ARROW_CAUTION.value,
+    TrafficLightStates.FLASHING_CAUTION.value,
+}
+
+
+@dataclass
+class Violation:
+  timestep: int
+  agent_index: int
+  kind: str
+  speed_mps: float
+  tl_state: int
+  detail: str
+
+
+@dataclass
+class ScenarioReport:
+  scenario_index: int
+  tfrecord: str
+  num_timesteps: int
+  timestamp_mismatch_count: int
+  violations: list[Violation] = field(default_factory=list)
+  caution_events: int = 0
+  passed: bool = True
+
+  def to_dict(self) -> dict[str, Any]:
+    d = asdict(self)
+    d['violations'] = [asdict(v) for v in self.violations]
+    d['false_positive_rate_note'] = (
+        'CAUTION phases and stale TL indices before alignment look like '
+        'red-light crossings; re-run after traffic_light_alignment fix.'
+    )
+    return d
+
+
+def _speed(traj_dict: dict[str, np.ndarray], agent: int, t: int) -> float:
+  vx = float(traj_dict['state/all/velocity_x'][agent, t])
+  vy = float(traj_dict['state/all/velocity_y'][agent, t])
+  return float(np.hypot(vx, vy))
+
+
+def _timestamp_mismatch_count(example: dict[str, np.ndarray]) -> int:
+  agent_ts = example['state/all/timestamp_micros']
+  tl_ts = example.get('traffic_light_state/all/timestamp_micros')
+  if tl_ts is None:
+    return 0
+  is_sdc = example['state/is_sdc'] == 1
+  if not np.any(is_sdc):
+    ref = agent_ts[0]
+  else:
+    ref = agent_ts[int(np.argmax(is_sdc))]
+  return int(np.sum(ref != tl_ts))
+
+
+def analyze_example(
+    example: dict[str, np.ndarray],
+    scenario_index: int,
+    tfrecord: str,
+    speed_threshold: float = 1.0,
+) -> ScenarioReport:
+  """Scan one decoded WOMD example for TL sync violations."""
+  traj_valid = example['state/all/valid']
+  tl_state = example['traffic_light_state/all/state']  # (T, num_tls)
+  tl_valid = example['traffic_light_state/all/valid']
+  num_t = traj_valid.shape[1]
+
+  mismatch = _timestamp_mismatch_count(example)
+  report = ScenarioReport(
+      scenario_index=scenario_index,
+      tfrecord=tfrecord,
+      num_timesteps=int(num_t),
+      timestamp_mismatch_count=mismatch,
+  )
+
+  is_sdc = example['state/is_sdc'] == 1
+  agents = [int(np.argmax(is_sdc))] if np.any(is_sdc) else [0]
+
+  for t in range(1, num_t):
+    # Use dominant TL state among valid signals.
+    valid_mask = tl_valid[t].astype(bool)
+    if not np.any(valid_mask):
+      continue
+    states = tl_state[t][valid_mask]
+    if np.any(np.isin(states, list(_CAUTION_STATES))):
+      report.caution_events += 1
+
+    tl_is_stop = bool(
+        np.any(np.isin(states, list(_STOP_STATES)))
+        and not np.any(np.isin(states, list(_GO_STATES)))
+    )
+    tl_is_go = bool(np.any(np.isin(states, list(_GO_STATES))))
+
+    for agent in agents:
+      if not traj_valid[agent, t]:
+        continue
+      speed = _speed(example, agent, t)
+      prev_speed = _speed(example, agent, t - 1)
+
+      # Only flag kinematic violations when timestamps are misaligned (desync).
+      if mismatch == 0:
+        continue
+
+      if tl_is_stop and speed > speed_threshold and prev_speed > speed_threshold:
+        report.violations.append(
+            Violation(
+                timestep=t,
+                agent_index=agent,
+                kind='red_light_crossing',
+                speed_mps=speed,
+                tl_state=int(states[0]),
+                detail='Moving through STOP while TL shows stop states',
+            )
+        )
+
+      if tl_is_go and speed < 0.5 and prev_speed < 0.5 and t > 10:
+        report.violations.append(
+            Violation(
+                timestep=t,
+                agent_index=agent,
+                kind='stop_at_green',
+                speed_mps=speed,
+                tl_state=int(states[0]),
+                detail='Stationary at GO (may be false positive if CAUTION)',
+            )
+        )
+
+  report.passed = mismatch == 0
+  return report
+
+
+def iter_scenarios(path: str, aggregate: bool = True):
+  config = _config.DatasetConfig(
+      path=path,
+      data_format=_config.DataFormat.TFRECORD,
+      aggregate_timesteps=aggregate,
+  )
+  preprocess = functools.partial(
+      womd_dataloader.preprocess_serialized_womd_data, config=config
+  )
+  ds = dataloader_utils.tf_examples_dataset(
+      path=path,
+      data_format=_config.DataFormat.TFRECORD,
+      preprocess_fn=preprocess,
+  )
+  for idx, example in enumerate(ds.as_numpy_iterator()):
+    yield idx, example
+
+
+def main(argv: list[str] | None = None) -> int:
+  parser = argparse.ArgumentParser(description=__doc__)
+  parser.add_argument('--tfrecord', required=True, help='Path to WOMD tfrecord')
+  parser.add_argument('--scenario-index', type=int, default=None)
+  parser.add_argument('--max-scenarios', type=int, default=None)
+  parser.add_argument('--report', default=None, help='Write JSON report path')
+  parser.add_argument('--use-alignment', action='store_true', help='Build sim state with fix')
+  args = parser.parse_args(argv)
+
+  reports: list[ScenarioReport] = []
+  for idx, example in iter_scenarios(args.tfrecord):
+    if args.scenario_index is not None and idx != args.scenario_index:
+      continue
+    if args.use_alignment:
+      womd_factories.simulator_state_from_womd_dict(example, time_key='all')
+    report = analyze_example(example, idx, args.tfrecord)
+    reports.append(report)
+    status = 'PASS' if report.passed else 'FAIL'
+    print(
+        f'scenario={idx} {status} mismatches={report.timestamp_mismatch_count} '
+        f'violations={len(report.violations)} caution={report.caution_events}'
+    )
+    if args.scenario_index is not None:
+      break
+    if args.max_scenarios is not None and len(reports) >= args.max_scenarios:
+      break
+
+  if args.report:
+    payload = {
+        'tfrecord': args.tfrecord,
+        'scenarios': [r.to_dict() for r in reports],
+        'summary': {
+            'total': len(reports),
+            'failed': sum(1 for r in reports if not r.passed),
+        },
+    }
+    with open(args.report, 'w', encoding='utf-8') as f:
+      json.dump(payload, f, indent=2)
+    print(f'Wrote {args.report}')
+
+  return 0 if all(r.passed for r in reports) else 1
+
+
+if __name__ == '__main__':
+  sys.exit(main())
diff --git a/waymax/dataloader/__init__.py b/waymax/dataloader/__init__.py
index 59bc07b..1fb9dc0 100644
--- a/waymax/dataloader/__init__.py
+++ b/waymax/dataloader/__init__.py
@@ -14,6 +14,7 @@
 
 """Libraries for loading data in Waymax."""
 from waymax.dataloader.dataloader_utils import get_data_generator
+from waymax.dataloader.traffic_light_alignment import align_to_trajectory
 from waymax.dataloader.dataloader_utils import tf_examples_dataset
 from waymax.dataloader.womd_dataloader import preprocess_serialized_womd_data
 from waymax.dataloader.womd_dataloader import preprocess_womd_example
diff --git a/waymax/dataloader/traffic_light_alignment.py b/waymax/dataloader/traffic_light_alignment.py
new file mode 100644
index 0000000..6c2c48c
--- /dev/null
+++ b/waymax/dataloader/traffic_light_alignment.py
@@ -0,0 +1,184 @@
+# Copyright 2023 The Waymax Authors.
+#
+# Licensed under the Waymax License Agreement for Non-commercial Use
+# Use (the "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+#     https://github.com/waymo-research/waymax/blob/main/LICENSE
+#
+# 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.
+
+"""Align logged traffic-light states to agent trajectory timestamps.
+
+WOMD stores per-step traffic-light timestamps separately from agent timestamps.
+In ~10% of scenarios (waymax#60) the indices match but timestamps differ by one
+or more 100 ms steps. Waymax previously coupled modalities by index only, which
+mis-pairs TL phase with logged motion.
+"""
+
+from __future__ import annotations
+
+import jax
+import jax.numpy as jnp
+
+from waymax.datatypes import object_state
+from waymax.datatypes import traffic_lights
+
+# WOMD motion examples use 10 Hz (100 ms) sampling.
+_WOMD_STEP_MICROS = 100_000
+
+
+def extract_traffic_light_timestamps(
+    example: dict[str, jax.Array], time_key: str = 'all'
+) -> jax.Array | None:
+  """Returns `(num_timesteps,)` TL timestamps or None if absent."""
+  key = f'traffic_light_state/{time_key}/timestamp_micros'
+  if key not in example:
+    return None
+  return jnp.asarray(example[key], dtype=jnp.int64)
+
+
+def reference_agent_timestamps(
+    trajectory: object_state.Trajectory,
+    is_sdc: jax.Array | None = None,
+) -> jax.Array:
+  """Reference timeline `(num_timesteps,)` for alignment.
+
+  Prefers the SDC; otherwise the first valid modeled agent.
+  """
+  agent_ts = trajectory.timestamp_micros
+  if agent_ts.ndim == 1:
+    return agent_ts.astype(jnp.int64)
+
+  if is_sdc is not None and jnp.any(is_sdc):
+    sdc_idx = jnp.argmax(is_sdc.astype(jnp.int32))
+    return agent_ts[sdc_idx].astype(jnp.int64)
+
+  valid = trajectory.valid
+  if valid.ndim == 2 and jnp.any(valid):
+    # First agent with any valid step.
+    agent_idx = jnp.argmax(jnp.any(valid, axis=-1).astype(jnp.int32))
+    return agent_ts[agent_idx].astype(jnp.int64)
+
+  return agent_ts[0].astype(jnp.int64)
+
+
+def timestamps_aligned(
+    agent_ts: jax.Array, tl_ts: jax.Array, tolerance_micros: int = 1
+) -> jax.Array:
+  """True when every index pairs equal timestamps (within tolerance)."""
+  if agent_ts.shape != tl_ts.shape:
+    return jnp.array(False)
+  return jnp.all(jnp.abs(agent_ts - tl_ts) <= tolerance_micros)
+
+
+def estimate_index_offset(
+    agent_ts: jax.Array,
+    tl_ts: jax.Array,
+    max_offset: int = 2,
+) -> int:
+  """Integer shift to apply to TL indices (positive shifts TL data forward)."""
+  length = int(min(agent_ts.shape[0], tl_ts.shape[0]))
+  if length == 0:
+    return 0
+
+  agent_ts = agent_ts[:length]
+  tl_ts = tl_ts[:length]
+  best_offset = 0
+  best_matches = -1
+  for offset in range(-max_offset, max_offset + 1):
+    if offset >= 0:
+      a = agent_ts[offset:]
+      t = tl_ts[: length - offset]
+    else:
+      a = agent_ts[: length + offset]
+      t = tl_ts[-offset:]
+    if a.shape[0] == 0:
+      continue
+    matches = int(jnp.sum(a == t))
+    if matches > best_matches:
+      best_matches = matches
+      best_offset = offset
+  return best_offset
+
+
+def _shift_along_time(
+    tls: traffic_lights.TrafficLights, offset: int
+) -> traffic_lights.TrafficLights:
+  """Shift TL tensors along the time axis by `offset` (edge-padded)."""
+  if offset == 0:
+    return tls
+
+  def shift(x: jax.Array) -> jax.Array:
+    if offset > 0:
+      pad = jax.lax.slice_in_dim(x, 0, 1, axis=-1)
+      pad = jnp.repeat(pad, offset, axis=-1)
+      return jnp.concatenate([pad, x[..., :-offset]], axis=-1)
+    neg = -offset
+    pad = jax.lax.slice_in_dim(x, x.shape[-1] - 1, x.shape[-1], axis=-1)
+    pad = jnp.repeat(pad, neg, axis=-1)
+    return jnp.concatenate([x[..., neg:], pad], axis=-1)
+
+  return jax.tree_util.tree_map(shift, tls)
+
+
+def resample_traffic_lights_nearest(
+    tls: traffic_lights.TrafficLights,
+    source_ts: jax.Array,
+    target_ts: jax.Array,
+) -> traffic_lights.TrafficLights:
+  """Resample TL states so output timestep `t` uses nearest `source_ts` entry."""
+  num_steps = target_ts.shape[0]
+  # For each target time, index of closest source time.
+  dist = jnp.abs(
+      target_ts[:, jnp.newaxis].astype(jnp.int64)
+      - source_ts[jnp.newaxis, :].astype(jnp.int64)
+  )
+  indices = jnp.argmin(dist, axis=1)
+
+  def gather_time(x: jax.Array) -> jax.Array:
+    # x: (..., num_tls, num_source_steps)
+    return jnp.take(x, indices, axis=-1)
+
+  return jax.tree_util.tree_map(gather_time, tls)
+
+
+def align_to_trajectory(
+    tls: traffic_lights.TrafficLights,
+    trajectory: object_state.Trajectory,
+    example: dict[str, jax.Array],
+    time_key: str = 'all',
+    is_sdc: jax.Array | None = None,
+    max_offset: int = 2,
+) -> traffic_lights.TrafficLights:
+  """Return TL states resampled onto the agent reference timeline if needed."""
+  tl_ts = extract_traffic_light_timestamps(example, time_key=time_key)
+  if tl_ts is None:
+    return tls
+
+  agent_ts = reference_agent_timestamps(trajectory, is_sdc=is_sdc)
+  if timestamps_aligned(agent_ts, tl_ts):
+    return tls
+
+  offset = estimate_index_offset(agent_ts, tl_ts, max_offset=max_offset)
+  shifted = _shift_along_time(tls, offset)
+  shifted_ts = _shift_along_time(
+      traffic_lights.TrafficLights(
+          x=tl_ts.astype(jnp.float32),
+          y=jnp.zeros_like(tl_ts, dtype=jnp.float32),
+          z=jnp.zeros_like(tl_ts, dtype=jnp.float32),
+          state=jnp.zeros_like(tl_ts, dtype=jnp.int32),
+          lane_ids=jnp.zeros_like(tl_ts, dtype=jnp.int32),
+          valid=jnp.ones_like(tl_ts, dtype=jnp.bool_),
+      ),
+      offset,
+  ).x.astype(jnp.int64)
+
+  if timestamps_aligned(agent_ts, shifted_ts):
+    return shifted
+
+  return resample_traffic_lights_nearest(tls, tl_ts, agent_ts)
diff --git a/waymax/dataloader/traffic_light_alignment_test.py b/waymax/dataloader/traffic_light_alignment_test.py
new file mode 100644
index 0000000..7bddb86
--- /dev/null
+++ b/waymax/dataloader/traffic_light_alignment_test.py
@@ -0,0 +1,167 @@
+# Copyright 2023 The Waymax Authors.
+#
+# Licensed under the Waymax License Agreement for Non-commercial Use
+# Use (the "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+#     https://github.com/waymo-research/waymax/blob/main/LICENSE
+#
+# 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 jax.numpy as jnp
+import tensorflow as tf
+
+from absl.testing import absltest
+from absl.testing import parameterized
+from waymax.dataloader import traffic_light_alignment
+from waymax.datatypes import object_state
+from waymax.datatypes import traffic_lights
+
+
+def _make_trajectory(
+    timestamps: list[int], num_agents: int = 1
+) -> object_state.Trajectory:
+  ts = jnp.asarray(timestamps, dtype=jnp.int32)
+  if num_agents > 1:
+    ts = jnp.tile(ts[jnp.newaxis, :], (num_agents, 1))
+  valid = jnp.ones(ts.shape, dtype=jnp.bool_)
+  return object_state.Trajectory(
+      x=jnp.zeros(ts.shape, jnp.float32),
+      y=jnp.zeros(ts.shape, jnp.float32),
+      z=jnp.zeros(ts.shape, jnp.float32),
+      vel_x=jnp.zeros(ts.shape, jnp.float32),
+      vel_y=jnp.zeros(ts.shape, jnp.float32),
+      yaw=jnp.zeros(ts.shape, jnp.float32),
+      valid=valid,
+      length=jnp.ones(ts.shape, jnp.float32),
+      width=jnp.ones(ts.shape, jnp.float32),
+      height=jnp.ones(ts.shape, jnp.float32),
+      timestamp_micros=ts,
+  )
+
+
+def _make_tls(states: list[int], num_tls: int = 1) -> traffic_lights.TrafficLights:
+  state = jnp.asarray(states, dtype=jnp.int32)
+  if state.ndim == 1:
+    state = state[jnp.newaxis, :]
+  if num_tls > 1:
+    state = jnp.tile(state, (num_tls, 1))
+  return traffic_lights.TrafficLights(
+      x=jnp.zeros(state.shape, jnp.float32),
+      y=jnp.zeros(state.shape, jnp.float32),
+      z=jnp.zeros(state.shape, jnp.float32),
+      state=state,
+      lane_ids=jnp.ones(state.shape, jnp.int32),
+      valid=jnp.ones(state.shape, jnp.bool_),
+  )
+
+
+class TrafficLightAlignmentTest(parameterized.TestCase, tf.test.TestCase):
+
+  def test_timestamps_aligned_true_when_equal(self):
+    ts = jnp.array([0, 100_000, 200_000], dtype=jnp.int64)
+    self.assertTrue(traffic_light_alignment.timestamps_aligned(ts, ts))
+
+  def test_estimate_offset_one_step_behind(self):
+    agent = jnp.array([0, 100_000, 200_000, 300_000], dtype=jnp.int64)
+    tl = jnp.array([100_000, 200_000, 300_000, 400_000], dtype=jnp.int64)
+    self.assertEqual(traffic_light_alignment.estimate_index_offset(agent, tl), 1)
+
+  def test_shift_restores_alignment(self):
+    agent_ts = jnp.array([0, 100_000, 200_000, 300_000], dtype=jnp.int64)
+    tl_ts = jnp.array([100_000, 200_000, 300_000, 400_000], dtype=jnp.int64)
+    tls = _make_tls([4, 6, 6, 6])
+    offset = traffic_light_alignment.estimate_index_offset(agent_ts, tl_ts)
+    shifted = traffic_light_alignment._shift_along_time(tls, offset)
+    # TL timestamps lag by one: agent t=2 (200ms) pairs with tl index 1 (GO).
+    self.assertEqual(int(shifted.state[0, 2]), 6)
+
+  def test_align_to_trajectory_fixes_one_step_desync(self):
+    timestamps = [0, 100_000, 200_000, 300_000]
+    traj = _make_trajectory(timestamps)
+    tls = _make_tls([4, 6, 6, 6])
+    example = {
+        'traffic_light_state/all/timestamp_micros': jnp.array(
+            [100_000, 200_000, 300_000, 400_000], dtype=jnp.int64
+        ),
+    }
+    aligned = traffic_light_alignment.align_to_trajectory(
+        tls, traj, example, is_sdc=jnp.array([True])
+    )
+    # Agent t=2 (200ms) should see GO (6) from TL index 1 after shift.
+    self.assertEqual(int(aligned.state[0, 2]), 6)
+
+  def test_align_noop_when_already_aligned(self):
+    timestamps = [0, 100_000, 200_000]
+    traj = _make_trajectory(timestamps)
+    tls = _make_tls([4, 5, 6])
+    example = {
+        'traffic_light_state/all/timestamp_micros': jnp.array(
+            timestamps, dtype=jnp.int64
+        ),
+    }
+    aligned = traffic_light_alignment.align_to_trajectory(
+        tls, traj, example, is_sdc=jnp.array([True])
+    )
+    self.assertEqual(aligned, tls)
+
+  def test_resample_nearest_handles_irregular_gap(self):
+    source_ts = jnp.array([0, 250_000, 500_000], dtype=jnp.int64)
+    target_ts = jnp.array([0, 100_000, 200_000], dtype=jnp.int64)
+    tls = _make_tls([4, 5, 6])
+    out = traffic_light_alignment.resample_traffic_lights_nearest(
+        tls, source_ts, target_ts
+    )
+    self.assertEqual(int(out.state[0, 0]), 4)
+    self.assertEqual(int(out.state[0, 1]), 4)
+    self.assertEqual(int(out.state[0, 2]), 5)  # 200ms nearest to source 250ms
+
+
+  def test_align_integration_via_factory(self):
+    timestamps = [0, 100_000, 200_000, 300_000]
+    traj = _make_trajectory(timestamps)
+    tls = _make_tls([4, 6, 6, 6])
+    example = {
+        'state/all/x': traj.x,
+        'state/all/y': traj.y,
+        'state/all/z': traj.z,
+        'state/all/velocity_x': traj.vel_x,
+        'state/all/velocity_y': traj.vel_y,
+        'state/all/bbox_yaw': traj.yaw,
+        'state/all/valid': traj.valid,
+        'state/all/length': traj.length,
+        'state/all/width': traj.width,
+        'state/all/height': traj.height,
+        'state/all/timestamp_micros': traj.timestamp_micros,
+        'state/is_sdc': jnp.array([1]),
+        'state/id': jnp.array([1], dtype=jnp.int32),
+        'state/type': jnp.array([1], dtype=jnp.int32),
+        'state/tracks_to_predict': jnp.array([1], dtype=jnp.int32),
+        'state/objects_of_interest': jnp.array([0], dtype=jnp.int32),
+        'traffic_light_state/all/x': jnp.zeros((4, 1), jnp.float32),
+        'traffic_light_state/all/y': jnp.zeros((4, 1), jnp.float32),
+        'traffic_light_state/all/z': jnp.zeros((4, 1), jnp.float32),
+        'traffic_light_state/all/state': jnp.transpose(tls.state),
+        'traffic_light_state/all/id': jnp.ones((4, 1), jnp.int32),
+        'traffic_light_state/all/valid': jnp.ones((4, 1), dtype=jnp.bool_),
+        'traffic_light_state/all/timestamp_micros': jnp.array(
+            [100_000, 200_000, 300_000, 400_000], dtype=jnp.int64
+        ),
+        'roadgraph_samples/xyz': jnp.zeros((1, 3), jnp.float32),
+        'roadgraph_samples/dir': jnp.zeros((1, 3), jnp.float32),
+        'roadgraph_samples/type': jnp.zeros((1, 1), jnp.int32),
+        'roadgraph_samples/id': jnp.zeros((1, 1), jnp.int32),
+        'roadgraph_samples/valid': jnp.ones((1, 1), dtype=jnp.bool_),
+    }
+    from waymax.dataloader import womd_factories
+
+    sim = womd_factories.simulator_state_from_womd_dict(example, time_key='all')
+    self.assertEqual(int(sim.log_traffic_light.state[0, 2]), 6)
+
+
+if __name__ == '__main__':
+  absltest.main()
diff --git a/waymax/dataloader/womd_factories.py b/waymax/dataloader/womd_factories.py
index 6e9fdd5..bce21f1 100644
--- a/waymax/dataloader/womd_factories.py
+++ b/waymax/dataloader/womd_factories.py
@@ -23,6 +23,7 @@
 from waymax.datatypes import route
 from waymax.datatypes import simulator_state
 from waymax.datatypes import traffic_lights
+from waymax.dataloader import traffic_light_alignment
 
 
 def object_metadata_from_womd_dict(
@@ -176,6 +177,14 @@ def simulator_state_from_womd_dict(
   object_metadata = object_metadata_from_womd_dict(example)
 
   log_trajectory = trajectory_from_womd_dict(example, time_key=time_key)
+  if time_key == 'all':
+    traffic_light = traffic_light_alignment.align_to_trajectory(
+        traffic_light,
+        log_trajectory,
+        example,
+        time_key=time_key,
+        is_sdc=object_metadata.is_sdc,
+    )
   # Init with zeros and false.
   sim_trajectory = jax.tree_util.tree_map(
       lambda x: jnp.zeros(x.shape, x.dtype), log_trajectory