Add deterministic trajectory lookahead

iltools/datasets/manager.py

@@ -133,6 +133,7 @@ def __init__(
         reset_start_step: int = 0,
         wrap_steps: bool = False,
         device: torch.device | str | None = None,
+        reset_generator: torch.Generator | None = None,
         target_joint_names: Optional[Sequence[str]] = None,
         reference_joint_names: Optional[Sequence[str]] = None,
     ) -> None:
@@ -160,6 +161,14 @@ def __init__(
             else torch.device("cpu")
         )
         self._state_device = self._device or self._storage_device
+        if reset_generator is not None:
+            generator_device = torch.device(reset_generator.device)
+            if generator_device != self._state_device:
+                raise ValueError(
+                    "reset_generator device must match the trajectory-manager "
+                    f"state device; got {generator_device} and {self._state_device}."
+                )
+        self._reset_generator = reset_generator
 
         try:
             start = torch.as_tensor(
@@ -303,6 +312,11 @@ def state_device(self) -> torch.device:
         """Device holding the trajectory-manager state tensors."""
         return self._state_device
 
+    @property
+    def reset_generator(self) -> torch.Generator | None:
+        """Dedicated generator for trajectory/reset selection, if configured."""
+        return self._reset_generator
+
     def get_env_traj_info(self, env_id: int) -> tuple[str, str, str]:
         """Return (dataset, motion, trajectory) tuple for the env's current rank."""
         r = int(self.env_traj_rank[int(env_id)])
@@ -342,6 +356,7 @@ def _choose_new_ranks(self, env_ids: Tensor) -> Tensor:
                 size=(n,),
                 dtype=torch.int64,
                 device=env_ids.device,
+                generator=self._reset_generator,
             )
 
         if self.reset_schedule == ResetSchedule.SEQUENTIAL:
@@ -499,7 +514,9 @@ def _set_reference_fields(
         td.set(make_key("joint_pos"), joint_pos_out)
         td.set(make_key("joint_vel"), joint_vel_out)
 
-    def attach_reference_fields(self, td: TensorDict, *, use_buffers: bool) -> TensorDict:
+    def attach_reference_fields(
+        self, td: TensorDict, *, use_buffers: bool
+    ) -> TensorDict:
         """Attach root/joint reference fields to a sampled transition.
 
         Public API; ``_attach_reference_fields`` remains as the private
@@ -609,6 +626,52 @@ def sample(
         )
         return self._attach_reference_fields(td, use_buffers=use_buffers)
 
+    def current_global_indices(
+        self, env_ids: Sequence[int] | Tensor | None = None
+    ) -> Tensor:
+        """Return replay indices for the current cursors without reading storage."""
+        env_ids_t = (
+            self._all_env_ids if env_ids is None else self._normalize_env_ids(env_ids)
+        )
+        ranks = self.env_traj_rank.index_select(0, env_ids_t)
+        steps = self.env_step.index_select(0, env_ids_t)
+        return get_global_index(ranks, self._start, self._end, steps)
+
+    def global_indices_for(self, ranks: Tensor, steps: Tensor) -> Tensor:
+        """Map an explicit cursor batch to replay indices without mutating state.
+
+        This is the public look-ahead boundary used by asynchronous consumers:
+        callers may plan reset cursors before committing them to environments,
+        while this manager remains the sole owner of trajectory bounds and the
+        local-to-global index convention.
+        """
+        ranks_t = torch.as_tensor(ranks, device=self._state_device, dtype=torch.int64)
+        steps_t = torch.as_tensor(steps, device=self._state_device, dtype=torch.int64)
+        if ranks_t.ndim != 1 or steps_t.ndim != 1 or ranks_t.shape != steps_t.shape:
+            raise ValueError("ranks and steps must be matching 1D tensors.")
+        if torch.any((ranks_t < 0) | (ranks_t >= self.num_trajectories)):
+            raise ValueError("ranks contains an out-of-range value.")
+        lengths = self._length.index_select(0, ranks_t)
+        if torch.any((steps_t < 0) | (steps_t >= lengths)):
+            raise ValueError("steps contains an out-of-range value.")
+        return get_global_index(ranks_t, self._start, self._end, steps_t)
+
+    def advance_cursors(
+        self, env_ids: Sequence[int] | Tensor | None = None
+    ) -> tuple[Tensor, Tensor]:
+        """Advance cursors without reading storage and return steps and indices.
+
+        The returned tensors are clones: reset handling may mutate the live
+        cursors while a caller asynchronously gathers the planned rows.
+        """
+        env_ids_t = (
+            self._all_env_ids if env_ids is None else self._normalize_env_ids(env_ids)
+        )
+        self._advance_steps(env_ids_t)
+        steps = self.env_step.index_select(0, env_ids_t).clone()
+        indices = self.current_global_indices(env_ids_t).clone()
+        return steps, indices
+
     def sample_slice(
         self,
         batch_size: int,

iltools/datasets/reset_sampling.py

@@ -73,14 +73,13 @@ def __init__(
         random_step_max: int = 0,
         weight_fn: WeightFunction | None = None,
         device: torch.device | str | None = None,
+        generator: torch.Generator | None = None,
     ) -> None:
         lengths = torch.as_tensor(
             trajectory_lengths, dtype=torch.long, device=trajectory_lengths.device
         ).reshape(-1)
         if lengths.numel() == 0:
-            raise ValueError(
-                "trajectory_lengths must contain at least one trajectory."
-            )
+            raise ValueError("trajectory_lengths must contain at least one trajectory.")
         if torch.any(lengths <= 0):
             raise ValueError("trajectory_lengths must all be positive.")
 
@@ -101,25 +100,25 @@ def __init__(
                 "adaptive starting-frame mode requires a weight_fn callable."
             )
         if mode != self.ADAPTIVE and weight_fn is not None:
-            raise ValueError(
-                "weight_fn is only used in adaptive starting-frame mode."
-            )
+            raise ValueError("weight_fn is only used in adaptive starting-frame mode.")
 
-        self._device = (
-            torch.device(device) if device is not None else lengths.device
-        )
+        self._device = torch.device(device) if device is not None else lengths.device
         self.trajectory_lengths = lengths.to(self._device)
         self.mode = mode
         self.fixed_step = int(fixed_step)
         self.random_step_min = int(random_step_min)
         self.random_step_max = int(random_step_max)
         self.weight_fn = weight_fn
+        if generator is not None and torch.device(generator.device) != self._device:
+            raise ValueError(
+                "generator device must match the sampler device; got "
+                f"{torch.device(generator.device)} and {self._device}."
+            )
+        self.generator = generator
 
     def _clamp_steps(self, ranks: torch.Tensor, steps: torch.Tensor) -> torch.Tensor:
         max_steps = self.trajectory_lengths.index_select(0, ranks) - 1
-        return torch.minimum(
-            torch.maximum(steps, torch.zeros_like(steps)), max_steps
-        )
+        return torch.minimum(torch.maximum(steps, torch.zeros_like(steps)), max_steps)
 
     def sample_steps(self, trajectory_ranks: torch.Tensor) -> torch.Tensor:
         """Sample one local starting frame per requested trajectory rank.
@@ -151,6 +150,7 @@ def sample_steps(self, trajectory_ranks: torch.Tensor) -> torch.Tensor:
                     (n,),
                     device=self._device,
                     dtype=torch.long,
+                    generator=self.generator,
                 )
             else:
                 steps = torch.full(
@@ -195,7 +195,7 @@ def _sample_adaptive(self, ranks: torch.Tensor) -> torch.Tensor:
             fallback[zero_rows] = valid[zero_rows].to(torch.float32)
             weights = torch.where(zero_rows[:, None], fallback, weights)
         probs = weights / weights.sum(dim=-1, keepdim=True)
-        return torch.multinomial(probs, 1).squeeze(-1)
+        return torch.multinomial(probs, 1, generator=self.generator).squeeze(-1)
 
 
 class SonicAdaptiveResetSampler:
@@ -229,16 +229,15 @@ def __init__(
         uniform_sampling_rate: float = 0.1,
         pre_failure_sample_window: int = 200,
         failure_rate_max_over_mean: float = 200.0,
+        generator: torch.Generator | None = None,
     ) -> None:
         lengths = torch.as_tensor(
             trajectory_lengths,
             dtype=torch.long,
             device=trajectory_lengths.device,
         ).reshape(-1)
         if lengths.numel() == 0:
-            raise ValueError(
-                "trajectory_lengths must contain at least one trajectory."
-            )
+            raise ValueError("trajectory_lengths must contain at least one trajectory.")
         if torch.any(lengths <= 0):
             raise ValueError("trajectory_lengths must all be positive.")
         if int(bin_size) <= 0:
@@ -259,6 +258,12 @@ def __init__(
         self.uniform_sampling_rate = float(uniform_sampling_rate)
         self.pre_failure_sample_window = int(pre_failure_sample_window)
         self.failure_rate_max_over_mean = float(failure_rate_max_over_mean)
+        if generator is not None and torch.device(generator.device) != self.device:
+            raise ValueError(
+                "generator device must match the sampler device; got "
+                f"{torch.device(generator.device)} and {self.device}."
+            )
+        self.generator = generator
 
         bins: list[torch.Tensor] = []
         trajectory_bin_ids: list[torch.Tensor] = []
@@ -325,9 +330,7 @@ def _bin_ids(
         if torch.any((ranks < 0) | (ranks >= self.trajectory_lengths.numel())):
             raise ValueError("trajectory_ranks contains an out-of-range value.")
         max_steps = self.trajectory_lengths.index_select(0, ranks) - 1
-        steps = torch.minimum(
-            torch.maximum(steps, torch.zeros_like(steps)), max_steps
-        )
+        steps = torch.minimum(torch.maximum(steps, torch.zeros_like(steps)), max_steps)
         local_bins = torch.div(steps, self.bin_size, rounding_mode="floor")
         return self.first_bin_ids.index_select(0, ranks) + local_bins
 
@@ -381,41 +384,71 @@ def weights(
         bin_ids = self._bin_ids(trajectory_ranks, frame_steps)
         bin_probs = self.sampling_probabilities()
         bin_lengths = self.bin_lengths
-        return (
-            bin_probs.index_select(0, bin_ids)
-            / bin_lengths.index_select(0, bin_ids).to(dtype=torch.float32)
-        )
+        return bin_probs.index_select(0, bin_ids) / bin_lengths.index_select(
+            0, bin_ids
+        ).to(dtype=torch.float32)
 
     def __call__(
         self, trajectory_ranks: torch.Tensor, frame_steps: torch.Tensor
     ) -> torch.Tensor:
         """Callable alias of :meth:`weights` for use as a ``weight_fn``."""
         return self.weights(trajectory_ranks, frame_steps)
 
-    def sample(self, count: int) -> tuple[torch.Tensor, torch.Tensor]:
-        """Sample trajectory ranks and random local starts with SONIC's lead-in."""
+    def sample(
+        self,
+        count: int,
+        *,
+        probabilities: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Sample trajectory ranks and local starts with SONIC's lead-in.
+
+        ``probabilities`` may hold a caller-owned snapshot of the bin
+        distribution. This makes the sampling-time contract explicit for
+        asynchronous consumers instead of racing later failure updates.
+        """
         count = int(count)
         if count < 0:
             raise ValueError("count must be >= 0.")
         if count == 0:
             empty = torch.empty(0, device=self.device, dtype=torch.long)
             return empty, empty
+        if probabilities is None:
+            probabilities = self.sampling_probabilities()
+        else:
+            probabilities = torch.as_tensor(
+                probabilities, device=self.device, dtype=torch.float32
+            )
+            if probabilities.shape != (self.num_bins,):
+                raise ValueError(
+                    "probabilities must have one entry per SONIC bin; expected "
+                    f"{(self.num_bins,)}, got {tuple(probabilities.shape)}."
+                )
+            if not torch.all(torch.isfinite(probabilities)):
+                raise ValueError("probabilities must be finite")
+            if torch.any(probabilities < 0) or probabilities.sum() <= 0:
+                raise ValueError("probabilities must be non-negative with positive sum")
+            probabilities = probabilities / probabilities.sum()
         sampled_bin_ids = torch.multinomial(
-            self.sampling_probabilities(), count, replacement=True
+            probabilities,
+            count,
+            replacement=True,
+            generator=self.generator,
         )
         sampled_bins = self.bins.index_select(0, sampled_bin_ids)
         trajectory_ranks = sampled_bins[:, 0]
         bin_starts = sampled_bins[:, 1]
         bin_ends = sampled_bins[:, 2]
         frame_steps = (
-            torch.rand(count, device=self.device) * (bin_ends - bin_starts)
+            torch.rand(count, device=self.device, generator=self.generator)
+            * (bin_ends - bin_starts)
         ).floor().to(torch.long) + bin_starts
         if self.pre_failure_sample_window > 0:
             lead_in = torch.randint(
                 self.pre_failure_sample_window,
                 (count,),
                 device=self.device,
                 dtype=torch.long,
+                generator=self.generator,
             )
             frame_steps = (frame_steps - lead_in).clamp_min(0)
         return trajectory_ranks, frame_steps

tests/datasets/test_parallel_trajectory_manager.py

@@ -151,6 +151,99 @@ def test_parallel_trajectory_manager_direct_indexing_and_reset(tmp_path):
     assert mgr.env_traj_rank.tolist()[0] in [0, 1]
 
 
+def test_parallel_trajectory_manager_advance_cursors_without_sampling(tmp_path):
+    from iltools.datasets.manager import ParallelTrajectoryManager, ResetSchedule
+
+    rb, traj_info = _make_step_rb_and_traj_info(tmp_path, lengths=(3, 5))
+    mgr = ParallelTrajectoryManager(
+        rb=rb,
+        traj_info=traj_info,
+        num_envs=2,
+        reset_schedule=ResetSchedule.SEQUENTIAL,
+        wrap_steps=False,
+        target_joint_names=["joint1", "joint2"],
+        reference_joint_names=["joint1", "joint2"],
+    )
+    mgr.set_env_cursor(
+        env_ids=[0, 1], ranks=torch.tensor([0, 1]), steps=torch.tensor([1, 4])
+    )
+
+    steps, indices = mgr.advance_cursors()
+
+    assert steps.tolist() == [2, 4]
+    assert indices.tolist() == [2, 7]
+    assert mgr.global_indices_for(
+        torch.tensor([1, 0]), torch.tensor([3, 1])
+    ).tolist() == [6, 1]
+    # Returned plans must not alias the live cursors: reset handling is allowed
+    # to mutate the manager while an asynchronous reader consumes the plan.
+    mgr.reset_envs([0], ranks=torch.tensor([1]), steps=torch.tensor([0]))
+    assert steps.tolist() == [2, 4]
+    assert indices.tolist() == [2, 7]
+
+
+def test_parallel_trajectory_manager_maps_planned_cursors_without_mutation(tmp_path):
+    from iltools.datasets.manager import ParallelTrajectoryManager, ResetSchedule
+
+    rb, traj_info = _make_step_rb_and_traj_info(tmp_path, lengths=(3, 5))
+    mgr = ParallelTrajectoryManager(
+        rb=rb,
+        traj_info=traj_info,
+        num_envs=2,
+        reset_schedule=ResetSchedule.SEQUENTIAL,
+        target_joint_names=["joint1", "joint2"],
+        reference_joint_names=["joint1", "joint2"],
+    )
+    original_ranks = mgr.env_traj_rank.clone()
+    original_steps = mgr.env_step.clone()
+
+    indices = mgr.global_indices_for(torch.tensor([1, 0]), torch.tensor([4, 2]))
+
+    assert indices.tolist() == [7, 2]
+    torch.testing.assert_close(mgr.env_traj_rank, original_ranks)
+    torch.testing.assert_close(mgr.env_step, original_steps)
+    with pytest.raises(ValueError, match="matching 1D"):
+        mgr.global_indices_for(torch.tensor([[0]]), torch.tensor([0]))
+    with pytest.raises(ValueError, match="ranks"):
+        mgr.global_indices_for(torch.tensor([2]), torch.tensor([0]))
+    with pytest.raises(ValueError, match="steps"):
+        mgr.global_indices_for(torch.tensor([0]), torch.tensor([3]))
+
+
+def test_parallel_trajectory_manager_reset_generator_is_dedicated(tmp_path):
+    from iltools.datasets.manager import ParallelTrajectoryManager, ResetSchedule
+
+    rb, traj_info = _make_dummy_rb_and_traj_info(tmp_path, lengths=(3, 5, 4))
+
+    def _manager(seed: int) -> ParallelTrajectoryManager:
+        generator = torch.Generator(device="cpu")
+        generator.manual_seed(seed)
+        return ParallelTrajectoryManager(
+            rb=rb,
+            traj_info=traj_info,
+            num_envs=8,
+            reset_schedule=ResetSchedule.RANDOM,
+            reset_generator=generator,
+            target_joint_names=["joint1", "joint2"],
+            reference_joint_names=["joint1", "joint2"],
+        )
+
+    torch.manual_seed(1)
+    first = _manager(123)
+    first_initial = first.env_traj_rank.clone()
+    first.reset_envs(torch.arange(8))
+    first_next = first.env_traj_rank.clone()
+
+    torch.manual_seed(9999)
+    second = _manager(123)
+    second_initial = second.env_traj_rank.clone()
+    second.reset_envs(torch.arange(8))
+    second_next = second.env_traj_rank.clone()
+
+    torch.testing.assert_close(first_initial, second_initial)
+    torch.testing.assert_close(first_next, second_next)
+
+
 def test_parallel_trajectory_manager_round_robin(tmp_path):
     from iltools.datasets.manager import ParallelTrajectoryManager, ResetSchedule