From f1e6ad106dea22f5ebc5cd4c4316fc990429935f Mon Sep 17 00:00:00 2001 From: Feiyang Wu Date: Wed, 5 Aug 2026 16:20:25 -0400 Subject: [PATCH] Add deterministic trajectory lookahead --- iltools/datasets/manager.py | 65 ++++++++++++- iltools/datasets/reset_sampling.py | 87 +++++++++++------ .../test_parallel_trajectory_manager.py | 93 +++++++++++++++++++ tests/datasets/test_reset_sampling.py | 72 ++++++++++++-- 4 files changed, 282 insertions(+), 35 deletions(-) diff --git a/iltools/datasets/manager.py b/iltools/datasets/manager.py index c8c3467..3838478 100644 --- a/iltools/datasets/manager.py +++ b/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, diff --git a/iltools/datasets/reset_sampling.py b/iltools/datasets/reset_sampling.py index 2fe17ce..888b701 100644 --- a/iltools/datasets/reset_sampling.py +++ b/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,6 +229,7 @@ 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, @@ -236,9 +237,7 @@ def __init__( 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,10 +384,9 @@ 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 @@ -392,23 +394,53 @@ def __call__( """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( @@ -416,6 +448,7 @@ def sample(self, count: int) -> tuple[torch.Tensor, torch.Tensor]: (count,), device=self.device, dtype=torch.long, + generator=self.generator, ) frame_steps = (frame_steps - lead_in).clamp_min(0) return trajectory_ranks, frame_steps diff --git a/tests/datasets/test_parallel_trajectory_manager.py b/tests/datasets/test_parallel_trajectory_manager.py index 19a0345..0b862b3 100644 --- a/tests/datasets/test_parallel_trajectory_manager.py +++ b/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 diff --git a/tests/datasets/test_reset_sampling.py b/tests/datasets/test_reset_sampling.py index 429051c..1436360 100644 --- a/tests/datasets/test_reset_sampling.py +++ b/tests/datasets/test_reset_sampling.py @@ -98,6 +98,47 @@ def test_random_full_trajectory_starts_apply_sonic_lead_in() -> None: assert torch.any(lead_in_steps == 0) +def test_sonic_dedicated_generator_is_independent_of_global_rng() -> None: + lengths = torch.tensor([500, 260]) + + def _sample(global_seed: int) -> tuple[torch.Tensor, torch.Tensor]: + torch.manual_seed(global_seed) + generator = torch.Generator(device="cpu") + generator.manual_seed(77) + sampler = SonicAdaptiveResetSampler(lengths, generator=generator) + return sampler.sample(512) + + first_ranks, first_steps = _sample(1) + second_ranks, second_steps = _sample(9999) + torch.testing.assert_close(first_ranks, second_ranks) + torch.testing.assert_close(first_steps, second_steps) + + +def test_sonic_probability_snapshot_is_not_changed_by_later_failures() -> None: + generator = torch.Generator(device="cpu") + generator.manual_seed(19) + sampler = SonicAdaptiveResetSampler( + torch.tensor([100, 100]), + pre_failure_sample_window=0, + generator=generator, + ) + snapshot = sampler.sampling_probabilities().clone() + sampler.num_visits.fill_(100.0) + sampler.num_failures.fill_(1.0) + sampler.num_failures[-1] = 100.0 + + ranks, steps = sampler.sample(4096, probabilities=snapshot) + # The frozen initial distribution is balanced across the equal-length + # motions even though the live distribution now overwhelmingly favors 1. + first_fraction = (ranks == 0).float().mean() + assert 0.45 < first_fraction < 0.55 + assert torch.all(steps >= 0) + assert torch.all(steps < 100) + + with pytest.raises(ValueError, match="one entry per SONIC bin"): + sampler.sample(1, probabilities=torch.ones(3)) + + # --------------------------------------------------------------------------- # StartFrameSampler: fixed / random modes. # --------------------------------------------------------------------------- @@ -143,6 +184,25 @@ def test_random_mode_with_single_value_is_fixed() -> None: assert sampler.sample_steps(torch.tensor([0])).tolist() == [7] +def test_start_frame_dedicated_generator_is_independent_of_global_rng() -> None: + ranks = torch.tensor([0, 1, 0, 1, 0, 1, 0, 1]) + + def _sample(global_seed: int) -> torch.Tensor: + torch.manual_seed(global_seed) + generator = torch.Generator(device="cpu") + generator.manual_seed(91) + sampler = StartFrameSampler( + torch.tensor([100, 100]), + mode="random", + random_step_min=10, + random_step_max=20, + generator=generator, + ) + return sampler.sample_steps(ranks) + + torch.testing.assert_close(_sample(2), _sample(2000)) + + def test_empty_rank_batch_returns_empty() -> None: sampler = StartFrameSampler(torch.tensor([100])) assert sampler.sample_steps(torch.empty(0, dtype=torch.long)).numel() == 0 @@ -204,14 +264,14 @@ def zero_for_first(ranks: torch.Tensor, steps: torch.Tensor) -> torch.Tensor: torch.manual_seed(3) steps = sampler.sample_steps(torch.tensor([0, 1, 0, 1])) assert torch.all(steps >= 0) - assert torch.all(steps < torch.tensor([100, 60]).index_select(0, torch.tensor([0, 1, 0, 1]))) + assert torch.all( + steps < torch.tensor([100, 60]).index_select(0, torch.tensor([0, 1, 0, 1])) + ) def test_adaptive_mode_sanitizes_non_finite_weights() -> None: def nan_weights(ranks: torch.Tensor, steps: torch.Tensor) -> torch.Tensor: - return torch.full( - (ranks.numel(),), float("nan"), device=ranks.device - ) + return torch.full((ranks.numel(),), float("nan"), device=ranks.device) sampler = StartFrameSampler( torch.tensor([100]), @@ -289,9 +349,7 @@ def test_adaptive_frames_follow_recorded_failures() -> None: failures in one bin must shift the sampled frame distribution toward it. """ lengths = torch.tensor([200, 200]) - sonic = SonicAdaptiveResetSampler( - lengths, bin_size=50, pre_failure_sample_window=0 - ) + sonic = SonicAdaptiveResetSampler(lengths, bin_size=50, pre_failure_sample_window=0) frame_sampler = StartFrameSampler( lengths, mode="adaptive", weight_fn=sonic, device="cpu" )