Particle rendering prototype

scripts/environments/teleoperation/teleop_franka_pour_spacemouse.py

@@ -0,0 +1,223 @@
+# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""SpaceMouse teleoperation for the joint-position Franka Pour task."""
+
+from __future__ import annotations
+
+import argparse
+import sys
+
+import torch
+
+DEFAULT_TASK = "Isaac-Pour-Franka-Teleop-v0"
+
+
+def joint_targets_to_actions(
+    *,
+    joint_targets: torch.Tensor,
+    action_offset: torch.Tensor,
+    action_scale: float | torch.Tensor,
+    lower_limits: torch.Tensor,
+    upper_limits: torch.Tensor,
+) -> torch.Tensor:
+    """Encode bounded joint targets in the task's normalized joint-action coordinates."""
+    joint_targets = torch.clamp(joint_targets, min=lower_limits, max=upper_limits)
+    scale = torch.as_tensor(action_scale, dtype=joint_targets.dtype, device=joint_targets.device)
+    if torch.any(scale == 0.0):
+        raise ValueError("Joint-position action scale must be nonzero.")
+    return (joint_targets - action_offset) / scale
+
+
+def compose_env_action(arm_action: torch.Tensor, gripper_command: torch.Tensor) -> torch.Tensor:
+    """Append the normalized symmetric-gripper command to seven arm joint-position actions."""
+    if gripper_command.ndim == 1:
+        gripper_command = gripper_command.unsqueeze(-1)
+    return torch.cat((arm_action, gripper_command), dim=-1)
+
+
+def apply_tcp_offset_to_jacobian(
+    jacobian: torch.Tensor,
+    body_quat: torch.Tensor,
+    offset_pos: torch.Tensor,
+) -> torch.Tensor:
+    """Move a root-frame geometric Jacobian to a hand-local tool-centre point."""
+    from isaaclab.utils import math as math_utils
+
+    result = jacobian.clone()
+    offset_pos_root = math_utils.quat_apply(body_quat, offset_pos)
+    result[:, :3, :] += torch.bmm(-math_utils.skew_symmetric_matrix(offset_pos_root), result[:, 3:, :])
+    return result
+
+
+def _base_frame_tcp_state_and_jacobian(
+    robot,
+    *,
+    body_idx: int,
+    joint_ids: list[int],
+    offset_pos: torch.Tensor,
+    offset_rot: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Return TCP pose and Jacobian in the robot root frame."""
+    from isaaclab.utils import math as math_utils
+
+    body_pose_w = robot.data.body_link_pose_w.torch[:, body_idx]
+    root_pose_w = robot.data.root_link_pose_w.torch
+    body_pos_b, body_quat_b = math_utils.subtract_frame_transforms(
+        root_pose_w[:, :3],
+        root_pose_w[:, 3:7],
+        body_pose_w[:, :3],
+        body_pose_w[:, 3:7],
+    )
+    tcp_pos_b, tcp_quat_b = math_utils.combine_frame_transforms(
+        body_pos_b,
+        body_quat_b,
+        offset_pos,
+        offset_rot,
+    )
+
+    jacobian_idx = body_idx - 1 if robot.is_fixed_base else body_idx
+    jacobian_joint_ids = [joint_id + robot.num_base_dofs for joint_id in joint_ids]
+    jacobian_w = robot.data.body_link_jacobian_w.torch[:, jacobian_idx, :, jacobian_joint_ids]
+    root_rot_w = math_utils.matrix_from_quat(math_utils.quat_inv(root_pose_w[:, 3:7]))
+    jacobian_b = jacobian_w.clone()
+    jacobian_b[:, :3, :] = torch.bmm(root_rot_w, jacobian_b[:, :3, :])
+    jacobian_b[:, 3:, :] = torch.bmm(root_rot_w, jacobian_b[:, 3:, :])
+    jacobian_b = apply_tcp_offset_to_jacobian(jacobian_b, body_quat_b, offset_pos)
+    return tcp_pos_b, tcp_quat_b, jacobian_b
+
+
+def _build_arg_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description="SpaceMouse teleoperation for Franka Pour.")
+    parser.add_argument("--task", default=DEFAULT_TASK, help="Registered Franka Pour task name.")
+    parser.add_argument("--num_envs", type=int, default=1, help="Number of environments driven by one device.")
+    parser.add_argument("--max_steps", type=int, default=-1, help="Stop after N steps; negative runs until close.")
+    parser.add_argument("--mock", action="store_true", help="Use zero device commands for a hardware-free smoke test.")
+    parser.add_argument("--pos_sensitivity", type=float, default=0.05, help="Translation sensitivity [m].")
+    parser.add_argument("--rot_sensitivity", type=float, default=0.05, help="Rotation sensitivity [rad].")
+    parser.add_argument("--ik_damping", type=float, default=0.05, help="Damped-least-squares regularization.")
+    return parser
+
+
+def main() -> None:
+    """Launch the task and convert SpaceMouse Cartesian deltas to joint-position actions."""
+    import gymnasium as gym
+
+    from isaaclab.app import add_launcher_args, launch_simulation
+    from isaaclab.controllers import DifferentialIKController, DifferentialIKControllerCfg
+    from isaaclab.devices import Se3SpaceMouse, Se3SpaceMouseCfg
+
+    import isaaclab_tasks  # noqa: F401
+    from isaaclab_tasks.utils import resolve_task_config, setup_preset_cli
+
+    parser = _build_arg_parser()
+    add_launcher_args(parser)
+    parser.set_defaults(visualizer=["kit"])
+    args_cli, hydra_args = setup_preset_cli(parser)
+    sys.argv = [sys.argv[0], *hydra_args]
+    env_cfg, _ = resolve_task_config(args_cli.task, "")
+
+    with launch_simulation(env_cfg, args_cli):
+        env_cfg.scene.num_envs = args_cli.num_envs
+        if args_cli.device is not None:
+            env_cfg.sim.device = args_cli.device
+        env = gym.make(args_cli.task, cfg=env_cfg).unwrapped
+
+        try:
+            robot = env.scene["robot"]
+            arm_cfg = env.cfg.actions.arm_action
+            arm_term = env.action_manager.get_term("arm_action")
+            joint_ids, joint_names = robot.find_joints(
+                arm_cfg.joint_names,
+                preserve_order=arm_cfg.preserve_order,
+            )
+            if len(joint_ids) != 7:
+                raise RuntimeError(f"Expected seven Franka arm joints, found {joint_names}.")
+            body_ids, body_names = robot.find_bodies(env.cfg.tcp_body_name)
+            if len(body_ids) != 1:
+                raise RuntimeError(f"Expected one TCP body, found {body_names}.")
+            ik = DifferentialIKController(
+                DifferentialIKControllerCfg(
+                    command_type="pose",
+                    use_relative_mode=True,
+                    ik_method="dls",
+                    ik_params={"lambda_val": args_cli.ik_damping},
+                ),
+                num_envs=env.num_envs,
+                device=env.device,
+            )
+            spacemouse = None
+            if not args_cli.mock:
+                spacemouse = Se3SpaceMouse(
+                    Se3SpaceMouseCfg(
+                        pos_sensitivity=args_cli.pos_sensitivity,
+                        rot_sensitivity=args_cli.rot_sensitivity,
+                        gripper_term=True,
+                        sim_device=env.device,
+                    )
+                )
+            reset_requested = False
+
+            def request_reset() -> None:
+                nonlocal reset_requested
+                reset_requested = True
+
+            if spacemouse is not None:
+                spacemouse.add_callback("R", request_reset)
+            env.reset()
+            ik.reset()
+            if spacemouse is not None:
+                spacemouse.reset()
+                print(spacemouse)
+                print("SpaceMouse drives the TCP through joint-position actions; R resets the environment.")
+
+            offset_pos = torch.tensor(env.cfg.tcp_offset_pos, device=env.device).repeat(env.num_envs, 1)
+            offset_rot = torch.tensor(env.cfg.tcp_offset_rot, device=env.device).repeat(env.num_envs, 1)
+
+            step = 0
+            with torch.inference_mode():
+                while args_cli.max_steps < 0 or step < args_cli.max_steps:
+                    if args_cli.max_steps < 0 and env.sim.visualizers:
+                        if not any(v.is_running() and not v.is_closed for v in env.sim.visualizers):
+                            break
+                    if reset_requested:
+                        env.reset()
+                        if spacemouse is not None:
+                            spacemouse.reset()
+                        ik.reset()
+                        reset_requested = False
+
+                    command = (
+                        spacemouse.advance()
+                        if spacemouse is not None
+                        else torch.zeros(7, device=env.device, dtype=torch.float32)
+                    )
+                    tcp_pos_b, tcp_quat_b, jacobian_b = _base_frame_tcp_state_and_jacobian(
+                        robot,
+                        body_idx=body_ids[0],
+                        joint_ids=joint_ids,
+                        offset_pos=offset_pos,
+                        offset_rot=offset_rot,
+                    )
+                    ik.set_command(command[:6].repeat(env.num_envs, 1), tcp_pos_b, tcp_quat_b)
+                    joint_pos = robot.data.joint_pos.torch[:, joint_ids]
+                    joint_targets = ik.compute(tcp_pos_b, tcp_quat_b, jacobian_b, joint_pos)
+                    joint_limits = robot.data.soft_joint_pos_limits.torch[:, joint_ids]
+                    arm_action = joint_targets_to_actions(
+                        joint_targets=joint_targets,
+                        action_offset=arm_term.action_offset,
+                        action_scale=arm_term.action_scale,
+                        lower_limits=joint_limits[..., 0],
+                        upper_limits=joint_limits[..., 1],
+                    )
+                    gripper_command = command[6].repeat(env.num_envs)
+                    env.step(compose_env_action(arm_action, gripper_command))
+                    step += 1
+        finally:
+            env.close()
+
+
+if __name__ == "__main__":
+    main()

scripts/tools/generate_franka_pour_reset_dataset.py

@@ -0,0 +1,77 @@
+# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Generate an oversampled candidate pool for the Franka Pour reset dataset."""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+from isaaclab.app import add_launcher_args, launch_simulation
+
+import isaaclab_tasks  # noqa: F401
+from isaaclab_tasks.contrib.franka_pour.pour_env import FrankaPourResetSamplerEnv
+from isaaclab_tasks.contrib.franka_pour.reset_dataset_generator import (
+    FRANKA_POUR_RESET_DATASET_TASK_ID,
+    FrankaPourResetDatasetGenerator,
+    FrankaPourResetDatasetGeneratorCfg,
+    save_reset_dataset,
+)
+from isaaclab_tasks.utils.parse_cfg import parse_env_cfg
+
+_REPO_ROOT = Path(__file__).resolve().parents[2]
+_DEFAULT_OUTPUT = _REPO_ROOT / "datasets/franka_pour/reset_dataset_candidates.pt"
+_CANDIDATE_GRASPING_COUNT = 14_000
+_CANDIDATE_NON_GRASPING_COUNT = 12_000
+_CANDIDATE_NEAR_POUR_COUNT = 2_000
+
+
+def main() -> None:
+    """Launch one task world, sample exact category quotas, and atomically save them."""
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--task", default=FRANKA_POUR_RESET_DATASET_TASK_ID)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--batch_size", type=int, default=256)
+    parser.add_argument("--max_attempt_multiplier", type=int, default=100)
+    parser.add_argument("--grasping_count", type=int, default=_CANDIDATE_GRASPING_COUNT)
+    parser.add_argument("--non_grasping_count", type=int, default=_CANDIDATE_NON_GRASPING_COUNT)
+    parser.add_argument("--near_pour_grasp_count", type=int, default=_CANDIDATE_NEAR_POUR_COUNT)
+    parser.add_argument(
+        "--output",
+        type=Path,
+        default=_DEFAULT_OUTPUT,
+        help="Candidate dataset written before dynamic validation.",
+    )
+    add_launcher_args(parser)
+    args = parser.parse_args()
+
+    # One scene world is sufficient: IK and collision candidates are replicated in sampler-owned
+    # Newton models.  This avoids allocating an MPM training batch merely to generate reset data.
+    env_cfg = parse_env_cfg(args.task, device=args.device, num_envs=1)
+    env_cfg.seed = args.seed
+    with launch_simulation(env_cfg, args):
+        env = FrankaPourResetSamplerEnv(env_cfg)
+        try:
+            sampler_cfg = FrankaPourResetDatasetGeneratorCfg(
+                grasping_count=args.grasping_count,
+                non_grasping_count=args.non_grasping_count,
+                near_pour_grasp_count=args.near_pour_grasp_count,
+                batch_size=args.batch_size,
+                seed=args.seed,
+                max_attempt_multiplier=args.max_attempt_multiplier,
+            )
+            payload = FrankaPourResetDatasetGenerator(env, sampler_cfg).generate()
+            save_reset_dataset(payload, args.output)
+            print(f"[INFO] Wrote {payload['metadata']['state_count']} candidate states to {args.output.resolve()}.")
+            print(f"[INFO] Contract SHA-256: {payload['contract_sha256']}")
+            print(f"[INFO] Content SHA-256:  {payload['content_sha256']}")
+            print("[INFO] Next: validate_franka_pour_reset_dataset.py")
+        finally:
+            env.close()
+
+
+if __name__ == "__main__":
+    main()