# ----------------------------------------------------------------------------
# Copyright (c) 2021-2026 DexForce Technology Co., Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ----------------------------------------------------------------------------
import numpy as np
import torch
from typing import Union
try:
from scipy.stats import qmc
SCIPY_AVAILABLE = True
except ImportError:
SCIPY_AVAILABLE = False
from embodichain.lab.sim.motion.workspace.configs.sampling_config import (
SamplingStrategy,
)
from embodichain.lab.sim.motion.workspace.samplers.base_sampler import (
BaseSampler,
)
from embodichain.utils import logger
[docs]
class SobolSampler(BaseSampler):
"""Sobol sequence sampler using quasi-random low-discrepancy sequences.
The Sobol sequence is a low-discrepancy sequence that provides excellent
uniformity in high-dimensional spaces. It's widely used in finance,
engineering, and scientific computing for Monte Carlo simulations.
Advantages:
- Excellent uniformity in high dimensions (up to ~40 dimensions)
- Industry standard for quasi-Monte Carlo methods
- Better convergence than random sampling (O(1/n) vs O(1/√n))
- Well-suited for integration and optimization
Disadvantages:
- Requires scipy library
- Sequential generation (but can be scrambled for randomization)
- Initial points may not be well-distributed (use skip parameter)
Attributes:
scramble: Whether to scramble the sequence for better randomization.
skip: Number of initial samples to skip.
Notes:
This implementation uses scipy.stats.qmc.Sobol for efficient generation.
Falls back to a basic implementation if scipy is not available.
"""
[docs]
def __init__(
self,
seed: int = 42,
device: torch.device | None = None,
scramble: bool = True,
skip: int = 0,
):
"""Initialize the Sobol sampler.
Args:
seed: Random seed for scrambling. Defaults to 42.
device: PyTorch device (cpu/cuda). Defaults to cpu.
scramble: Whether to scramble the sequence. Defaults to True.
Scrambling improves randomization while maintaining low discrepancy.
skip: Number of initial samples to skip. Defaults to 0.
Recommended: 0 for scrambled, >0 (e.g., 100) for unscrambled.
constraint: Optional geometric constraint for sampling (e.g., SphereConstraint).
"""
super().__init__(seed, device)
self.scramble = scramble
self.skip = skip
if not SCIPY_AVAILABLE:
logger.log_warning(
"scipy is not available. Sobol sampler will use a basic fallback implementation. "
"For optimal performance, install scipy: pip install scipy"
)
[docs]
def sample(
self, bounds: torch.Tensor | np.ndarray, num_samples: int
) -> torch.Tensor:
"""Generate Sobol sequence samples within the given bounds.
Args:
bounds: Tensor/Array of shape (n_dims, 2) containing [lower, upper] bounds.
num_samples: Number of samples to generate.
Returns:
Tensor of shape (num_samples, n_dims) containing the sampled points.
Raises:
ValueError: If bounds are invalid or num_samples is non-positive.
Examples:
>>> sampler = SobolSampler(scramble=True, seed=42)
>>> bounds = torch.tensor([[-1.0, 1.0], [-1.0, 1.0]], dtype=torch.float32)
>>> samples = sampler.sample(bounds, num_samples=100)
>>> samples.shape
torch.Size([100, 2])
"""
bounds = self._validate_bounds(bounds)
if num_samples <= 0:
raise ValueError(f"num_samples must be positive, got {num_samples}")
n_dims = bounds.shape[0]
if n_dims > 21201: # Maximum dimension for scipy's Sobol
raise ValueError(
f"Sobol sequence supports up to 21201 dimensions, got {n_dims}"
)
# Generate Sobol sequence
if SCIPY_AVAILABLE:
samples_unit = self._generate_sobol_scipy(n_dims, num_samples)
else:
samples_unit = self._generate_sobol_fallback(n_dims, num_samples)
# Convert to tensor and scale to bounds
samples_unit_tensor = self._to_tensor(samples_unit)
samples = self._scale_samples(samples_unit_tensor, bounds)
# Validate samples
self._validate_samples(samples, bounds)
return samples
def _generate_sobol_scipy(self, n_dims: int, num_samples: int) -> np.ndarray:
"""Generate Sobol sequence using scipy.
Args:
n_dims: Number of dimensions.
num_samples: Number of samples to generate.
Returns:
Array of shape (num_samples, n_dims) with values in [0, 1].
"""
# Create Sobol engine
sobol_engine = qmc.Sobol(d=n_dims, scramble=self.scramble, seed=self.seed)
# Skip initial samples if requested
if self.skip > 0:
sobol_engine.fast_forward(self.skip)
# Generate samples
samples = sobol_engine.random(n=num_samples)
return samples.astype(np.float32)
def _generate_sobol_fallback(self, n_dims: int, num_samples: int) -> np.ndarray:
"""Fallback Sobol generator when scipy is not available.
This is a basic implementation and may not match scipy's quality.
For production use, install scipy.
Args:
n_dims: Number of dimensions.
num_samples: Number of samples to generate.
Returns:
Array of shape (num_samples, n_dims) with values in [0, 1].
"""
logger.log_warning(
"Using fallback Sobol generator. Results may not match scipy quality."
)
# Simple fallback: use stratified random sampling
# This is NOT a true Sobol sequence but provides reasonable coverage
samples = np.zeros((num_samples, n_dims), dtype=np.float32)
for dim in range(n_dims):
# Divide [0, 1] into num_samples intervals
intervals = np.linspace(0, 1, num_samples + 1)
# Sample randomly within each interval
samples[:, dim] = intervals[:-1] + self.rng.rand(num_samples) / num_samples
# Shuffle to reduce correlation
for dim in range(n_dims):
self.rng.shuffle(samples[:, dim])
return samples
[docs]
def get_strategy_name(self) -> str:
"""Get the name of the sampling strategy.
Returns:
String identifier for the sampling strategy.
"""
return SamplingStrategy.SOBOL.value
def __repr__(self) -> str:
"""String representation of the sampler."""
return (
f"{self.__class__.__name__}("
f"strategy={self.get_strategy_name()}, "
f"scramble={self.scramble}, "
f"skip={self.skip}, "
f"seed={self.seed})"
)
__all__ = ["SobolSampler"]