Improve workspace batching and analytic IK buffer reuse

agent_context/topics/ik-solvers/ik-solvers.md

@@ -140,3 +140,20 @@ solver interfaces. The compute kernels do not import simulation modules.
 `utils/warp/kinematics/*_solver.py` are compatibility aliases.
 Validate kernel import/compilation with `tests/compute/test_imports.py` and
 solver behavior with the corresponding `tests/sim/motion/solvers/` tests.
+
+## Batch adapters and analytic scratch memory
+
+`BaseSolver.get_fk_batch()` and `get_ik_batch()` flatten and restore arbitrary
+leading batch axes in the chain-root frame. Nearest IK returns a boolean success
+mask with the leading shape and qpos with a final `dof` axis; existing concrete
+`get_fk`/`get_ik` signatures remain available. Robot owns local-arena/root frame
+conversion and broadcasts root transforms without materializing per-target copies.
+
+UR/OPW reuse internal candidate buffers through `solvers/_buffers.py`.
+`prepare_buffers(max_batch)` reserves a high-water capacity; allocation is lazy
+and later larger calls grow it. Public results remain independent of subsequent
+calls, including `return_all_solutions=True`. Calls are serialized by a lock and
+CUDA events; Warp kernels use the current Torch stream. Buffer storage is released
+with the solver. UR retains all 512 periodic candidates and the existing nearest
+selection; OPW retains eight candidates. OPW packs live joint limits in one host
+transfer per call, so limit updates do not require cache invalidation.

agent_context/topics/robot-system/robot-system.md

@@ -47,6 +47,9 @@ path retains independent per-target seeds and does not require this capability.
 - `motion` subpackages load lazily, and workspace analyzer/visualization APIs
   retain a separate lazy export boundary. Do not introduce offline analysis
   imports into the Robot initialization path.
+- `compute_batch_fk`/`compute_batch_ik` broadcast per-environment root transforms
+  and delegate batch shape handling to solver adapters; see
+  [solver batch contracts](../ik-solvers/ik-solvers.md#batch-adapters-and-analytic-scratch-memory).
 - Focused workspace coverage lives under `tests/sim/motion/workspace/`; solver
   tests live under `tests/sim/motion/solvers/`.
 

agent_context/topics/robot-workspace/analysis-and-cache.md

@@ -4,7 +4,7 @@
 
 ### Offline analysis, cache, and runtime path
 
-1. `WorkspaceAnalysisConfig` chooses `JOINT_SPACE`, `CARTESIAN_SPACE`, or `PLANE_SAMPLING`,
+1. `WorkspaceAnalyzerConfig` chooses `JOINT_SPACE`, `CARTESIAN_SPACE`, or `PLANE_SAMPLING`,
    plus sampler seed/count/batch size, cache, constraints, visualization, metrics, IK
    seeds/reference pose, control part, and plane settings.
 2. Analysis is designed for one simulation environment; with multiple environments it warns
@@ -14,7 +14,7 @@
 3. Joint-space mode samples within selected joint limits and applies FK; stored
    `workspace_points` and `joint_configurations` contain only valid aligned rows.
 4. Cartesian mode samples requested XYZ bounds, or infers bounds from 1,000 random FK samples
-   plus margin/fallback, then runs seeded IK. It stores all sampled points and reachability
+   plus margin/fallback using one batched FK call, then runs seeded IK. It stores all sampled points and reachability
    outputs plus best joint configurations aligned to reachable points.
 5. Plane mode creates samples on the configured plane and follows the same IK/result alignment path.
 6. `WorkspaceAnalyzer.analyze()` checks `ResultsCache` before computation unless force is set,
@@ -72,3 +72,33 @@
 | Sampled pose is reachable but motion collides | Workspace is kinematic only; add motion planning/collision validation at the consuming layer. |
 | `workspace-cache list/clean` cannot see ResultsCache entries | It targets legacy session caches; use analyzer preview/cache directory semantics for result entries. |
 | Concurrent analyzers expose corrupt/incomplete cache | `ResultsCache.save()` writes result then metadata directly with no observed lock or temp-file replacement. Add atomic write/lock coverage before relying on shared concurrent writers. |
+
+### Sampling and allocation controls
+
+- `SamplingConfig` and `analyze-workspace --sampler` default to scrambled Sobol.
+  Explicit random/grid selections remain available. Sobol draws continue a
+  dimension-specific sequence; recreate the sampler to replay it. Samplers use
+  private RNGs and do not reset the application's Torch random state.
+- `sample_within_constraints=True` (`--sample-within-constraints`) refills random,
+  Sobol or LHS proposals inside the permitted domain. Cartesian box/ground bounds
+  are intersected before generation; sphere proposals use a volume-uniform
+  transform; plane proposals are filtered after projection. Exclusion zones are
+  respected. Exhausting `max_sampling_rounds` raises instead of returning an
+  incomplete cache. Reachability is conditional on this permitted domain, so its
+  percentage is not interchangeable with the default full-domain statistic.
+- `retain_diagnostics=False` (`--compact-results`) retains reachable positions,
+  aligned qpos and scores, while dropping all-target points and masks from the
+  returned result and archive. Counts and metrics remain available. Runtime and
+  preview accept both formats. This reduces retained/output storage; analysis
+  still materializes the original sampled domain before compaction.
+- IK batches allocate target poses once across point/seed axes. One seed bypasses
+  seed-axis reductions and selection indexing. The analyzer never changes the
+  process-wide logging level. Analyzer FK/IK explicitly selects environment zero
+  and defaults to the robot device when no simulation manager is supplied.
+- Cache keys include a sampling revision, compact/domain settings, geometry and
+  excluded zones. These changes intentionally invalidate older result entries.
+- Focused logic tests: `tests/sim/motion/workspace/test_analysis_batching.py`.
+  Real analytic solver tests (including CUDA stream reuse) are in
+  `tests/sim/motion/solvers/test_analytic_batching.py`. The renderer-free benchmark
+  is `scripts/benchmark/workspace_analyzer/benchmark_robot_workspace.py`; it uses
+  real asset chains and Robot methods with fixed root snapshots and eager FK.

agent_context/topics/robot-workspace/robot-workspace.md

@@ -10,7 +10,7 @@ Paths below are relative to `embodichain/lab/sim/` unless qualified.
 
 | Request | Owner |
 |---|---|
-| Offline analysis | `workspace/analyzer.py`: `WorkspaceAnalyzer`, `WorkspaceAnalysisConfig` |
+| Offline analysis | `workspace/analyzer.py`: `WorkspaceAnalyzer`, `WorkspaceAnalyzerConfig` |
 | Result-cache identity and persistence | `workspace/caches/results_cache.py`: `ResultsCache` |
 | Cache loading / point or voxel sampling | `workspace/runtime.py`: `RobotWorkspace`, `WorkspaceSample` |
 | Runtime binding config | `workspace/cfg.py`: `RobotWorkspaceCfg`; `cfg.py`: `RobotCfg.workspace_cfg` |

embodichain/lab/scripts/analyze_workspace.py

@@ -419,6 +419,8 @@ def build_analyzer_config(
         constraint=constraint,
         ik_samples_per_point=args.ik_samples_per_point,
         control_part_name=control_part_name,
+        retain_diagnostics=not getattr(args, "compact_results", False),
+        sample_within_constraints=getattr(args, "sample_within_constraints", False),
     )
 
     if mode == AnalysisMode.PLANE_SAMPLING:
@@ -565,6 +567,8 @@ def _preview_points_and_colors(
         points = np.asarray(arrays["workspace_points"])
     elif "all_points" in arrays:
         points = np.asarray(arrays["all_points"])
+    elif "reachable_points" in arrays:
+        points = np.asarray(arrays["reachable_points"])
     else:
         points = np.asarray(arrays[next(iter(arrays))])
 
@@ -904,8 +908,8 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
         "--sampler",
         type=str,
         choices=["random", "sobol", "halton", "lhs", "uniform", "gaussian"],
-        default="random",
-        help="Sampling strategy (default: random).",
+        default="sobol",
+        help="Sampling strategy (default: sobol).",
     )
     sampling.add_argument(
         "--seed", type=int, default=42, help="Random seed (default: 42)."
@@ -914,6 +918,17 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
         "--batch-size", type=int, default=1000, help="FK/IK batch size (default: 1000)."
     )
 
+    sampling.add_argument(
+        "--sample-within-constraints",
+        action="store_true",
+        help="Refill samples inside the permitted domain; changes reachability denominator.",
+    )
+    sampling.add_argument(
+        "--compact-results",
+        action="store_true",
+        help="Store reachable points, qpos and aligned scores without rejected-point diagnostics.",
+    )
+
     # --- Workspace / plane --------------------------------------------------
     space = parser.add_argument_group("Workspace bounds & plane")
     space.add_argument(

embodichain/lab/sim/motion/solvers/_buffers.py

@@ -0,0 +1,92 @@
+# ----------------------------------------------------------------------------
+# 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.
+# ----------------------------------------------------------------------------
+
+"""Bounded analytic IK scratch storage with caller-owned results."""
+
+from __future__ import annotations
+
+from contextlib import contextmanager, nullcontext
+from functools import wraps
+from threading import RLock
+
+import torch
+import warp as wp
+
+__all__: list[str] = []
+
+
+class _IKBuffers:
+    """Reuse candidate storage, serializing callers across threads and streams."""
+
+    def __init__(self, device: torch.device) -> None:
+        self.device = torch.device(device)
+        self.capacity = 0
+        self.arrays: dict[str, torch.Tensor] = {}
+        self.lock = RLock()
+        self.event: torch.cuda.Event | None = None
+
+    def reserve(self, max_batch: int) -> None:
+        if max_batch < 1:
+            raise ValueError("max_batch must be positive")
+        with self.lock:
+            self.capacity = max(self.capacity, max_batch)
+
+    @contextmanager
+    def borrow(self):
+        with self.lock:
+            stream = (
+                torch.cuda.current_stream(self.device)
+                if self.device.type == "cuda"
+                else None
+            )
+            if stream is not None and self.event is not None:
+                stream.wait_event(self.event)
+            scope = (
+                wp.ScopedStream(wp.stream_from_torch(stream))
+                if stream is not None
+                else nullcontext()
+            )
+            with scope:
+                try:
+                    yield
+                finally:
+                    if stream is not None:
+                        self.event = torch.cuda.Event()
+                        self.event.record(stream)
+
+    def zeros(self, name: str, batch: int, width: int, dtype: torch.dtype) -> wp.array:
+        self.capacity = max(self.capacity, batch)
+        size = self.capacity * width
+        tensor = self.arrays.get(name)
+        if tensor is None or tensor.numel() < size:
+            tensor = torch.empty(size, dtype=dtype, device=self.device)
+            self.arrays[name] = tensor
+        if self.device.type == "cuda":
+            tensor.record_stream(torch.cuda.current_stream(self.device))
+        view = tensor[: batch * width]
+        view.zero_()
+        return wp.from_torch(view)
+
+
+def _with_ik_buffers(method):
+    """Keep candidate buffers borrowed until result copies have been queued."""
+
+    @wraps(method)
+    def wrapped(self, *args, **kwargs):
+        with self._ik_buffers.borrow():
+            return method(self, *args, **kwargs)
+
+    return wrapped

embodichain/lab/sim/motion/solvers/base_solver.py

@@ -29,6 +29,8 @@
 
 from embodichain.lab.sim.utility.solver_utils import create_pk_serial_chain
 
+from ._buffers import _IKBuffers
+
 __all__ = ["BaseSolver", "SolverCfg"]
 
 
@@ -157,6 +159,8 @@ def __init__(self, cfg: SolverCfg = None, device: str = None, **kwargs):
         else:
             self.device = device
 
+        self._ik_buffers = _IKBuffers(self.device)
+
         self.urdf_path = cfg.urdf_path
 
         self.joint_names = cfg.joint_names
@@ -220,6 +224,57 @@ def __init__(self, cfg: SolverCfg = None, device: str = None, **kwargs):
 
         self._init_qpos_limits()
 
+    def prepare_buffers(self, max_batch: int) -> None:
+        """Reserve reusable analytic IK scratch capacity.
+
+        Args:
+            max_batch: Maximum flattened target count, including IK seeds.
+                Storage is allocated lazily by analytic backends and grows when
+                necessary. Returned FK/IK tensors remain owned by the caller.
+        """
+        self._ik_buffers.reserve(max_batch)
+
+    def get_fk_batch(self, qpos: torch.Tensor) -> torch.Tensor:
+        """Compute chain-root FK while preserving arbitrary leading batch axes.
+
+        Args:
+            qpos: Joint positions with shape ``(..., dof)`` on the solver device.
+
+        Returns:
+            TCP transforms with shape ``(..., 4, 4)``.
+        """
+        if qpos.ndim < 2 or qpos.shape[-1] != self.dof:
+            raise ValueError(f"Expected batched qpos ending in {self.dof} joints")
+        shape = qpos.shape[:-1]
+        return self.get_fk(qpos.reshape(-1, self.dof)).reshape(*shape, 4, 4)
+
+    def get_ik_batch(
+        self, target_xpos: torch.Tensor, qpos_seed: torch.Tensor | None = None
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Compute nearest-solution IK with a uniform batch result contract.
+
+        Args:
+            target_xpos: Chain-root TCP transforms with shape ``(..., 4, 4)``.
+            qpos_seed: Optional seeds with matching leading axes and ``dof`` joints.
+
+        Returns:
+            Boolean success ``(...)`` and joint positions ``(..., dof)``. No
+            candidate axis is exposed; concrete ``get_ik`` APIs stay unchanged.
+        """
+        if target_xpos.ndim < 3 or target_xpos.shape[-2:] != (4, 4):
+            raise ValueError("Expected batched target poses ending in (4, 4)")
+        shape = target_xpos.shape[:-2]
+        if qpos_seed is not None:
+            if qpos_seed.shape != (*shape, self.dof):
+                raise ValueError("Joint seed batch axes must match target poses")
+            qpos_seed = qpos_seed.reshape(-1, self.dof)
+        success, qpos = self.get_ik(
+            target_xpos=target_xpos.reshape(-1, 4, 4),
+            qpos_seed=qpos_seed,
+            return_all_solutions=False,
+        )
+        return success.bool().reshape(shape), qpos.reshape(*shape, self.dof)
+
     def set_ik_nearest_weight(
         self, ik_weight: np.ndarray, joint_ids: np.ndarray | None = None
     ) -> bool: