fix opw solver

[matafela]

Description

• Fix opw cpu batch inverse kinematic qpos seed problem. (use warp cpu)
• Fix opw gpu inverse kinematic.
• Add more test case in opw unittest.

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 focuses on fixing OPW inverse kinematics behavior across CPU/GPU by routing more execution through the Warp-based solver path and expanding the unit tests to exercise batched IK/FK behavior.

Changes:

• Update Warp OPW implementation to normalize angles using trig-based wrapping (instead of modulo).
• Modify the lab OPW solver to initialize/use the Warp solver path (including on CPU) and adjust IK selection weighting defaults.
• Expand OPW solver tests to generate batched joint samples from joint limits and validate batch FK/IK consistency (including CUDA test coverage).

Reviewed changes

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

File	Description
tests/sim/solvers/test_opw_solver.py	Adds grid-based joint sampling and converts the test to batched FK/IK, plus enables CUDA variant.
embodichain/utils/warp/kinematics/opw_solver.py	Changes angle normalization to a Warp-compatible implementation and leaves debug comments in-kernel.
embodichain/lab/sim/solvers/opw_solver.py	Forces Warp solver init path, refactors IK path to always call Warp, and tweaks IK weighting/default handling.

Comments suppressed due to low confidence (5)

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

• The default joint_weight is now a length-6 tensor, but opw_best_ik_kernel currently uses only joint_weights[0] for all joints. This makes per-joint weights ineffective/misleading to callers. Either change the kernel to use joint_weights[t] per DOF, or change the API to accept a single scalar weight (and document it).

        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],
            joint_weight[3],
            joint_weight[4],
            joint_weight[5],
        )

tests/sim/solvers/test_opw_solver.py:37

• Docstring mismatch: steps_per_joint is documented as defaulting to 2 (low/high), but the function signature defaults it to 4. Please update the docstring or the default value so they agree.

    """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.

tests/sim/solvers/test_opw_solver.py:187

• The CUDA test class now always runs self.setup_simulation("cuda") unconditionally. On machines/CI jobs without CUDA available, this will fail the test suite. Consider adding a @pytest.mark.skipif(not torch.cuda.is_available(), ...) guard (or re-adding a skip) so CPU-only environments can still run the test suite.

class TestOPWSolverCUDA(BaseSolverTest):
    def setup_method(self):
        self.setup_simulation("cuda")

embodichain/utils/warp/kinematics/opw_solver.py:440

• There are commented-out debug assignments left in the kernel (qpos[...] = ..., qpos[...] = 0.0). These add noise and can be confusing when maintaining the solver. Please remove the commented-out lines (or gate them behind a proper debug flag if they’re still needed).

            qpos[qpos_start + k] = normalize_to_pi(
                (theta[idx] + offsets[k]) * sign_corrections[k]
            )
            # qpos[qpos_start + k] = (theta[idx] + offsets[k]) * sign_corrections[k]
            # qpos[qpos_start + k] = 0.0

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

• all_qpos / all_ik_valid are converted from Warp arrays to Torch tensors unconditionally, and then converted again inside the if return_all_solutions: block. This duplicates host/device copies and adds overhead on every IK call. Convert only in the branch that actually returns them (or reuse the already-converted tensors).

        all_qpos = wp.to_torch(all_qpos_wp).reshape(n_sample, N_SOL, DOF)
        all_ik_valid = wp.to_torch(all_ik_valid_wp).reshape(n_sample, N_SOL)

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

---

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

[Copilot]

Pull request overview

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

Comments suppressed due to low confidence (3)

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

• joint_weight is treated as a per-joint vector, but opw_best_ik_kernel currently uses only joint_weights[0] for every joint when computing the weighted distance. Changing the default here won’t enable true per-joint weighting until the kernel indexes joint_weights[t]. Consider fixing the kernel (and, if desired, validate joint_weight length/dtype here).

        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],

tests/sim/solvers/test_opw_solver.py:37

• The docstring for steps_per_joint says the default is 2 (low/high), but the function signature defaults it to 4. Please make these consistent (either change the default or update the docstring) to avoid confusion in future test maintenance.

    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.

embodichain/utils/warp/kinematics/opw_solver.py:30

• The TODO comment is now misleading: the modulo-based implementation is commented out, and the function does work in Warp using the atan2(sin, cos) formulation. Please replace the TODO with an explanatory comment (e.g., that Warp doesn’t support % on floats) or remove it to avoid suggesting the function is broken.

    # TODO: Cannot work in warp.
    # return (angle + wp.pi) % (2.0 * wp.pi) - wp.pi
    return wp.atan2(wp.sin(angle), wp.cos(angle))

---

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

[Copilot]

ik_qpos returned from compute_batch_ik here has shape (n_envs, n_batch, dof) and later in this test it is passed as joint_seed to the non-batched compute_ik(...) call. That API expects a per-env seed (n_envs, dof), so this relies on downstream .reshape(-1) behavior and can add a lot of unnecessary data movement. Consider extracting a single seed (e.g., ik_qpos[:, 0, :]) before reusing it as a non-batched seed.

[Copilot]

This test generates up to 65,536 IK/FK samples per arm (steps_per_joint=8 on a 6-DOF grid, then capped). That is a large workload for a unit test and can significantly slow CI or lead to timeouts. Consider lowering steps_per_joint / max_samples, or sampling a smaller randomized subset that still covers the edge cases you care about.

Suggested change

	# generate a small grid of qpos samples from the joint limits (low/high)
	sample_qpos = grid_sample_qpos_from_limits(
	qpos_limit, steps_per_joint=8, device=self.robot.device, max_samples=65536
	# generate a grid of qpos samples from the joint limits (low/high)
	# Use a reduced grid size to keep the unit test lightweight for CI.
	sample_qpos = grid_sample_qpos_from_limits(
	qpos_limit, steps_per_joint=4, device=self.robot.device, max_samples=2048