Fix opw solver

[matafela]

Description

• Fix opw ik solver in cpu mode.
• Fix gpu memory out of bound if directly using opw solver.
• Add safe_margin to OPW solver cfg.
• SolverCfg supports qpos limits now.

TODO:

• Findout what operator in warp cpu failed.

Type of change

• Bug fix (non-breaking change which fixes an issue)

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.

[Copilot]

Pull request overview

This PR adjusts the OPW inverse-kinematics Warp backend to better handle “CPU mode” execution by selecting a separate device for Warp kernel execution and managing tensor/device transfers accordingly.

Changes:

• Introduces a kernel_device to run Warp kernels on CUDA even when self.device is CPU.
• Moves target_xpos / qpos_seed inputs and certain Warp arrays (offsets, sign_corrections) onto kernel_device.
• Moves best-solution IK outputs back to self.device after kernel execution.

Comments suppressed due to low confidence (1)

embodichain/lab/sim/solvers/opw_solver.py:290

• In return_all_solutions mode, all_qpos / all_ik_valid are converted with wp.to_torch(...) but not moved back to self.device. With the new kernel_device logic, CPU solvers will now return CUDA tensors here, while the non-return_all_solutions path returns tensors on self.device. Make the return devices consistent by moving these outputs to self.device as well.

        )

        if return_all_solutions:
            all_qpos = wp.to_torch(all_qpos_wp).reshape(n_sample, N_SOL, DOF)

---

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

[Copilot]

kernel_device is forced to CUDA when self.device is CPU, but there’s no availability check. On machines without CUDA (or where Warp/PyTorch CUDA isn’t available), this will raise at .to(kernel_device) / wp.launch. Consider selecting CUDA only if available (e.g., torch.cuda.is_available() / Warp device query), otherwise keep CPU or raise a clear, actionable error message.

Suggested change

	kernel_device = torch.device("cuda")
	warp_cuda_available = getattr(wp, "is_cuda_available", lambda: False)()
	if torch.cuda.is_available() and warp_cuda_available:
	kernel_device = torch.device("cuda")
	else:
	kernel_device = self.device

[Copilot]

n_sample is computed from target_xpos.shape[0] before the (4, 4) special-case, which makes n_sample==4 for a single pose and causes allocation/launch dimensions to mismatch target_xpos_batch (which becomes (1,4,4)). Set n_sample after batching (e.g., from target_xpos_batch.shape[0]) or explicitly set it to 1 in the (4,4) branch.

[Copilot]

offsets_ = self.offsets.to(...) and sign_corrections_ = ... will potentially allocate/copy Warp arrays on every IK call. If this function is called at high frequency, consider caching device-specific copies (or only copying when kernel_device differs) to reduce per-call overhead.

[Copilot]

PR description/checklist indicates documentation and tests were updated/added, but this PR only changes opw_solver.py and doesn’t modify docs or tests. If tests/docs exist in separate commits/files, please include them here; otherwise update the PR description/checklist to match what’s actually changed.

[Copilot]

Pull request overview

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

Comments suppressed due to low confidence (1)

embodichain/lab/sim/solvers/base_solver.py:254

• set_position_limits() no longer validates that lower_position_limits and upper_position_limits lengths match the solver DOF. Because the subsequent checks use zip(), mismatched inputs will be silently truncated, which can later cause index errors or incorrect limit filtering. Add an explicit length check (preferably against self.dof) and return False / raise on mismatch.

    def set_position_limits(
        self,
        lower_position_limits: List[float],
        upper_position_limits: List[float],
    ) -> bool:
        r"""Sets the upper and lower joint position limits.

        Parameters:
            lower_position_limits (List[float]): A list of lower limits for each joint.
            upper_position_limits (List[float]): A list of upper limits for each joint.

        Returns:
            bool: True if limits are successfully set, False if the input is invalid.
        """

        if any(
            lower > upper
            for lower, upper in zip(lower_position_limits, upper_position_limits)
        ):

---

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

[Copilot]

The default joint_weight uses torch.ones(..., dtype=float), but dtype must be a torch.dtype (e.g., torch.float32). As written, this will raise a TypeError when joint_weight is not provided. Also consider extracting values via .item() (and/or ensuring the tensor is on CPU) before constructing wp_vec6f to avoid issues with 0-dim CUDA tensors.

[Copilot]

normalize_in_limit() takes an upper argument but never uses it, and it only shifts angles to be >= lower (potentially leaving results above upper). This is both confusing and can incorrectly discard otherwise-valid solutions when the joint range is < 2π. Either implement full wrapping into [lower, upper] (using both bounds) or rename/remove the unused parameter to reflect the actual behavior.

Suggested change

	result = angle + k * two_pi
	result = angle + k * two_pi
	if result > upper:
	result = result - two_pi

[Copilot]

Minor readability: high=np.array(UPPER_LIMITS) + -5.1/180*pi is equivalent to subtracting the margin but is easy to misread. Consider changing it to an explicit subtraction for clarity.

Suggested change

	high=np.array(UPPER_LIMITS) + -5.1 / 180.0 * np.pi,
	high=np.array(UPPER_LIMITS) - 5.1 / 180.0 * np.pi,

[Copilot]

n_sample is computed before the (4, 4) special-case handling. If target_xpos is a single pose with shape (4, 4), n_sample becomes 4 (not 1), leading to incorrect buffer sizes and likely out-of-bounds indexing in the kernels. Move the n_sample computation to after target_xpos_batch is created (or explicitly set n_sample = 1 for the (4, 4) case).

[Copilot]

In the pk_serial_chain provided branch, the code ignores the provided chain and instead sets default limits. This is also unsafe when self.dof == 0 (e.g., joint_names=None), because it produces empty limit arrays that will later be indexed as if they had 6 elements (OPW). Consider: if the provided pk_serial_chain exposes low/high, use those; otherwise require explicit DOF/limits or set a non-empty default based on the solver’s DOF.

Suggested change

	# If a pk_serial_chain is provided, we assume it has the necessary information to set position limits.
	self.set_position_limits(
	lower_position_limits=[-2 * np.pi] * self.dof,
	upper_position_limits=[2 * np.pi] * self.dof,
	)
	# If a pk_serial_chain is provided, prefer its own joint limits.
	chain_low = getattr(self.pk_serial_chain, "low", None)
	chain_high = getattr(self.pk_serial_chain, "high", None)
	if chain_low is not None and chain_high is not None and len(chain_low) > 0 and len(chain_high) > 0:
	self.set_position_limits(
	lower_position_limits=chain_low,
	upper_position_limits=chain_high,
	)
	elif self.dof > 0:
	# Fall back to default limits only when the solver DOF is known.
	self.set_position_limits(
	lower_position_limits=[-2 * np.pi] * self.dof,
	upper_position_limits=[2 * np.pi] * self.dof,
	)
	else:
	logger.log_error(
	"A provided pk_serial_chain does not expose non-empty low/high limits, "
	"and the solver DOF could not be inferred from joint_names."
	)
	raise ValueError(
	"Cannot set position limits from pk_serial_chain: missing non-empty low/high "
	"and self.dof == 0. Provide joint_names or a chain with limits."
	)

[Copilot]

If qpos_seed has an unexpected shape, this logs an error but still proceeds to build qpos_seed_wp from qpos_seed_, which will be undefined and raise at runtime. After logging, return/raise (or fall back to a zeros seed) to avoid an UnboundLocalError.

Suggested change

	logger.log_error(
	f"Invalid shape for qpos_seed: {qpos_seed.shape}. Expected ({n_sample}, {DOF}) or ({DOF},)."
	)
	error_msg = (
	f"Invalid shape for qpos_seed: {qpos_seed.shape}. "
	f"Expected ({n_sample}, {DOF}) or ({DOF},)."
	)
	logger.log_error(error_msg)
	raise ValueError(error_msg)

[Copilot]

Pull request overview

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

Comments suppressed due to low confidence (1)

tests/sim/solvers/test_opw_solver.py:39

• The steps_per_joint docstring says it defaults to 2 (low/high), but the function default is 4. Update the docstring (and mention the new safe_margin parameter) so callers know what sampling they get by default.

    steps_per_joint: int = 4,
    device=None,
    max_samples: int = 4096,
    safe_margin: float = 5 / 180 * np.pi,  # 5 degrees in radians
) -> torch.Tensor:
    """Generate grid samples for qpos from qpos_limits.

    Args:
        qpos_limits: tensor of shape (1, n, 2) or (n, 2) where each row is [low, high].
        steps_per_joint: number of values per joint (defaults to 2: low and high).
        device: torch device to place the samples on.
        max_samples: cap the number of returned samples (take first N if grid is larger).

---

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

[Copilot]

In opw_ik_select_kernel, the distance metric multiplies by joint_weights[t] twice (delta * w * delta * w), effectively using w^2 * delta^2. The Python-side scoring logic uses w * delta^2 (single weight factor). Align the Warp kernel with the intended weighting (likely delta*delta*joint_weights[t]) so solution selection matches CPU behavior.

[Copilot]

self._data does not appear to be defined on Robot (this class uses self.body_data elsewhere). This will raise AttributeError during solver initialization. Use self.body_data.qpos_limits (or self.get_qpos_limits(...)) to fetch joint limits instead.

Suggested change

	joint_limits = self._data.qpos_limits[0][joint_ids]
	joint_limits = self.body_data.qpos_limits[0][joint_ids]

[Copilot]

set_qpos_limits() no longer validates that the provided limits match the solver DOF. If the lengths differ, zip() will silently truncate and later indexing (e.g., in OPW kernels) can go out-of-bounds or use wrong limits. Add a length/shape check against self.dof (and return False / raise) before storing the limits.

[Copilot]

qpos_limits is typed as List[float] | None, but it is treated as a 2xDOF structure (e.g., torch.tensor(self.cfg.qpos_limits) and checked against (2, self.dof)). Update the type annotation (and docstring) to reflect the expected nested/2D shape to prevent misuse.

Suggested change

	qpos_limits: List[float] | None = None
	"""Joint position limits [2, DOF] for the solver."""
	qpos_limits: List[List[float]] | None = None
	"""Joint position limits for the solver as a nested list with shape [2, DOF].
	The first row contains the lower limits for each joint, and the second row
	contains the corresponding upper limits.
	"""

[Copilot]

The comment says the config limits are expected to have shape [DOF, 2], but the actual validation enforces (2, self.dof) and later indexes [0]/[1] as (lower, upper). Update the comment to match the enforced shape to avoid misconfigured solvers.

[Copilot]

high=np.array(UPPER_LIMITS) + -5.1/180*π is equivalent to subtracting the margin but is easy to misread. Use an explicit subtraction (and consider hoisting the margin into a named constant) to make the sampling bounds unambiguous.

Suggested change

	qpos_np = np.random.uniform(
	low=np.array(LOWER_LIMITS)
	+ 5.1 / 180.0 * np.pi, # add a margin to avoid sampling near the joint limits
	high=np.array(UPPER_LIMITS) + -5.1 / 180.0 * np.pi,
	joint_limit_margin = 5.1 / 180.0 * np.pi
	qpos_np = np.random.uniform(
	low=np.array(LOWER_LIMITS)
	+ joint_limit_margin, # add a margin to avoid sampling near the joint limits
	high=np.array(UPPER_LIMITS) - joint_limit_margin,

[Copilot]

lower_limits_/upper_limits_ are constructed from elements of self.lower_qpos_limits / self.upper_qpos_limits, which are torch tensors (often on CUDA). Passing 0-d torch tensors into wp_vec6f(...) is likely to fail (or trigger implicit device sync). Convert to Python floats (e.g., .item()) and also add an explicit error if the limits were never initialized (currently this will crash with a non-obvious NoneType/index error).

Suggested change

	lower_limits_ = wp_vec6f(
	self.lower_qpos_limits[0],
	self.lower_qpos_limits[1],
	self.lower_qpos_limits[2],
	self.lower_qpos_limits[3],
	self.lower_qpos_limits[4],
	self.lower_qpos_limits[5],
	)
	upper_limits_ = wp_vec6f(
	self.upper_qpos_limits[0],
	self.upper_qpos_limits[1],
	self.upper_qpos_limits[2],
	self.upper_qpos_limits[3],
	self.upper_qpos_limits[4],
	self.upper_qpos_limits[5],
	if self.lower_qpos_limits is None or self.upper_qpos_limits is None:
	raise ValueError(
	"Joint position limits must be initialized before calling the OPW IK solver."
	)
	lower_limits_ = wp_vec6f(
	float(self.lower_qpos_limits[0].item() if isinstance(self.lower_qpos_limits[0], torch.Tensor) else self.lower_qpos_limits[0]),
	float(self.lower_qpos_limits[1].item() if isinstance(self.lower_qpos_limits[1], torch.Tensor) else self.lower_qpos_limits[1]),
	float(self.lower_qpos_limits[2].item() if isinstance(self.lower_qpos_limits[2], torch.Tensor) else self.lower_qpos_limits[2]),
	float(self.lower_qpos_limits[3].item() if isinstance(self.lower_qpos_limits[3], torch.Tensor) else self.lower_qpos_limits[3]),
	float(self.lower_qpos_limits[4].item() if isinstance(self.lower_qpos_limits[4], torch.Tensor) else self.lower_qpos_limits[4]),
	float(self.lower_qpos_limits[5].item() if isinstance(self.lower_qpos_limits[5], torch.Tensor) else self.lower_qpos_limits[5]),
	)
	upper_limits_ = wp_vec6f(
	float(self.upper_qpos_limits[0].item() if isinstance(self.upper_qpos_limits[0], torch.Tensor) else self.upper_qpos_limits[0]),
	float(self.upper_qpos_limits[1].item() if isinstance(self.upper_qpos_limits[1], torch.Tensor) else self.upper_qpos_limits[1]),
	float(self.upper_qpos_limits[2].item() if isinstance(self.upper_qpos_limits[2], torch.Tensor) else self.upper_qpos_limits[2]),
	float(self.upper_qpos_limits[3].item() if isinstance(self.upper_qpos_limits[3], torch.Tensor) else self.upper_qpos_limits[3]),
	float(self.upper_qpos_limits[4].item() if isinstance(self.upper_qpos_limits[4], torch.Tensor) else self.upper_qpos_limits[4]),
	float(self.upper_qpos_limits[5].item() if isinstance(self.upper_qpos_limits[5], torch.Tensor) else self.upper_qpos_limits[5]),

[Copilot]

Pull request overview

Copilot reviewed 13 out of 13 changed files in this pull request and generated 8 comments.

Comments suppressed due to low confidence (1)

embodichain/lab/sim/solvers/opw_solver.py:334

• joint_weight_wp is constructed from joint_weight[0] etc., which are Torch scalar tensors by default. To avoid type conversion issues when building wp_vec6f (and to keep this device-agnostic), convert these elements to Python floats (e.g., .item()) before passing them into wp_vec6f.

        joint_weight = kwargs.get("joint_weight", torch.ones(size=(DOF,), dtype=float))
        joint_weight_wp = wp_vec6f(
            joint_weight[0],
            joint_weight[1],
            joint_weight[2],

---

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

[Copilot]

user_qpos_limits is described/consumed as a 2D limit matrix ([2, DOF] or [DOF, 2]), but the type annotation is List[float] | None (1D). Update the annotation to a 2D type (e.g., Sequence[Sequence[float]] | torch.Tensor | np.ndarray | None) to reflect the accepted shapes and catch misuse earlier.

[Copilot]

When user_qpos_limits has an invalid shape, the code only logs an error and leaves lower_qpos_limits/upper_qpos_limits as None. Downstream code (e.g., OPW warp IK) indexes these tensors and will crash with a NoneType error. Consider raising a ValueError (or falling back to URDF limits) and returning early so the solver is never left in a partially-initialized state.

Suggested change

	logger.log_error(
	raise ValueError(

[Copilot]

The solver config dict uses qpos_limits, but the new solver config field introduced here is user_qpos_limits. With the current merge logic this may be ignored (or set as an unused dynamic attribute), so the OPW solver won’t actually apply these limits. Rename the key to user_qpos_limits in this test config.

[Copilot]

Same issue for right_arm: the config uses qpos_limits instead of user_qpos_limits, so the limits may not be applied by the solver. Update the key to match the new solver config API.

[Copilot]

This uses high=np.array(UPPER_LIMITS) + -5.1 / 180.0 * np.pi, which is equivalent to subtracting the margin but is easy to misread and looks like a typo. Use - 5.1 / 180.0 * np.pi for clarity (and to reduce the chance of accidentally flipping the sign during future edits).

Suggested change

	high=np.array(UPPER_LIMITS) + -5.1 / 180.0 * np.pi,
	high=np.array(UPPER_LIMITS) - 5.1 / 180.0 * np.pi,

[Copilot]

set_qpos_limits() doesn’t validate that the provided limit arrays have length self.dof. If a caller passes a mismatched length, the solver can end up with incorrect limit tensors and later indexing/clamping bugs. Add an explicit length/shape check (and fail fast) before storing the tensors.

[Copilot]

The SRSSolver config example uses upper_qpos_limits=..., but SRSSolverCfg/SolverCfg exposes user_qpos_limits (and the solver derives lower_qpos_limits/upper_qpos_limits internally). As written, the example won’t run. Update it to user_qpos_limits=arm_params.qpos_limits (or introduce/describe the correct field if different).

Suggested change

	upper_qpos_limits=arm_params.qpos_limits,
	user_qpos_limits=arm_params.qpos_limits,

[Copilot]

update_with_robot_limit() clamps lower_qpos_limits and upper_qpos_limits independently, but doesn’t re-check that the resulting limits still satisfy lower <= upper. If user limits are inconsistent with robot limits, this can produce an impossible interval that later IK code can’t satisfy. Add a final validation step (and error/fallback) after clamping.