Improve Franka Lift resets and continuous Reach tracking

source/isaaclab_tasks/changelog.d/maximiliank-franka-lift-soft.changed.rst

@@ -0,0 +1,3 @@
+* **Breaking:** Changed Franka Reach and Reach-OSC to continuous pose tracking: success remained a
+  reported metric but no longer ended the episode or awarded the terminal success bonus. Episodes
+  ran until timeout. Requalify existing checkpoints because reward totals and episode lengths changed.

source/isaaclab_tasks/changelog.d/maximiliank-franka-lift-soft.fixed.rst

@@ -0,0 +1,5 @@
+* **Breaking:** Corrected rigid Lift reset sampling and success-driven motion regularization without changing
+  Kuka-Allegro rewards. Requalify existing Franka Lift checkpoints because the reset distribution changed.
+  Training and play mode now propose aligned pre-grasps with probability 0.75 before bank rejection and
+  sampling. Reported success covers this mixed reset distribution. For table-only evaluation, set
+  ``env.events.conditional_reset.params.terms.reset_object_to_target.params.probability=0`` before startup.

source/isaaclab_tasks/isaaclab_tasks/core/lift/config/franka/franka_env_cfg.py

@@ -6,6 +6,7 @@
 """Configuration for the Franka lift environment."""
 
 from isaaclab.assets import ArticulationCfg
+from isaaclab.managers import CurriculumTermCfg as CurrTerm
 from isaaclab.managers import EventTermCfg as EventTerm
 from isaaclab.managers import ObservationTermCfg as ObsTerm
 from isaaclab.managers import RewardTermCfg as RewTerm
@@ -80,6 +81,18 @@
 FINGER_SENSORS = [f"{name}_object_s" for name in FINGERTIP_LIST if name != "panda_leftfinger"]
 """Contact sensors of the remaining fingers."""
 
+GRASPABLE_OBJECT_PREGRASPS = [
+    (MeshCuboidCfg(size=(0.05, 0.05, 0.05), **lift.OBJECT_PHYSICS), 0.026, (0.0, 0.0, 0.0, 1.0)),
+    (MeshCuboidCfg(size=(0.025, 0.05, 0.05), **lift.OBJECT_PHYSICS), 0.026, (0.0, 0.0, 0.0, 1.0)),
+    (MeshCuboidCfg(size=(0.025, 0.025, 0.05), **lift.OBJECT_PHYSICS), 0.0135, (0.0, 0.0, 0.0, 1.0)),
+    (MeshCuboidCfg(size=(0.01, 0.05, 0.05), **lift.OBJECT_PHYSICS), 0.026, (0.0, 0.0, 0.0, 1.0)),
+    (MeshSphereCfg(radius=0.02, **lift.OBJECT_PHYSICS), 0.021, (0.0, 0.0, 0.0, 1.0)),
+    (MeshCapsuleCfg(radius=0.025, height=0.1, **lift.OBJECT_PHYSICS), 0.026, (0.0, 2.0**-0.5, 0.0, 2.0**-0.5)),
+    (MeshCapsuleCfg(radius=0.025, height=0.2, **lift.OBJECT_PHYSICS), 0.026, (0.0, 2.0**-0.5, 0.0, 2.0**-0.5)),
+    (MeshCapsuleCfg(radius=0.01, height=0.2, **lift.OBJECT_PHYSICS), 0.011, (0.0, 2.0**-0.5, 0.0, 2.0**-0.5)),
+]
+"""Object shapes, finger openings [m], and object orientations in the hand frame [xyzw]."""
+
 
 ##
 # Scene definition
@@ -115,16 +128,7 @@ def __post_init__(self):
                     filter_prim_paths_expr=["{ENV_REGEX_NS}/Object"],
                 ),
             )
-        graspable_shape_assets_cfg = [
-            MeshCuboidCfg(size=(0.05, 0.05, 0.05), **lift.OBJECT_PHYSICS),
-            MeshCuboidCfg(size=(0.025, 0.05, 0.05), **lift.OBJECT_PHYSICS),
-            MeshCuboidCfg(size=(0.025, 0.025, 0.05), **lift.OBJECT_PHYSICS),
-            MeshCuboidCfg(size=(0.01, 0.05, 0.05), **lift.OBJECT_PHYSICS),
-            MeshSphereCfg(radius=0.02, **lift.OBJECT_PHYSICS),
-            MeshCapsuleCfg(radius=0.025, height=0.1, **lift.OBJECT_PHYSICS),
-            MeshCapsuleCfg(radius=0.025, height=0.2, **lift.OBJECT_PHYSICS),
-            MeshCapsuleCfg(radius=0.01, height=0.2, **lift.OBJECT_PHYSICS),
-        ]
+        graspable_shape_assets_cfg = [clone(shape) for shape, _, _ in GRASPABLE_OBJECT_PREGRASPS]
         self.object.spawn.shapes.assets_cfg = graspable_shape_assets_cfg
         self.object.spawn.default.assets_cfg = graspable_shape_assets_cfg
 
@@ -159,6 +163,9 @@ class FrankaRelJointPosActionCfg:
 class FrankaReorientRewardCfg(lift.RewardsCfg):
     """Reward terms for the MDP, with the Franka finger contact sensors filled in."""
 
+    action_rate = RewTerm(func=mdp.action_rate_l2, weight=-1e-4)
+    joint_vel = RewTerm(func=mdp.joint_vel_l2, weight=-1e-4, params={"asset_cfg": SceneEntityCfg("robot")})
+
     good_finger_contact = RewTerm(
         func=mdp.contacts,
         weight=0.75,
@@ -185,6 +192,28 @@ def __post_init__(self):
         self.success.params["finger_names"] = FINGER_SENSORS
 
 
+@configclass
+class FrankaLiftCurriculumCfg(lift.CurriculumCfg):
+    """Franka-specific motion regularization as grasp difficulty increases."""
+
+    action_rate = CurrTerm(
+        func=mdp.modify_term_cfg,
+        params={
+            "address": "rewards.action_rate.weight",
+            "modify_fn": mdp.difficulty_interpolate_float,
+            "modify_params": {"initial_value": -1e-4, "final_value": -1e-1},
+        },
+    )
+    joint_vel = CurrTerm(
+        func=mdp.modify_term_cfg,
+        params={
+            "address": "rewards.joint_vel.weight",
+            "modify_fn": mdp.difficulty_interpolate_float,
+            "modify_params": {"initial_value": -1e-4, "final_value": -1e-1},
+        },
+    )
+
+
 @configclass
 class FrankaEventCfg(lift.EventCfg):
     """Franka-specific event configuration."""
@@ -224,8 +253,12 @@ def __post_init__(self):
         to_target = reset_terms["reset_object_to_target"].params
         to_target["target_cfg"] = SceneEntityCfg("robot", body_names="panda_hand")
         to_target["pose_range"] = {"x": [-0.02, 0.02], "y": [-0.02, 0.02], "z": [0.08, 0.12]}
-        # every link but the ground-mounted base (a base-link ground check is unsatisfiable)
+        # The ground-mounted base is excluded; all enabled arm and gripper colliders are checked.
         criteria["robot_table_clearance"].body_names = ["panda_link[1-7]", "panda_hand", ".*finger"]
+
+        # Allow prefill even with one environment per shape and a low acceptance rate.
+        self.conditional_reset.params["max_prefill_iters"] = 20_000
+
         # spread the reset bank over the grasp geometry, same bodies as fingers_to_object
         diversity_feature = self.conditional_reset.params.get("diversity_feature")
         if diversity_feature is not None:
@@ -242,7 +275,7 @@ def __post_init__(self):
 
 @configclass
 class FrankaMixinCfg:
-    """Franka-specific scene, observation, action, reward and event terms, mixed into the task configurations."""
+    """Franka scene, observation, action, reward and event terms for the lift task."""
 
     scene: FrankaSceneCfg = FrankaSceneCfg(num_envs=4096, env_spacing=3, replicate_physics=True)
     rewards: FrankaReorientRewardCfg = FrankaReorientRewardCfg()
@@ -255,12 +288,45 @@ def __post_init__(self):
         self.commands.object_pose.body_name = "panda_hand"
         # Franka base is rotated 180 deg about z, so the workspace mirrors to positive x.
         self.commands.object_pose.ranges.pos_x = (0.3, 0.7)
+        # The actuator limits are nominal policy limits, not evidence of unstable physics.
+        # Reserve the abnormal-state termination for velocities beyond the solver contract.
+        self.terminations.abnormal_robot.func = mdp.abnormal_robot_state
         self.terminations.abnormal_robot.params["asset_cfg"] = SceneEntityCfg("robot", joint_names="panda_joint.*")
 
 
 @configclass
 class FrankaLiftEnvCfg(FrankaMixinCfg, lift.LiftEnvCfg):
-    """Franka object lifting environment."""
+    """Franka object lifting with a mixed aligned-grasp and table reset bank.
+
+    Training and play mode propose aligned pre-grasps with probability 0.75 before clearance
+    rejection and bank sampling. Success is measured on this mixture, rather than table-only picks.
+    Set ``events.conditional_reset.params.terms.reset_object_to_target.params.probability=0``
+    before initialization to build a table-only evaluation bank.
+    """
+
+    curriculum: FrankaLiftCurriculumCfg | None = FrankaLiftCurriculumCfg()
+
+    def __post_init__(self):
+        super().__post_init__()
+        reset = self.events.conditional_reset.params
+        terms = reset["terms"]
+
+        # The aligned opening must be written after the generic gripper-width reset.
+        pregrasp = terms.pop("reset_object_to_target")
+        terms["reset_object_to_target"] = pregrasp
+        pregrasp.func = mdp.reset_to_grasp
+        pregrasp.params.update(
+            probability=0.75,
+            gripper_cfg=SceneEntityCfg("robot", joint_names="panda_finger_joint.*"),
+            pose_range={"x": (-0.002, 0.002), "y": (-0.002, 0.002), "z": (0.1, 0.1)},
+            grasp_configs=[
+                (clone(shape), opening, orientation) for shape, opening, orientation in GRASPABLE_OBJECT_PREGRASPS
+            ],
+        )
+        pregrasp.params.pop("velocity_range")
+
+        # Farthest-point thinning discards valid near-grasp starts.
+        reset["diversity_feature"] = None
 
     def play_mode(self):
         super().play_mode()

source/isaaclab_tasks/isaaclab_tasks/core/lift/mdp/__init__.pyi

@@ -31,6 +31,7 @@ __all__ = [
     "deformable_ee_distance",
     "deformable_lifting",
     "deformable_outside_bounds",
+    "difficulty_interpolate_float",
     "ee_below_minimum",
     "fingers_contact_force_b",
     "get_reset_state",
@@ -51,6 +52,7 @@ __all__ = [
     "reset_cable_state_uniform",
     "reset_deformable_over_support",
     "reset_joints_shared_offset",
+    "reset_to_grasp",
     "reset_to_target",
     "set_reset_state",
     "slab_clearance",
@@ -61,14 +63,15 @@ __all__ = [
 from isaaclab_tasks.utils.success_monitor import SuccessMonitor, SuccessMonitorCfg
 
 from .commands import CableUniformPoseCommandCfg, DeformableUniformPoseCommandCfg, ObjectUniformPoseCommandCfg
-from .curriculums import gravity_range_linear
+from .curriculums import difficulty_interpolate_float, gravity_range_linear
 from .events import (
     conditional_reset,
     grasp_travel_distance,
     mesh_clearance,
     reset_cable_state_uniform,
     reset_deformable_over_support,
     reset_joints_shared_offset,
+    reset_to_grasp,
     reset_to_target,
     slab_clearance,
 )

source/isaaclab_tasks/isaaclab_tasks/core/lift/mdp/curriculums.py

@@ -14,6 +14,19 @@
     from isaaclab.envs import ManagerBasedRLEnv
 
 
+def difficulty_interpolate_float(
+    env: ManagerBasedRLEnv,
+    _env_ids: Sequence[int],
+    _data: float,
+    initial_value: float,
+    final_value: float,
+) -> float:
+    """Interpolate a scalar continuously with an ADR term's success-driven difficulty."""
+    difficulty_term = env.curriculum_manager.cfg.adr.func
+    fraction = min(max(difficulty_term.difficulty_frac, 0.0), 1.0)
+    return initial_value + fraction * (final_value - initial_value)
+
+
 def gravity_range_linear(
     env: ManagerBasedRLEnv,
     _env_ids: Sequence[int],

source/isaaclab_tasks/isaaclab_tasks/core/lift/mdp/events.py

@@ -18,7 +18,7 @@
 from isaaclab import cloner
 from isaaclab.managers import EventTermCfg, ManagerTermBase, ManagerTermBaseCfg, SceneEntityCfg
 from isaaclab.utils import instantiate
-from isaaclab.utils.math import quat_apply, random_orientation, sample_uniform, sample_uniform_from_ranges
+from isaaclab.utils.math import quat_apply, quat_mul, random_orientation, sample_uniform, sample_uniform_from_ranges
 
 from isaaclab_tasks.utils.success_monitor import SuccessMonitor, SuccessMonitorCfg
 
@@ -146,6 +146,125 @@ def reset_to_target(
     asset.write_root_velocity_to_sim_index(root_velocity=velocities, env_ids=picked)
 
 
+class reset_to_grasp(ManagerTermBase):
+    """Place selected object variants in aligned parallel-gripper pre-grasps."""
+
+    def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv) -> None:
+        super().__init__(cfg, env)
+
+        asset_cfg: SceneEntityCfg = cfg.params["asset_cfg"]
+        gripper_cfg: SceneEntityCfg = cfg.params["gripper_cfg"]
+        target_cfg: SceneEntityCfg = cfg.params["target_cfg"]
+        self._asset = env.scene[asset_cfg.name]
+        self._gripper = env.scene[gripper_cfg.name]
+        self._target = env.scene[target_cfg.name]
+        self._target_body_ids = target_cfg.body_ids
+
+        object_cfg = getattr(env.cfg.scene, asset_cfg.name)
+        plan = env.scene.clone_plan
+        object_prototypes = cloner.path.get_asset_prototypes(plan, object_cfg.prim_path)
+        if len(object_prototypes) == 0:
+            raise ValueError(f"Could not find clone-plan prototypes for asset '{asset_cfg.name}'.")
+
+        grasp_configs = cfg.params["grasp_configs"]
+        grasp_by_geometry = {_grasp_geometry(shape): index for index, (shape, _, _) in enumerate(grasp_configs)}
+        if len(grasp_by_geometry) != len(grasp_configs):
+            raise ValueError("reset_to_grasp requires one unambiguous pre-grasp per object geometry.")
+
+        variant_ids = np.full(env.num_envs, -1, dtype=np.int64)
+        for prototype_id in object_prototypes:
+            spawn = plan.asset_cfgs[prototype_id].spawn
+            shape = spawn.assets_cfg[0] if isinstance(spawn, sim_utils.MultiAssetSpawnerCfg) else spawn
+            geometry = _grasp_geometry(shape)
+            if geometry not in grasp_by_geometry:
+                raise ValueError(f"reset_to_grasp has no configured pre-grasp for object geometry {geometry}.")
+
+            world_ids, _ = cloner.query.get_asset_prototype_unique_world_index(plan.topology, int(prototype_id))
+            world_ids = world_ids[world_ids >= 0]
+            if (variant_ids[world_ids] >= 0).any():
+                raise ValueError(f"Multiple object variants are assigned to asset '{asset_cfg.name}' in one world.")
+            variant_ids[world_ids] = grasp_by_geometry[geometry]
+        if (variant_ids < 0).any():
+            raise ValueError(f"No object variant is assigned to asset '{asset_cfg.name}' in some worlds.")
+
+        self._variant_ids = torch.as_tensor(variant_ids, device=env.device)
+        self._gripper_joint_ids = self._gripper.find_joints(gripper_cfg.joint_names)[0]
+        self._gripper_joint_positions = torch.tensor([opening for _, opening, _ in grasp_configs], device=env.device)
+        self._asset_orientations = torch.tensor([orientation for _, _, orientation in grasp_configs], device=env.device)
+        self._offset_ranges = torch.tensor(
+            [cfg.params["pose_range"].get(axis, (0.0, 0.0)) for axis in ("x", "y", "z")], device=env.device
+        )
+        self._zero_joint_velocities = torch.zeros((env.num_envs, len(self._gripper_joint_ids)), device=env.device)
+        self._zero_root_velocities = torch.zeros((env.num_envs, 6), device=env.device)
+
+    def __call__(
+        self,
+        env: ManagerBasedEnv,
+        env_ids: Sequence[int] | slice | torch.Tensor,
+        pose_range: dict[str, tuple[float, float]],
+        probability: float,
+        target_cfg: SceneEntityCfg,
+        gripper_cfg: SceneEntityCfg,
+        grasp_configs: list[tuple[sim_utils.SpawnerCfg, float, tuple[float, float, float, float]]],
+        asset_cfg: SceneEntityCfg = SceneEntityCfg("object"),
+    ) -> None:
+        """Reset a fraction of environments to configured pre-grasps.
+
+        Args:
+            env: The environment.
+            env_ids: Environments to reset.
+            pose_range: Object-position offsets in the target frame [m].
+            probability: Per-environment probability of applying the pre-grasp.
+            target_cfg: Body defining the pre-grasp frame.
+            gripper_cfg: Parallel gripper joints to pose.
+            grasp_configs: Shape configs paired with a finger opening [m or rad, depending on joint type]
+                and object orientation in the target frame in ``(x, y, z, w)`` order. Geometry, not list order,
+                associates each pre-grasp with its cloned object. Static values are cached at initialization.
+            asset_cfg: Object asset to reset.
+        """
+        env_ids = env.scene._ALL_INDICES[env_ids]
+        picked = env_ids[torch.rand(len(env_ids), device=env.device) < probability]
+        if len(picked) == 0:
+            return
+
+        variant_ids = self._variant_ids[picked]
+
+        target_pos = self._target.data.body_pos_w.torch[picked][:, self._target_body_ids, :].flatten(1)[:, :3]
+        target_quat = self._target.data.body_quat_w.torch[picked][:, self._target_body_ids, :].flatten(1)[:, :4]
+        local_offsets = sample_uniform(
+            self._offset_ranges[:, 0], self._offset_ranges[:, 1], (len(picked), 3), device=env.device
+        )
+        positions = target_pos + quat_apply(target_quat, local_offsets)
+        local_orientations = self._asset_orientations[variant_ids]
+        orientations = quat_mul(target_quat, local_orientations)
+
+        joint_positions = self._gripper_joint_positions[variant_ids, None].expand(-1, len(self._gripper_joint_ids))
+        self._gripper.write_joint_position_to_sim_index(
+            position=joint_positions, joint_ids=self._gripper_joint_ids, env_ids=picked
+        )
+        self._gripper.write_joint_velocity_to_sim_index(
+            velocity=self._zero_joint_velocities[: len(picked)], joint_ids=self._gripper_joint_ids, env_ids=picked
+        )
+        self._gripper.set_joint_position_target_index(
+            target=joint_positions, joint_ids=self._gripper_joint_ids, env_ids=picked
+        )
+        self._asset.write_root_pose_to_sim_index(root_pose=torch.cat((positions, orientations), dim=-1), env_ids=picked)
+        self._asset.write_root_velocity_to_sim_index(
+            root_velocity=self._zero_root_velocities[: len(picked)], env_ids=picked
+        )
+
+
+def _grasp_geometry(shape: sim_utils.SpawnerCfg) -> tuple:
+    """Identify parallel-gripper geometry independently of mesh authoring and prototype order."""
+    if isinstance(shape, (sim_utils.CuboidCfg, sim_utils.MeshCuboidCfg)):
+        return ("cuboid", *shape.size)
+    if isinstance(shape, (sim_utils.SphereCfg, sim_utils.MeshSphereCfg)):
+        return ("sphere", shape.radius)
+    if isinstance(shape, (sim_utils.CapsuleCfg, sim_utils.MeshCapsuleCfg)):
+        return ("capsule", shape.radius, shape.height, shape.axis)
+    raise ValueError(f"reset_to_grasp does not support shape {type(shape).__name__}.")
+
+
 class conditional_reset(ManagerTermBase):
     """Run wrapped reset terms and guarantee the resulting states satisfy a criterion.
 

source/isaaclab_tasks/isaaclab_tasks/core/lift/mdp/terminations.py

@@ -21,7 +21,7 @@
 
 
 class out_of_bound(ManagerTermBase):
-    """Termination condition for when the object falls out of bound.
+    """Terminate when the rigid object state is non-finite or its position leaves the workspace bounds.
 
     The world-space bounds are cached and rebuilt per axis only when the corresponding
     ``in_bound_range`` entry changes. This keeps the hot path free of host-to-device
@@ -58,7 +58,13 @@ def __call__(
                 self._cached_axis[i] = bounds
 
         pos_w = self._object.data.root_pos_w.torch
-        return ((pos_w < self._lower) | (pos_w > self._upper)).any(dim=1)
+        quat_w = self._object.data.root_quat_w.torch
+        vel_w = self._object.data.root_vel_w.torch
+        # NaNs compare false against both bounds; unstable object states must still terminate.
+        invalid = (
+            ~torch.isfinite(pos_w).all(dim=1) | ~torch.isfinite(quat_w).all(dim=1) | ~torch.isfinite(vel_w).all(dim=1)
+        )
+        return invalid | ((pos_w < self._lower) | (pos_w > self._upper)).any(dim=1)
 
 
 def abnormal_robot_state(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
@@ -67,8 +73,8 @@ def abnormal_robot_state(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = Sce
     Such violations indicate unstable physics, typically caused by aggressive actions.
     """
     robot: Articulation = env.scene[asset_cfg.name]
-    joint_vel = robot.data.joint_vel.torch
-    joint_vel_limits = robot.data.joint_vel_limits.torch
+    joint_vel = robot.data.joint_vel.torch[:, asset_cfg.joint_ids]
+    joint_vel_limits = robot.data.joint_vel_limits.torch[:, asset_cfg.joint_ids]
     return (joint_vel.abs() > (joint_vel_limits * 2)).any(dim=1)