diff --git a/examples/visualization/curobo_world_obstacles.py b/examples/visualization/curobo_world_obstacles.py index 69634da..19a2754 100644 --- a/examples/visualization/curobo_world_obstacles.py +++ b/examples/visualization/curobo_world_obstacles.py @@ -1,9 +1,9 @@ """Visualize cuRobo world obstacles in the Isaac Sim viewport. -This script reads the static obstacle world that cuRobo uses for collision checking -(as built inside `autosim`'s `CuroboPlanner._initialize_static_world()`), converts -all obstacle poses from the planner's reference frame (robot root) to the world frame, -and draws simple wireframes. +This script reads the obstacle world currently held inside cuRobo's collision checker +(`planner.motion_gen.world_coll_checker.world_model`) and draws simple wireframes. +Before drawing, it synchronizes dynamic rigid object poses so the visualization matches +the obstacle poses currently used by cuRobo for collision checking. Usage ----- @@ -13,10 +13,8 @@ Notes ----- -* cuRobo's `UsdHelper.get_obstacles_from_stage(..., reference_prim_path=robot_prim_path)` - applies an inverse reference transform, so obstacle poses are returned in the - **robot root frame** (not world). We therefore compose with the robot root pose - in world to get world-frame obstacle poses for visualization. +* This script reads from cuRobo's live collision world instead of rebuilding from USD, + so reset-time pose changes for dynamic rigid objects are reflected in the visualization. * Wireframe drawing uses debug lines; keep the app running to inspect the scene. """ @@ -38,7 +36,6 @@ app_launcher = AppLauncher(vars(args_cli)) simulation_app = app_launcher.app - import isaaclab.utils.math as PoseUtils import autosim_examples # noqa: F401 @@ -148,11 +145,11 @@ def _visualize_world_obstacles( max_mesh_vertices: int = 200000, ): planner = pipeline._motion_planner - if not hasattr(planner, "_static_world_config"): - raise RuntimeError( - "CuroboPlanner has no `_static_world_config`. Ensure planner initialized and warmup complete." - ) - world_r = planner._static_world_config # WorldConfig in robot root frame + planner.sync_dynamic_objects() + + world_model = planner.motion_gen.world_coll_checker.world_model + if world_model is None: + raise RuntimeError("cuRobo collision checker has no world model loaded.") robot_root_pose_w = pipeline._robot.data.root_pose_w[env_id].detach() device = robot_root_pose_w.device @@ -160,8 +157,10 @@ def _visualize_world_obstacles( clear_debug_drawing() + # NOTE: poses in the world_model are in the robot-root frame + # Cuboids (OBB pose + dims) - for cub in world_r.cuboid: + for cub in world_model.cuboid: pose_r = _as_pose7(cub.pose, device=device, dtype=dtype) pose_w = _compose_robot_to_world(robot_root_pose_w=robot_root_pose_w, pose_r=pose_r) dims = torch.as_tensor(cub.dims, device=device, dtype=dtype).view(3) @@ -169,7 +168,7 @@ def _visualize_world_obstacles( _draw_oriented_box(pose_w=pose_w, half_dims_xyz=half_dims, color=color, thickness=thickness, z_lift=z_lift) # Spheres (pose center + radius) -> draw as cube approximation - for sph in world_r.sphere: + for sph in world_model.sphere: pose_r = _as_pose7(sph.pose, device=device, dtype=dtype) pose_w = _compose_robot_to_world(robot_root_pose_w=robot_root_pose_w, pose_r=pose_r) r = float(sph.radius) @@ -177,7 +176,7 @@ def _visualize_world_obstacles( _draw_oriented_box(pose_w=pose_w, half_dims_xyz=half_dims, color=color, thickness=thickness, z_lift=z_lift) # Cylinders (pose + radius + height) -> draw as oriented bounding box approximation - for cyl in world_r.cylinder: + for cyl in world_model.cylinder: pose_r = _as_pose7(cyl.pose, device=device, dtype=dtype) pose_w = _compose_robot_to_world(robot_root_pose_w=robot_root_pose_w, pose_r=pose_r) half_dims = torch.tensor( @@ -186,7 +185,7 @@ def _visualize_world_obstacles( _draw_oriented_box(pose_w=pose_w, half_dims_xyz=half_dims, color=color, thickness=thickness, z_lift=z_lift) # Capsules (pose + radius + base/tip) -> draw as oriented bounding box approximation - for cap in world_r.capsule: + for cap in world_model.capsule: pose_r = _as_pose7(cap.pose, device=device, dtype=dtype) pose_w = _compose_robot_to_world(robot_root_pose_w=robot_root_pose_w, pose_r=pose_r) base = torch.as_tensor(cap.base, device=device, dtype=dtype).view(3) @@ -196,7 +195,7 @@ def _visualize_world_obstacles( _draw_oriented_box(pose_w=pose_w, half_dims_xyz=half_dims, color=color, thickness=thickness, z_lift=z_lift) # Mesh obstacles: prefer a tight OBB via cuRobo's `Mesh.get_cuboid()`. - for mesh in world_r.mesh: + for mesh in world_model.mesh: verts = getattr(mesh, "vertices", None) if verts is not None and len(verts) > int(max_mesh_vertices): continue diff --git a/source/autosim/autosim/capabilities/motion_planning/curobo/curobo_planner.py b/source/autosim/autosim/capabilities/motion_planning/curobo/curobo_planner.py index c184dd9..e4734cc 100644 --- a/source/autosim/autosim/capabilities/motion_planning/curobo/curobo_planner.py +++ b/source/autosim/autosim/capabilities/motion_planning/curobo/curobo_planner.py @@ -6,6 +6,7 @@ import torch from curobo.cuda_robot_model.util import load_robot_yaml from curobo.geom.types import WorldConfig +from curobo.rollout.cost.pose_cost import PoseCostMetric from curobo.types.base import TensorDeviceType from curobo.types.file_path import ContentPath from curobo.types.math import Pose @@ -20,6 +21,7 @@ ) from isaaclab.assets import Articulation from isaaclab.envs import ManagerBasedEnv +from isaaclab.utils.math import subtract_frame_transforms from autosim.core.logger import AutoSimLogger @@ -96,6 +98,9 @@ def __init__( # Read static world geometry once self._initialize_static_world() + # Cache for dynamic world synchronization. + self._cached_object_mappings: dict[str, str] | None = None + # Define supported cuRobo primitive types for object discovery and pose synchronization self.primitive_types: list[str] = ["mesh", "cuboid", "sphere", "capsule", "cylinder", "voxel", "blox"] @@ -150,22 +155,102 @@ def _initialize_static_world(self) -> None: env_prim_path = f"/World/envs/env_{self._env_id}" robot_prim_path = self.cfg.robot_prim_path or f"{env_prim_path}/Robot" - only_paths = [f"{env_prim_path}/{sub}" for sub in self.cfg.world_only_subffixes] + world_only_subffixes = self.cfg.world_only_subffixes or [] + only_paths = [f"{env_prim_path}/{sub}" for sub in world_only_subffixes] - ignore_list = [f"{env_prim_path}/{sub}" for sub in self.cfg.world_ignore_subffixes] or [ + world_ignore_subffixes = self.cfg.world_ignore_subffixes or [] + ignore_list = [f"{env_prim_path}/{sub}" for sub in world_ignore_subffixes] or [ f"{env_prim_path}/target", "/World/defaultGroundPlane", "/curobo", ] ignore_list.append(robot_prim_path) - self._static_world_config = self.usd_helper.get_obstacles_from_stage( + world_cfg = self.usd_helper.get_obstacles_from_stage( only_paths=only_paths, reference_prim_path=robot_prim_path, ignore_substring=ignore_list, ) - self._static_world_config = self._static_world_config.get_collision_check_world() + self._static_world_config = world_cfg.get_collision_check_world() self.motion_gen.update_world(self._static_world_config) + self._invalidate_object_mapping_cache() + + def _get_object_mappings(self) -> dict[str, list[str]]: + """Map IsaacLab scene object names to cuRobo world obstacle names. + + Returns: + Dictionary mapping IsaacLab scene object names to list of cuRobo world obstacle names. + """ + + if self._cached_object_mappings is not None: + return self._cached_object_mappings + + world_model = self.motion_gen.world_coll_checker.world_model + rigid_objects = self._env.scene.rigid_objects + + env_prefix = f"/World/envs/env_{self._env_id}/Scene" + world_object_paths: list[str] = [] + for primitive_type in self.primitive_types: + primitive_list = getattr(world_model, primitive_type, None) + if not primitive_list: + continue + for primitive in primitive_list: + primitive_name = primitive.name + if env_prefix in primitive_name: + world_object_paths.append(primitive_name) + + mappings: dict[str, list[str]] = {} + for object_name in rigid_objects.keys(): + mappings[object_name] = [] + for world_path in world_object_paths: + if world_path.startswith(f"{env_prefix}/{object_name}"): + mappings[object_name].append(world_path) + + self._cached_object_mappings = mappings + self._logger.debug(f"Object mappings built: {mappings}") + return mappings + + def _invalidate_object_mapping_cache(self) -> None: + """Invalidate cached object-name mapping used for dynamic sync.""" + + self._cached_object_mappings = None + + def sync_dynamic_objects(self) -> int: + """Synchronize dynamic object poses into cuRobo world model. + + Returns: + Number of obstacles whose pose was updated. + """ + + object_mappings = self._get_object_mappings() + if not object_mappings: + return 0 + + rigid_objects = self._env.scene.rigid_objects + robot_root_pos_in_world, robot_root_quat_in_world = self._robot.data.root_pos_w, self._robot.data.root_quat_w + + updated_count = 0 + for object_name, world_obstacle_names in object_mappings.items(): + obj = rigid_objects[object_name] + # NOTE: cuRobo world model is in the robot-root frame + obj_pos_in_world, obj_quat_in_world = obj.data.root_pos_w, obj.data.root_quat_w + obj_pos_in_robot_root, obj_quat_in_robot_root = subtract_frame_transforms( + robot_root_pos_in_world, robot_root_quat_in_world, obj_pos_in_world, obj_quat_in_world + ) + obj_pose = Pose( + position=self._to_curobo_device(obj_pos_in_robot_root[self._env_id]), + quaternion=self._to_curobo_device(obj_quat_in_robot_root[self._env_id]), + ) + for world_obstacle_name in world_obstacle_names: + self.motion_gen.world_coll_checker.update_obstacle_pose( + world_obstacle_name, + obj_pose, + env_idx=self._env_id, + update_cpu_reference=True, + ) + updated_count += 1 + + return updated_count def plan_motion( self, @@ -189,6 +274,10 @@ def plan_motion( JointState of the trajectory or None if planning failed """ + # Dynamic object pose sync before planning (cheap, incremental). + if self.cfg.enable_dynamic_world_sync: + self.sync_dynamic_objects() + if current_qd is None: current_qd = torch.zeros_like(current_q) dof_needed = len(self.target_joint_names) @@ -233,11 +322,26 @@ def plan_motion( for link_name, pose in link_goals.items() } + # Build per-call plan config: clone only when we need to attach a pose_cost_metric + # so the shared self.plan_config is never mutated. + if self.cfg.reach_partial_pose_weight is not None: + weights = torch.tensor( + self.cfg.reach_partial_pose_weight, + device=self.tensor_args.device, + dtype=self.tensor_args.dtype, + ) + pose_metric = PoseCostMetric(reach_partial_pose=True, reach_vec_weight=weights) + active_plan_config = self.plan_config.clone() + active_plan_config.pose_cost_metric = pose_metric + self._logger.debug(f"reach_partial_pose_weight applied: {self.cfg.reach_partial_pose_weight}") + else: + active_plan_config = self.plan_config + # execute planning result = self.motion_gen.plan_single( current_joint_state.unsqueeze(0), goal, - self.plan_config, + active_plan_config, link_poses=link_poses, ) @@ -382,6 +486,82 @@ def solve_ik_batch( } return self.motion_gen.ik_solver.solve_batch(goal, link_poses=link_poses) + def plan_to_joint_config( + self, + target_q: torch.Tensor, + current_q: torch.Tensor, + current_qd: torch.Tensor | None = None, + ) -> JointState | None: + """Plan a joint-space trajectory to a target joint configuration. + + Unlike plan_motion which targets a Cartesian pose, this plans directly + in joint space using cuRobo's plan_single_js. Used by RetractSkill to + move the robot to a predefined retract configuration. + + Args: + target_q: Target joint positions, shape [dof]. + current_q: Current joint positions, shape [dof]. + current_qd: Current joint velocities, shape [dof]. Defaults to zeros. + + Returns: + JointState of the trajectory or None if planning failed. + """ + + if current_qd is None: + current_qd = torch.zeros_like(current_q) + + dof_needed = len(self.target_joint_names) + for name, q in [("current_q", current_q), ("target_q", target_q)]: + if len(q) < dof_needed: + pad = torch.zeros(dof_needed - len(q), dtype=q.dtype) + q = torch.concatenate([q, pad], axis=0) + elif len(q) > dof_needed: + q = q[:dof_needed] + if name == "current_q": + current_q = q + else: + target_q = q + + if len(current_qd) < dof_needed: + current_qd = torch.concatenate( + [current_qd, torch.zeros(dof_needed - len(current_qd), dtype=current_qd.dtype)] + ) + elif len(current_qd) > dof_needed: + current_qd = current_qd[:dof_needed] + + start_state = JointState( + position=self._to_curobo_device(current_q), + velocity=self._to_curobo_device(current_qd) * 0.0, + acceleration=self._to_curobo_device(current_qd) * 0.0, + jerk=self._to_curobo_device(current_qd) * 0.0, + joint_names=self.target_joint_names, + ) + goal_state = JointState( + position=self._to_curobo_device(target_q), + velocity=self._to_curobo_device(torch.zeros_like(target_q)), + acceleration=self._to_curobo_device(torch.zeros_like(target_q)), + jerk=self._to_curobo_device(torch.zeros_like(target_q)), + joint_names=self.target_joint_names, + ) + + start_state = start_state.get_ordered_joint_state(self.target_joint_names) + goal_state = goal_state.get_ordered_joint_state(self.target_joint_names) + + result = self.motion_gen.plan_single_js( + start_state.unsqueeze(0), + goal_state.unsqueeze(0), + self.plan_config.clone(), + ) + + if result.success.item(): + current_plan = result.get_interpolated_plan() + motion_plan = current_plan.get_ordered_joint_state(self.target_joint_names) + self._logger.debug(f"joint-space planning succeeded with {len(motion_plan.position)} waypoints") + return motion_plan + else: + self._logger.warning(f"joint-space planning failed: {result.status}") + return None + def reset(self): """reset the planner state""" diff --git a/source/autosim/autosim/capabilities/motion_planning/curobo/curobo_planner_cfg.py b/source/autosim/autosim/capabilities/motion_planning/curobo/curobo_planner_cfg.py index fb8daa7..6ec0b91 100644 --- a/source/autosim/autosim/capabilities/motion_planning/curobo/curobo_planner_cfg.py +++ b/source/autosim/autosim/capabilities/motion_planning/curobo/curobo_planner_cfg.py @@ -46,7 +46,9 @@ class CuroboPlannerCfg: """Maximum number of planning attempts.""" time_dilation_factor: float = 0.5 """Time dilation factor for planning.""" - + reach_partial_pose_weight: list[float] | None = None + """Per-axis weights [rx, ry, rz, px, py, pz] for partial-pose reaching via cuRobo PoseCostMetric. + Setting a weight to 0.0 relaxes that axis (e.g. [0,0,0,1,1,1] for position-only reaching).""" # Optional prim path configuration robot_prim_path: str | None = None """Absolute USD prim path to the robot root for world extraction; None derives it from environment root.""" @@ -55,6 +57,10 @@ class CuroboPlannerCfg: world_ignore_subffixes: list[str] | None = None """List of subffixes to ignore when extracting world obstacles.""" + # World update strategy + enable_dynamic_world_sync: bool = False + """If True, synchronize dynamic object poses into cuRobo world before planning (fast incremental update).""" + # Debug and visualization debug_planner: bool = False """Enable detailed motion planning debug information.""" diff --git a/source/autosim/autosim/core/pipeline.py b/source/autosim/autosim/core/pipeline.py index bf06cda..c81475e 100644 --- a/source/autosim/autosim/core/pipeline.py +++ b/source/autosim/autosim/core/pipeline.py @@ -96,6 +96,9 @@ def initialize(self) -> None: # save generated actions self._generated_actions = [] + # full-size action buffer (action_space dims), used as base for every step + self._last_action = torch.zeros(self._env.action_space.shape, device=self._env.device) + # set the initialized flag self._initialized = True @@ -201,10 +204,12 @@ def _execute_single_skill(self, skill: Skill, goal: SkillGoal) -> tuple[bool, in output = skill.step(world_state) - action = torch.zeros(self._env.action_space.shape, device=self._env.device) - action[self._env_id, :] = self._action_adapter.apply(skill, output, self._env) + adapter_result = self._action_adapter.apply(skill, output, self._env) + action = self._last_action.clone() + action[self._env_id, : adapter_result.shape[0]] = adapter_result self._env.step(action) + self._last_action = action self._generated_actions.append(action) steps += 1 diff --git a/source/autosim/autosim/decomposers/llm_decomposer/llm_decomposer.py b/source/autosim/autosim/decomposers/llm_decomposer/llm_decomposer.py index 4d25bba..793a720 100644 --- a/source/autosim/autosim/decomposers/llm_decomposer/llm_decomposer.py +++ b/source/autosim/autosim/decomposers/llm_decomposer/llm_decomposer.py @@ -83,19 +83,27 @@ def decompose(self, extra_info: EnvExtraInfo) -> DecomposeResult: prompt = self._build_prompt(task_code, extra_info) self._logger.debug(f"prompt for llm composer: \n{prompt}") - self._logger.info("generate response from llm...") - response = self._llm_backend.generate( - prompt=prompt, temperature=self.cfg.temperature, max_tokens=self.cfg.max_tokens - ) - - # parse json response - try: - results = self._extract_json(response) - self._validate_result(results) - - return from_dict(DecomposeResult, results) - except json.JSONDecodeError as e: - raise ValueError(f"Failed to parse JSON response: {e}\nResponse: {response}") + max_retries = self.cfg.max_decompose_retries + last_error: Exception | None = None + valid_objects = set(extra_info.objects) if extra_info.objects else None + + for attempt in range(1, max_retries + 1): + self._logger.info(f"generate response from llm (attempt {attempt}/{max_retries})...") + response = self._llm_backend.generate( + prompt=prompt, temperature=self.cfg.temperature, max_tokens=self.cfg.max_tokens + ) + + try: + results = self._extract_json(response) + self._validate_result(results, valid_objects) + return from_dict(DecomposeResult, results) + except (json.JSONDecodeError, ValueError) as e: + last_error = e + self._logger.warning(f"Decomposition attempt {attempt} failed: {e}") + if attempt < max_retries: + self._logger.info("Retrying...") + + raise ValueError(f"Decomposition failed after {max_retries} attempts. Last error: {last_error}") def _load_task_code(self, task_name: str) -> str: """ @@ -202,12 +210,14 @@ def _extract_json(self, response: str) -> dict: raise json.JSONDecodeError("No valid JSON found in response", response, 0) - def _validate_result(self, result: dict) -> None: + def _validate_result(self, result: dict, valid_objects: set | None = None) -> None: """ Validate decomposition result structure Args: result: Decomposition result dictionary + valid_objects: Set of valid object names from the scene. If provided, target_object + fields are checked against this set (skills without a target, e.g. retract, are skipped). Raises: ValueError: If validation fails @@ -230,8 +240,15 @@ def _validate_result(self, result: dict) -> None: if field not in result: raise ValueError(f"Missing required field: {field}") - # Validate skill types + # Validate skill types and target objects for subtask in result["subtasks"]: for skill in subtask["skills"]: if skill["skill_type"] not in self._atomic_skills: raise ValueError(f"Invalid skill type: {skill['skill_type']}. Must be one of {self._atomic_skills}") + if valid_objects is not None: + target = skill.get("target_object", "") + if target and target not in valid_objects: + raise ValueError( + f"Invalid target_object '{target}' for skill '{skill['skill_type']}'. " + f"Must be one of: {sorted(valid_objects)}" + ) diff --git a/source/autosim/autosim/decomposers/llm_decomposer/llm_decomposer_cfg.py b/source/autosim/autosim/decomposers/llm_decomposer/llm_decomposer_cfg.py index 7d16f13..8f81a09 100644 --- a/source/autosim/autosim/decomposers/llm_decomposer/llm_decomposer_cfg.py +++ b/source/autosim/autosim/decomposers/llm_decomposer/llm_decomposer_cfg.py @@ -18,7 +18,7 @@ class LLMDecomposerCfg(DecomposerCfg): api_key: str = MISSING """The API key for the LLM.""" - base_url: str = "https://api.chatanywhere.org/v1" # TODO: change here. + base_url: str = "https://api.openai.com/v1" """The base URL for the LLM API.""" model: str = "gpt-3.5-turbo" @@ -30,6 +30,9 @@ class LLMDecomposerCfg(DecomposerCfg): max_tokens: int = 4000 """The maximum number of tokens to generate.""" + max_decompose_retries: int = 3 + """Maximum number of retries if decomposition fails (JSON parse error or validation error).""" + def __post_init__(self) -> None: super().__post_init__() api_key = os.environ.get("AUTOSIM_LLM_API_KEY") diff --git a/source/autosim/autosim/decomposers/llm_decomposer/prompts/task_decompose.jinja b/source/autosim/autosim/decomposers/llm_decomposer/prompts/task_decompose.jinja index 1709b13..f8904fd 100644 --- a/source/autosim/autosim/decomposers/llm_decomposer/prompts/task_decompose.jinja +++ b/source/autosim/autosim/decomposers/llm_decomposer/prompts/task_decompose.jinja @@ -49,7 +49,7 @@ Output the task decomposition result in JSON format with the following fields: - description: Description (string) - skills: Skill sequence, each skill containing: - step: Step number (integer, globally incrementing) - - skill_type: Skill type (must be one of the 6 atomic skills) + - skill_type: Skill type (must be one of the available atomic skills listed above) - target_object: Target object name (string) - target_type: Target type - "object", "fixture", "interactive_element", or "position" - description: Step description (string) @@ -85,6 +85,13 @@ Output the task decomposition result in JSON format with the following fields: 6. **Combined Patterns**: For complex tasks, combine multiple patterns in sequence. +7. **Retract Pattern** (for returning the arm to its home/safe configuration after completing a task): + ``` + ... (task skills) → retract(none) + ``` + Use `retract` when the task explicitly requires the arm to return to a safe/home pose after completion. + The `target_object` field for retract should be set to `"none"` or an empty string since it has no target. + ## Planning Strategy 1. **Start with success conditions**: Check the success conditions first to understand what needs to be achieved 2. **Minimize steps**: Only include skills that are necessary to satisfy the success conditions diff --git a/source/autosim/autosim/skills/__init__.py b/source/autosim/autosim/skills/__init__.py index cb246da..c3b04e7 100644 --- a/source/autosim/autosim/skills/__init__.py +++ b/source/autosim/autosim/skills/__init__.py @@ -12,6 +12,7 @@ PushSkill, PushSkillCfg, ) +from .retract import RetractSkill, RetractSkillCfg @configclass @@ -25,6 +26,7 @@ class AutoSimSkillsExtraCfg: pull: PullSkillCfg = PullSkillCfg() push: PushSkillCfg = PushSkillCfg() reach: ReachSkillCfg = ReachSkillCfg() + retract: RetractSkillCfg = RetractSkillCfg() def get(cls, skill_name: str): """Get the skill configuration by name.""" diff --git a/source/autosim/autosim/skills/reach.py b/source/autosim/autosim/skills/reach.py index d528ca6..80df1d3 100644 --- a/source/autosim/autosim/skills/reach.py +++ b/source/autosim/autosim/skills/reach.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import isaaclab.utils.math as PoseUtils import torch from isaaclab.envs import ManagerBasedEnv @@ -17,11 +19,19 @@ from .base_skill import CuroboSkillBase, CuroboSkillExtraCfg +@configclass +class ReachSkillExtraCfg(CuroboSkillExtraCfg): + """Extra configuration for the reach skill.""" + + corrective_reach: bool = False + """Whether to perform corrective reach.""" + + @configclass class ReachSkillCfg(SkillCfg): """Configuration for the reach skill.""" - extra_cfg: CuroboSkillExtraCfg = CuroboSkillExtraCfg() + extra_cfg: ReachSkillExtraCfg = ReachSkillExtraCfg() """Extra configuration for the reach skill.""" @@ -33,7 +43,7 @@ class ReachSkillCfg(SkillCfg): class ReachSkill(CuroboSkillBase): """Skill to reach to a target object or location""" - def __init__(self, extra_cfg: CuroboSkillExtraCfg) -> None: + def __init__(self, extra_cfg: ReachSkillExtraCfg) -> None: super().__init__(extra_cfg) self._logger = AutoSimLogger("ReachSkill") @@ -42,6 +52,12 @@ def __init__(self, extra_cfg: CuroboSkillExtraCfg) -> None: self._trajectory = None self._step_idx = 0 + self._corrective_reach_done = False + self._saved_env = None + self._saved_target_object = None + self._saved_reach_offset = None + self._saved_env_extra_info = None + def _build_activate_joint_state( self, full_sim_joint_names: list[str], full_sim_q: torch.Tensor, full_sim_qd: torch.Tensor | None = None ) -> tuple[torch.Tensor, torch.Tensor | None]: @@ -120,11 +136,11 @@ def _compute_relative_extra_target_poses( link_pos_in_primary, link_quat_in_primary, ) - self._logger.info( + self._logger.debug( f"Relative offset for {link_name} in primary frame: pos={link_pos_in_primary}," f" quat={link_quat_in_primary}" ) - self._logger.info( + self._logger.debug( f"Target pose for {link_name} in robot root frame: pos={link_target_pos_in_robot_root}," f" quat={link_target_quat_in_robot_root}" ) @@ -225,12 +241,12 @@ def _build_keep_initial_relative_offset_extra_target_poses( env_extra_info.cached_initial_extra_target_offsets = self._compute_relative_offsets_in_primary( primary_link_pose_in_robot_root, extra_link_poses_in_robot_root ) - self._logger.info( + self._logger.debug( "Cached initial relative offsets for extra links:" f" {list(env_extra_info.cached_initial_extra_target_offsets.keys())}" ) else: - self._logger.info( + self._logger.debug( "Reusing cached initial relative offsets for extra links:" f" {list(env_extra_info.cached_initial_extra_target_offsets.keys())}" ) @@ -239,42 +255,47 @@ def _build_keep_initial_relative_offset_extra_target_poses( primary_link_pose_in_robot_root, env_extra_info.cached_initial_extra_target_offsets, target_pose ) - def extract_goal_from_info( - self, skill_info: SkillInfo, env: ManagerBasedEnv, env_extra_info: EnvExtraInfo + def _compute_goal_from_offset( + self, + env: ManagerBasedEnv, + target_object: str, + reach_offset: torch.Tensor, + env_extra_info: EnvExtraInfo, ) -> SkillGoal: - """Return the target pose[x, y, z, qw, qx, qy, qz] in the robot root frame. - IMPORTANT: the robot root frame is not the same as the robot base frame. - """ + """Compute reach goal by transforming object-frame offsets into robot root frame. - target_object = skill_info.target_object - robot = env.scene[env_extra_info.robot_name] + Args: + env: The Isaac Lab environment. + target_object: Name of the target object in the scene. + reach_offset: [7] tensor (pos + quat) in object frame for the primary EE. + env_extra_info: Env info for cfg-driven extra target. + + Returns: + SkillGoal with target poses in robot root frame. + """ object_pose_in_env = env.scene[target_object].data.root_pose_w - object_pos_in_env, object_quat_in_env = object_pose_in_env[:, :3], object_pose_in_env[:, 3:] - - reach_target_pose = env_extra_info.get_next_reach_target_pose(target_object) - reach_target_pose = reach_target_pose.to(env.device) - reach_target_pose_in_object = reach_target_pose.unsqueeze(0) - reach_target_pos_in_object, reach_target_quat_in_object = ( - reach_target_pose_in_object[:, :3], - reach_target_pose_in_object[:, 3:], - ) + object_pos_in_env = object_pose_in_env[:, :3] + object_quat_in_env = object_pose_in_env[:, 3:] + + offset = reach_offset.to(env.device).unsqueeze(0) reach_target_pos_in_env, reach_target_quat_in_env = PoseUtils.combine_frame_transforms( - object_pos_in_env, object_quat_in_env, reach_target_pos_in_object, reach_target_quat_in_object + object_pos_in_env, object_quat_in_env, offset[:, :3], offset[:, 3:] ) - self._logger.info(f"Reach target position in environment: {reach_target_pos_in_env}") - self._logger.info(f"Reach target quaternion in environment: {reach_target_quat_in_env}") + self._logger.debug(f"Reach target position in environment: {reach_target_pos_in_env}") + self._logger.debug(f"Reach target quaternion in environment: {reach_target_quat_in_env}") self._target_poses["target_pose"] = torch.cat((reach_target_pos_in_env, reach_target_quat_in_env), dim=-1) - robot_root_pose_in_env = robot.data.root_pose_w - robot_root_pos_in_env, robot_root_quat_in_env = robot_root_pose_in_env[:, :3], robot_root_pose_in_env[:, 3:] + robot = env.scene[env_extra_info.robot_name] + robot_root_pos_in_env = robot.data.root_pose_w[:, :3] + robot_root_quat_in_env = robot.data.root_pose_w[:, 3:] - reach_target_pos_in_robot_root, reach_target_quat_in_robot_root = PoseUtils.subtract_frame_transforms( + reach_target_pos_in_root, reach_target_quat_in_root = PoseUtils.subtract_frame_transforms( robot_root_pos_in_env, robot_root_quat_in_env, reach_target_pos_in_env, reach_target_quat_in_env ) + target_pose = torch.cat((reach_target_pos_in_root, reach_target_quat_in_root), dim=-1).squeeze(0) - target_pose = torch.cat((reach_target_pos_in_robot_root, reach_target_quat_in_robot_root), dim=-1).squeeze(0) activate_q, _ = self._build_activate_joint_state( robot.data.joint_names, robot.data.joint_pos[0], robot.data.joint_vel[0] ) @@ -282,10 +303,46 @@ def extract_goal_from_info( return SkillGoal(target_object=target_object, target_pose=target_pose, extra_target_poses=extra_target_poses) + def _compute_corrective_goal(self) -> SkillGoal | None: + """Re-compute reach goal using the object's current actual pose. + + This is called after the first trajectory finishes. The same relative offset (in object + frame) is re-applied to the object's current pose, so if the object was nudged during + approach the robot corrects for it. + """ + + goal = self._compute_goal_from_offset( + self._saved_env, + self._saved_target_object, + self._saved_reach_offset, + self._saved_env_extra_info, + ) + if goal is not None: + self._logger.info("corrective_reach: recomputed target from current object pose") + return goal + + def extract_goal_from_info( + self, skill_info: SkillInfo, env: ManagerBasedEnv, env_extra_info: EnvExtraInfo + ) -> SkillGoal: + """Return the target pose[x, y, z, qw, qx, qy, qz] in the robot root frame. + IMPORTANT: the robot root frame is not the same as the robot base frame. + """ + + target_object = skill_info.target_object + reach_offset = env_extra_info.get_next_reach_target_pose(target_object).to(env.device) + + # Save state needed for corrective reach re-planning + self._saved_env = env + self._saved_target_object = target_object + self._saved_reach_offset = reach_offset + self._saved_env_extra_info = env_extra_info + + return self._compute_goal_from_offset(env, target_object, reach_offset, env_extra_info) + def execute_plan(self, state: WorldState, goal: SkillGoal) -> bool: """Execute the plan of the reach skill.""" - self._logger.info(f"Reach from pose in environment: {state.robot_ee_pose}") + self._logger.debug(f"Reach from pose in environment: {state.robot_ee_pose}") target_pose = goal.target_pose # target pose in the robot root frame target_pos, target_quat = target_pose[:3], target_pose[3:] @@ -328,6 +385,24 @@ def step(self, state: WorldState) -> SkillOutput: done = False self._step_idx += 1 + # Corrective reach: when the first trajectory finishes, re-plan using the object's + # actual current position in case it was nudged during approach. + need_corrective_reach = ( + done + and not self._corrective_reach_done + and type(self) is ReachSkill + and self.cfg.extra_cfg.corrective_reach + ) + if need_corrective_reach: + self._corrective_reach_done = True # prevent infinite loop + new_goal = self._compute_corrective_goal() + if new_goal is not None: + self._logger.info("corrective_reach: re-planning to corrected object pose") + self._step_idx = 0 + plan_success = self.execute_plan(state, new_goal) + if plan_success: + done = False # continue with corrective trajectory + curobo_joint_names = self._trajectory.joint_names sim_joint_names = state.sim_joint_names joint_pos = state.robot_joint_pos.clone() @@ -348,3 +423,8 @@ def reset(self) -> None: super().reset() self._step_idx = 0 self._trajectory = None + self._corrective_reach_done = False + self._saved_env = None + self._saved_target_object = None + self._saved_reach_offset = None + self._saved_env_extra_info = None diff --git a/source/autosim/autosim/skills/retract.py b/source/autosim/autosim/skills/retract.py new file mode 100644 index 0000000..6aa080f --- /dev/null +++ b/source/autosim/autosim/skills/retract.py @@ -0,0 +1,143 @@ +import torch +from isaaclab.envs import ManagerBasedEnv +from isaaclab.utils import configclass + +from autosim import register_skill +from autosim.core.logger import AutoSimLogger +from autosim.core.skill import SkillCfg +from autosim.core.types import ( + EnvExtraInfo, + SkillGoal, + SkillInfo, + SkillOutput, + WorldState, +) + +from .base_skill import CuroboSkillBase, CuroboSkillExtraCfg + + +@configclass +class RetractSkillExtraCfg(CuroboSkillExtraCfg): + """Extra configuration for the retract skill.""" + + pass + + +@configclass +class RetractSkillCfg(SkillCfg): + """Configuration for the retract skill.""" + + extra_cfg: RetractSkillExtraCfg = RetractSkillExtraCfg() + + +@register_skill( + name="retract", + cfg_type=RetractSkillCfg, + description="Move robot arm to retract (home) configuration", +) +class RetractSkill(CuroboSkillBase): + """Skill to move the robot arm to a predefined retract configuration. + + The retract configuration is read from the robot's cuRobo yaml file + (cspace.retract_config). This is useful for resetting the arm to a safe + pose between tasks or after a grasp. + """ + + def __init__(self, extra_cfg: RetractSkillExtraCfg) -> None: + super().__init__(extra_cfg) + + self._logger = AutoSimLogger("RetractSkill") + self._trajectory = None + self._step_idx = 0 + + def extract_goal_from_info( + self, skill_info: SkillInfo, env: ManagerBasedEnv, env_extra_info: EnvExtraInfo + ) -> SkillGoal: + """Retract has no target object — return an empty goal.""" + + return SkillGoal() + + def execute_plan(self, state: WorldState, goal: SkillGoal) -> bool: + """Plan a joint-space trajectory to the retract configuration.""" + + retract_q = self._get_retract_config(state) + if retract_q is None: + self._logger.warning("retract_config not found in robot yaml, cannot plan") + return False + + full_sim_joint_names = state.sim_joint_names + full_sim_q = state.robot_joint_pos + full_sim_qd = state.robot_joint_vel + planner_joint_names = self._planner.target_joint_names + + activate_q, activate_qd = [], [] + for joint_name in planner_joint_names: + if joint_name in full_sim_joint_names: + idx = full_sim_joint_names.index(joint_name) + activate_q.append(full_sim_q[idx]) + activate_qd.append(full_sim_qd[idx]) + else: + raise ValueError(f"Joint {joint_name} in planner joints is not in simulation joint names.") + activate_q = torch.stack(activate_q, dim=0) + activate_qd = torch.stack(activate_qd, dim=0) + + self._trajectory = self._planner.plan_to_joint_config( + retract_q, + activate_q, + activate_qd, + ) + + return self._trajectory is not None + + def _get_retract_config(self, state: WorldState) -> torch.Tensor | None: + """Extract retract_config from the loaded robot yaml and reorder to planner joint order.""" + + robot_cfg = self._planner.robot_cfg + try: + cspace = robot_cfg["robot_cfg"]["kinematics"]["cspace"] + yaml_joint_names = cspace["joint_names"] + yaml_retract = cspace["retract_config"] + except (KeyError, TypeError): + return None + + # Build a lookup from yaml joint name → retract value + retract_map = dict(zip(yaml_joint_names, yaml_retract)) + + # Reorder to match planner's joint order; use 0.0 for any missing joint + planner_joint_names = self._planner.target_joint_names + retract_values = [retract_map.get(name, 0.0) for name in planner_joint_names] + + return torch.tensor(retract_values, dtype=torch.float32, device=state.device) + + def step(self, state: WorldState) -> SkillOutput: + """Step through the planned retract trajectory.""" + + traj_positions = self._trajectory.position + if self._step_idx >= len(traj_positions): + traj_pos = traj_positions[-1] + done = True + else: + traj_pos = traj_positions[self._step_idx] + done = False + self._step_idx += 1 + + curobo_joint_names = self._trajectory.joint_names + sim_joint_names = state.sim_joint_names + joint_pos = state.robot_joint_pos.clone() + for curobo_idx, curobo_joint_name in enumerate(curobo_joint_names): + sim_idx = sim_joint_names.index(curobo_joint_name) + joint_pos[sim_idx] = traj_pos[curobo_idx] + + return SkillOutput( + action=joint_pos, + done=done, + success=True, + info={}, + ) + + def reset(self) -> None: + """Reset the retract skill.""" + + super().reset() + self._step_idx = 0 + self._trajectory = None