Enhance workspace analyzer computational efficiency

[yuecideng]

Description

This PR optimizes the workspace analyzer for large-scale sampling (10k-100k+ samples) by replacing per-sample Python loops with batched GPU operations and vectorized algorithms.

Changes

Batch FK/IK (P0 - highest impact):

• workspace_analyzer.py: Replaced per-sample compute_fk() loop with chunked compute_batch_fk() calls
• workspace_analyzer.py: Replaced per-sample compute_ik() loop with chunked compute_batch_ik() calls
• Extracted _get_reference_pose() helper for cleaner IK reference pose handling
• Added check_constraints() to WorkspaceConstraintChecker combining bounds + collision in one call

Vectorized metrics (P1):

• density_metric.py: Replaced O(N²) brute-force density with scipy.spatial.cKDTree (O(N log N)), with brute-force fallback
• reachability_metric.py: Replaced dict-based voxel counting with np.unique(return_counts=True)
• manipulability_metric.py: Replaced per-sample Jacobian loops with batch np.linalg.det() and np.linalg.svd()

Vectorized sampler (P1):

• halton_sampler.py: Replaced nested Python loops with vectorized van der Corput computation using numpy broadcasting

Benchmark script:

• scripts/benchmark/workspace_analyzer/benchmark_workspace_analyzer.py: Measures each optimization independently across sample sizes [100, 1000, 10000, 100000]

Benchmark Results (CPU-only metrics)

Component	10k samples	50k-100k samples
Halton sampler	8.5 ms	103 ms (100k)
Density metric	18 ms	135 ms (50k)
Voxelization	18 ms	208 ms (100k)
Manipulability	79 ms	395 ms (50k)

Batch FK/IK benchmarks require live robot/simulation and are expected to show 10-100x speedup.

Type of change

• Bug fix (non-breaking change which fixes an issue)
• Enhancement (non-breaking change which improves an existing functionality)
• New feature (non-breaking change which adds functionality)
• Breaking change (existing functionality will not work without user modification)
• Documentation update

Screenshots

N/A

Checklist

• I have run the black . command to format the code base.
• I have made corresponding changes to the documentation
• I have added tests that prove my fix is effective or that my feature works
• Dependencies have been updated, if applicable.

🤖 Generated with Claude Code

[Copilot]

Pull request overview

This PR focuses on speeding up workspace analysis at large sample counts by batching FK/IK calls and vectorizing several metrics and sampling components.

Changes:

• Updated WorkspaceAnalyzer to use batched FK (compute_batch_fk) and batched IK (compute_batch_ik), including a helper for IK reference pose selection.
• Vectorized/optimized metrics (density via KDTree, reachability voxelization via np.unique, manipulability via batched det/svd) and Halton sampling.
• Added a benchmark script to measure performance of individual optimizations.

Reviewed changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
scripts/benchmark/workspace_analyzer/benchmark_workspace_analyzer.py	Adds a standalone benchmarking entrypoint for sampler/metric optimizations and placeholders for FK/IK benchmarks.
embodichain/lab/sim/utility/workspace_analyzer/workspace_analyzer.py	Switches FK/IK computation to batched robot APIs; adds _get_reference_pose and updates constraint checking usage.
embodichain/lab/sim/utility/workspace_analyzer/samplers/halton_sampler.py	Replaces loop-based Halton generation with a vectorized van der Corput implementation.
embodichain/lab/sim/utility/workspace_analyzer/metrics/reachability_metric.py	Replaces dict-based voxel counting with np.unique(..., return_counts=True)-based counting.
embodichain/lab/sim/utility/workspace_analyzer/metrics/manipulability_metric.py	Uses batched linear algebra for manipulability/condition numbers; handles empty valid-score cases.
embodichain/lab/sim/utility/workspace_analyzer/metrics/density_metric.py	Uses scipy.spatial.cKDTree to avoid O(N²) radius-density computation, with brute-force fallback.
embodichain/lab/sim/utility/workspace_analyzer/constraints/workspace_constraint.py	Adds check_constraints() convenience method combining bounds + collision checks.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

If cartesian_points is empty, batch_size = min(batch_size, num_samples) becomes 0, which will later make range(0, num_samples, batch_size) raise (step cannot be 0). Add an early return for num_samples == 0 before computing/clamping batch_size, or clamp batch_size to at least 1.

Suggested change

	num_samples = len(cartesian_points)
	num_samples = len(cartesian_points)
	if num_samples == 0:
	return (
	cartesian_points,
	cartesian_points[:0],
	cartesian_points.new_empty((0,)),
	torch.zeros((0,), dtype=torch.bool, device=cartesian_points.device),
	cartesian_points.new_empty((0, 0)),
	)

[Copilot]

This log line divides by num_samples unconditionally: (valid_cartesian_mask.sum()/num_samples*100). If cartesian_points is empty (num_samples == 0), this will raise a ZeroDivisionError before the early-return fix above. Please guard the percentage calculation when num_samples == 0.

Suggested change

	logger.log_info(
	f"Pre-filtered Cartesian points: {valid_cartesian_mask.sum()}/{num_samples} "
	f"points pass workspace constraints ({(valid_cartesian_mask.sum()/num_samples*100):.1f}%)"
	valid_percentage = (
	valid_cartesian_mask.sum() / num_samples * 100 if num_samples > 0 else 0.0
	)
	logger.log_info(
	f"Pre-filtered Cartesian points: {valid_cartesian_mask.sum()}/{num_samples} "
	f"points pass workspace constraints ({valid_percentage:.1f}%)"

[Copilot]

batch_best_success is used as a success counter, but newly-succeeded entries are set to 1.0 and then immediately incremented again via batch_best_success += success.float(), which double-counts the first success for each point (inflating success_rates_batch). Consider tracking (1) a boolean has_best for storing first successful qpos and (2) a separate success_counts tensor for counting successes, or just remove the = 1.0 assignment and only accumulate via += success.float().

[Copilot]

cKDTree.query_ball_point(..., return_length=True) is not supported in some SciPy versions (it raises TypeError: unexpected keyword argument 'return_length'). Since SciPy isn’t version-pinned in this repo, consider catching TypeError here (in addition to ImportError) and falling back to the brute-force path or to np.fromiter((len(x) for x in tree.query_ball_point(...)), ...) for broader compatibility.

Suggested change

	# Count neighbors within radius for all points at once
	counts = tree.query_ball_point(points, r=radius, return_length=True)
	# Count neighbors within radius for all points at once.
	# Older SciPy versions do not support `return_length=True`,
	# so fall back to counting the returned neighbor lists.
	try:
	counts = tree.query_ball_point(points, r=radius, return_length=True)
	except TypeError:
	neighbors = tree.query_ball_point(points, r=radius)
	counts = np.fromiter(
	(len(neighbor_list) for neighbor_list in neighbors),
	dtype=np.int64,
	count=len(points),
	)

[Copilot]

When joint_configs is empty, batch_size = min(batch_size, num_samples) becomes 0 and then range(0, num_samples, batch_size) raises (step cannot be 0). Consider an early return for num_samples == 0 (empty workspace_points / valid_configs) before computing batch_size/creating the tqdm range, or ensure batch_size >= 1.

[Copilot]

Pull request overview

Copilot reviewed 12 out of 12 changed files in this pull request and generated 7 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

The docstring now states the grid size "will not exceed" num_samples, but with samples_per_dim = max(2, int(num_samples ** (1.0 / n_dims))) the actual grid can still exceed num_samples for small requests (e.g., num_samples=3, n_dims=2 -> samples_per_dim=2 -> 4 samples). Either adjust the wording to match the behavior, or change the logic so samples_per_dim**n_dims <= num_samples while still handling small num_samples cleanly.

[Copilot]

When no Jacobians pass the threshold, setting valid_scores = np.array([0.0]) makes num_valid_points report 1 and forces min/max to 0, which is misleading (it implies one valid sample). Prefer keeping num_valid_points = 0 and explicitly setting mean/std/min/max to 0.0 (or NaN) for the empty-valid set.

[Copilot]

This example now hard-codes sim_device="cuda", which will fail on CPU-only machines and in CI environments without a GPU. Consider keeping the default as cpu (or adding a simple runtime check / CLI flag) so the example remains runnable everywhere.

[Copilot]

Raising the default num_samples from 3,000 to 30,000 can make this example take significantly longer (and may be too slow for quick-start usage). Consider keeping a smaller default and documenting how to increase it for performance testing, or making it configurable via an argument/env var.

[Copilot]

batch_size = min(batch_size, num_samples) can produce batch_size == 0 when num_samples is 0 (or config batch_size is 0), which makes range(0, num_samples, batch_size) raise ValueError (step must be non-zero). Add an early return for num_samples == 0 (and/or validate batch_size >= 1) before constructing the progress range.

[Copilot]

This log line computes a percentage as a torch scalar and formats it with :.1f, which raises TypeError: unsupported format string passed to Tensor.__format__. Also, when num_samples == 0 it will divide by zero. Convert the mask sum to a Python number via .item() and guard the percentage computation for the empty-input case.

Suggested change

	logger.log_info(
	f"Pre-filtered Cartesian points: {valid_cartesian_mask.sum()}/{num_samples} "
	f"points pass workspace constraints ({(valid_cartesian_mask.sum()/num_samples*100):.1f}%)"
	num_valid_cartesian_points = valid_cartesian_mask.sum().item()
	valid_cartesian_percentage = (
	(num_valid_cartesian_points / num_samples * 100) if num_samples else 0.0
	)
	logger.log_info(
	f"Pre-filtered Cartesian points: {num_valid_cartesian_points}/{num_samples} "
	f"points pass workspace constraints ({valid_cartesian_percentage:.1f}%)"

[Copilot]

batch_size = min(batch_size, num_samples) can yield batch_size == 0 when num_samples is 0 (or batch_size is misconfigured as 0). That will make range(0, num_samples, batch_size) invalid (step=0). Consider early-returning a fully-empty result tuple for num_samples == 0 and validating batch_size >= 1 before constructing the tqdm range.

[Copilot]

Pull request overview

Copilot reviewed 12 out of 12 changed files in this pull request and generated 2 comments.

Comments suppressed due to low confidence (1)

embodichain/lab/sim/utility/workspace_analyzer/samplers/iniform_sampler.py:108

• The updated docstring/logic claims the generated grid size (samples_per_dim**n_dims) will not exceed num_samples, but samples_per_dim = max(2, int(num_samples ** (1/n_dims))) can still overshoot for small num_samples / high n_dims (e.g., num_samples=2, n_dims=6 yields samples_per_dim=2 => 64 samples). Either relax the docstring, or adjust the computation to ensure samples_per_dim**n_dims <= num_samples (e.g., allow samples_per_dim to be 1, or decrement until the constraint holds).

                Note: The actual number of samples (samples_per_dim^n_dims) will not
                exceed this value, but may be less to maintain a uniform grid.

        Returns:
            Tensor of shape (actual_num_samples, n_dims) containing the sampled points.
            The actual number of samples will be samples_per_dim^n_dims.

        Raises:
            ValueError: If bounds are invalid.

        Examples:
            >>> sampler = UniformSampler(samples_per_dim=3)
            >>> bounds = torch.tensor([[-1, 1], [-1, 1]], dtype=torch.float32)
            >>> samples = sampler.sample(bounds, num_samples=10)
            >>> samples.shape
            torch.Size([9, 2])  # 3^2 = 9 samples
        """
        bounds = self._validate_bounds(bounds)

        n_dims = bounds.shape[0]

        # Calculate samples per dimension if not provided
        if self.samples_per_dim is None:
            # Compute samples_per_dim to approximate the desired num_samples
            # Use floor to ensure actual grid size never exceeds num_samples
            samples_per_dim = max(2, int(num_samples ** (1.0 / n_dims)))
        else:
            samples_per_dim = self.samples_per_dim

        actual_num_samples = samples_per_dim**n_dims

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

compute_reachability() has the same empty-input issue as FK: if cartesian_points is empty then batch_size becomes 0 and range(0, num_samples, batch_size) will throw. Also the pre-filter log computes (valid_cartesian_mask.sum()/num_samples*100) which will divide by zero. Add a num_samples == 0 guard before logging / creating the progress bar and return empty outputs.

[Copilot]

_get_reference_pose() only moves reference_pose to self.device when it is a NumPy array. If callers provide a torch.Tensor on CPU (or wrong dtype), the returned pose can be on the wrong device and later compute_batch_ik will error due to device mismatch. Normalize reference_pose for both numpy and torch inputs (e.g., reference_pose = torch.as_tensor(..., device=self.device, dtype=torch.float32)), and keep the (1,4,4) shape handling.