diff --git a/docs/source/_static/css/environment-browser.js b/docs/source/_static/css/environment-browser.js index edb297ec0bf1..5b82864fb67b 100644 --- a/docs/source/_static/css/environment-browser.js +++ b/docs/source/_static/css/environment-browser.js @@ -61,6 +61,7 @@ ["IsaacContrib-AutoMate-Disassembly-Direct", "rl_games", "", "", "", "environment-previews/automate_disassembly.jpg"], ["IsaacContrib-Cartpole-Camera-Showcase-Direct", "skrl", "", "", "box_box,box_discrete,box_multidiscrete,dict_box,dict_discrete,dict_multidiscrete,tuple_box,tuple_discrete,tuple_multidiscrete", "", false, {}, {}, "", null, ["||", "||box_box", "||box_discrete", "||box_multidiscrete", "||dict_box", "||dict_discrete", "||dict_multidiscrete", "||tuple_box", "||tuple_discrete", "||tuple_multidiscrete"]], ["IsaacContrib-Cartpole-Showcase-Direct", "skrl", "isaacsim_physx,newton_kamino,newton_mjwarp,ovphysx", "", "box_box,box_discrete,box_multidiscrete,dict_box,dict_discrete,dict_multidiscrete,discrete_box,discrete_discrete,discrete_multidiscrete,multidiscrete_box,multidiscrete_discrete,multidiscrete_multidiscrete,tuple_box,tuple_discrete,tuple_multidiscrete", "environment-previews/cartpole.jpg", false, {}, {}, "", null, ["||", "||box_box", "||box_discrete", "||box_multidiscrete", "||dict_box", "||dict_discrete", "||dict_multidiscrete", "||discrete_box", "||discrete_discrete", "||discrete_multidiscrete", "||multidiscrete_box", "||multidiscrete_discrete", "||multidiscrete_multidiscrete", "||tuple_box", "||tuple_discrete", "||tuple_multidiscrete", "isaacsim_physx||", "isaacsim_physx||box_box", "isaacsim_physx||box_discrete", "isaacsim_physx||box_multidiscrete", "isaacsim_physx||dict_box", "isaacsim_physx||dict_discrete", "isaacsim_physx||dict_multidiscrete", "isaacsim_physx||discrete_box", "isaacsim_physx||discrete_discrete", "isaacsim_physx||discrete_multidiscrete", "isaacsim_physx||multidiscrete_box", "isaacsim_physx||multidiscrete_discrete", "isaacsim_physx||multidiscrete_multidiscrete", "isaacsim_physx||tuple_box", "isaacsim_physx||tuple_discrete", "isaacsim_physx||tuple_multidiscrete", "newton_kamino||", "newton_kamino||box_box", "newton_kamino||box_discrete", "newton_kamino||box_multidiscrete", "newton_kamino||dict_box", "newton_kamino||dict_discrete", "newton_kamino||dict_multidiscrete", "newton_kamino||discrete_box", "newton_kamino||discrete_discrete", "newton_kamino||discrete_multidiscrete", "newton_kamino||multidiscrete_box", "newton_kamino||multidiscrete_discrete", "newton_kamino||multidiscrete_multidiscrete", "newton_kamino||tuple_box", "newton_kamino||tuple_discrete", "newton_kamino||tuple_multidiscrete", "newton_mjwarp||", "newton_mjwarp||box_box", "newton_mjwarp||box_discrete", "newton_mjwarp||box_multidiscrete", "newton_mjwarp||dict_box", "newton_mjwarp||dict_discrete", "newton_mjwarp||dict_multidiscrete", "newton_mjwarp||discrete_box", "newton_mjwarp||discrete_discrete", "newton_mjwarp||discrete_multidiscrete", "newton_mjwarp||multidiscrete_box", "newton_mjwarp||multidiscrete_discrete", "newton_mjwarp||multidiscrete_multidiscrete", "newton_mjwarp||tuple_box", "newton_mjwarp||tuple_discrete", "newton_mjwarp||tuple_multidiscrete", "ovphysx||", "ovphysx||box_box", "ovphysx||box_discrete", "ovphysx||box_multidiscrete", "ovphysx||dict_box", "ovphysx||dict_discrete", "ovphysx||dict_multidiscrete", "ovphysx||discrete_box", "ovphysx||discrete_discrete", "ovphysx||discrete_multidiscrete", "ovphysx||multidiscrete_box", "ovphysx||multidiscrete_discrete", "ovphysx||multidiscrete_multidiscrete", "ovphysx||tuple_box", "ovphysx||tuple_discrete", "ovphysx||tuple_multidiscrete"]], + ["IsaacContrib-Conveyor-Racetrack-Transfer-v0", "rsl_rl", "", "", "", "./tasks/conveyor/racetrack_transfer.jpg"], ["IsaacContrib-Deploy-GearAssembly-Rizon4s-Grav", "rsl_rl", "", "", "", "environment-previews/rizon4s_gear_assembly.jpg"], ["IsaacContrib-Deploy-GearAssembly-Rizon4s-Grav-ROS-Inference", "rsl_rl", "", "", "", "environment-previews/rizon4s_gear_assembly.jpg"], ["IsaacContrib-Deploy-GearAssembly-UR10e-2F140", "rsl_rl", "", "", "", "environment-previews/ur10e_2f140_gear_assembly.jpg"], diff --git a/docs/source/_static/tasks/conveyor/racetrack_transfer.jpg b/docs/source/_static/tasks/conveyor/racetrack_transfer.jpg new file mode 100644 index 000000000000..3335ad4c15f6 Binary files /dev/null and b/docs/source/_static/tasks/conveyor/racetrack_transfer.jpg differ diff --git a/source/isaaclab_contrib/changelog.d/conveyor-clone-default.fixed.rst b/source/isaaclab_contrib/changelog.d/conveyor-clone-default.fixed.rst new file mode 100644 index 000000000000..b8217db029c1 --- /dev/null +++ b/source/isaaclab_contrib/changelog.d/conveyor-clone-default.fixed.rst @@ -0,0 +1 @@ +* Fixed conveyor namespace defaults to follow the shared clone configuration. diff --git a/source/isaaclab_contrib/changelog.d/maximiliank-conveyor-support.added.rst b/source/isaaclab_contrib/changelog.d/maximiliank-conveyor-support.added.rst new file mode 100644 index 000000000000..6d71071aca0e --- /dev/null +++ b/source/isaaclab_contrib/changelog.d/maximiliank-conveyor-support.added.rst @@ -0,0 +1 @@ +* Added experimental straight and curved conveyor surfaces using Newton's upstream force model. diff --git a/source/isaaclab_contrib/changelog.d/maximiliank-conveyor-support.minor b/source/isaaclab_contrib/changelog.d/maximiliank-conveyor-support.minor new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab_contrib/isaaclab_contrib/conveyors/README.md b/source/isaaclab_contrib/isaaclab_contrib/conveyors/README.md new file mode 100644 index 000000000000..4d7624b34844 --- /dev/null +++ b/source/isaaclab_contrib/isaaclab_contrib/conveyors/README.md @@ -0,0 +1,19 @@ +# Experimental conveyor support + +This adapter stages Newton conveyor integration in `isaaclab_contrib`; its API is +experimental and is not exported from Isaac Lab core or the backend packages. + +- `surface_velocity.SurfaceVelocitySpec` describes a straight or curved surface + in the collision prim's local frame. +- `newton.SurfaceVelocity` uses Newton's upstream `ConveyorForceModel`. The + adapter owns replicated surface/body bindings and solver lifecycle integration, + including hard resets and registration before CUDA graph capture. It supports + the Newton MJWarp solver; other Newton solvers are outside this adapter's scope. + Pass the simulation context's `backend_manager`; solved contact forces use + Newton's solver observables rather than the legacy contact-force attribute. + +Surface speeds and traction parameters come from the authored descriptions. +Full resets restart the startup ramp; episode resets clear stale forces for +only the selected worlds. + +No meshes, assets, task registrations, or policy checkpoints are bundled here. diff --git a/source/isaaclab_contrib/isaaclab_contrib/conveyors/__init__.py b/source/isaaclab_contrib/isaaclab_contrib/conveyors/__init__.py new file mode 100644 index 000000000000..9ceb327848ad --- /dev/null +++ b/source/isaaclab_contrib/isaaclab_contrib/conveyors/__init__.py @@ -0,0 +1,10 @@ +# 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 + +"""Experimental conveyor support for contributed tasks. + +The Newton adapter is imported explicitly from :mod:`.newton` so importing the +surface descriptions does not load the optional backend. +""" diff --git a/source/isaaclab_contrib/isaaclab_contrib/conveyors/newton.py b/source/isaaclab_contrib/isaaclab_contrib/conveyors/newton.py new file mode 100644 index 000000000000..29e426209bec --- /dev/null +++ b/source/isaaclab_contrib/isaaclab_contrib/conveyors/newton.py @@ -0,0 +1,258 @@ +# 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 + +"""Bind authored conveyor surfaces to Newton's upstream force model and solver lifecycle.""" + +from __future__ import annotations + +import re +from collections.abc import Sequence +from typing import Any + +import numpy as np +import torch +import warp as wp +from isaaclab_newton.physics import NewtonManager +from newton.examples.basic.example_basic_conveyor_forces import ConveyorForceModel +from newton.solvers import SolverBase + +from isaaclab.cloner.cloner_cfg import DEFAULT_ENV_TEMPLATE +from isaaclab.physics import PhysicsEvent + +from .surface_velocity import SurfaceVelocitySpec, resolve_surface_velocity_paths + + +@wp.kernel +def _filter_body_forces(body_is_tracked: wp.array[wp.int32], body_force: wp.array[wp.spatial_vector]): + body_id = wp.tid() + if body_is_tracked[body_id] == 0: + body_force[body_id] = wp.spatial_vector() + + +@wp.kernel +def _advance_startup_scale( + dt: wp.float32, duration: wp.float32, elapsed: wp.array[wp.float32], scale: wp.array[wp.float32] +): + elapsed[0] += dt + scale[0] = wp.min(1.0, elapsed[0] / duration) + + +@wp.kernel +def _clear_selected_body_forces( + body_world: wp.array[wp.int32], world_mask: wp.array[wp.bool], body_force: wp.array[wp.spatial_vector] +): + body_id = wp.tid() + world_id = body_world[body_id] + if world_id >= 0 and world_mask[world_id]: + body_force[body_id] = wp.spatial_vector() + + +def _world_vector(transform: wp.transform, vector: tuple[float, float, float]) -> wp.vec3: + """Rotate and normalize an authored direction or surface normal.""" + rotated = wp.transform_vector(transform, wp.vec3(*vector)) + values = np.asarray([float(rotated[index]) for index in range(3)], dtype=np.float32) + values /= float(np.linalg.norm(values)) + return wp.vec3(*(float(value) for value in values)) + + +class SurfaceVelocity: + """Own replicated conveyor bindings from model initialization through graph teardown. + + Register after the simulation context is created and before its first reset. + Speeds come from the authored surface descriptions. Newton computes traction + after each solver substep and applies it on the next one. Hard resets rebuild + the binding before CUDA graph capture; episode resets clear selected-world forces. + """ + + def __init__( + self, + num_envs: int, + surface_specs: Sequence[SurfaceVelocitySpec], + *, + manager: NewtonManager, + body_pattern: str, + body_count_per_env: int, + startup_duration_s: float = 1.0, + env_path_format: str = DEFAULT_ENV_TEMPLATE, + ) -> None: + """Register model and solver callbacks for the authored surfaces. + + Args: + num_envs: Number of replicated simulation environments. + surface_specs: Surface descriptions in within-environment order. + manager: Newton backend manager owned by the simulation context. + body_pattern: Regular expression selecting transported bodies. + body_count_per_env: Expected transported-body count. + startup_duration_s: Initial traction ramp duration [s]. + env_path_format: Clone-template format resolving environment roots. + """ + self._surface_specs = tuple(surface_specs) + self._prim_paths = resolve_surface_velocity_paths(num_envs, self._surface_specs, env_path_format) + self._body_pattern = re.compile(body_pattern) + if not np.isfinite(startup_duration_s) or startup_duration_s <= 0.0: + raise ValueError(f"Conveyor startup duration must be positive, got {startup_duration_s}.") + self._num_envs = num_envs + self._manager = manager + self._body_count_per_env = body_count_per_env + self._startup_duration_s = startup_duration_s + self._model = None + self._conveyor: ConveyorForceModel | None = None + self._closed = False + manager.request_solver_observables(SolverBase.ObservableKind.CONTACT_F) + self._physics_ready_handle = manager.register_callback( + self._queue_solver_binding, PhysicsEvent.PHYSICS_READY, name="surface_velocity_solver_init" + ) + + def _queue_solver_binding(self, _event: Any) -> None: + """Bind after solver allocation on initial construction and every hard reset.""" + self._manager.register_solver_init_callback(self._bind_solver) + + def _bind_solver(self) -> None: + """Resolve the newly finalized model and register its graph callbacks.""" + model = self._manager.get_model() + contacts = self._manager.get_contacts() + observables = self._manager.get_solver_observables() + if contacts is None or observables is None or observables.contact_f is None: + raise RuntimeError("The conveyor requires solved per-contact force observables.") + if model.world_count != self._num_envs: + raise RuntimeError(f"Newton model has {model.world_count} worlds, expected {self._num_envs}.") + + surfaces_per_env = len(self._surface_specs) + surface_shapes = [-1] * len(self._prim_paths) + shape_world = model.shape_world.numpy() + shape_body = model.shape_body.numpy() + for shape_id, label in enumerate(model.shape_label): + world_id = int(shape_world[shape_id]) + if not 0 <= world_id < self._num_envs: + continue + for spec_id in range(surfaces_per_env): + surface_id = world_id * surfaces_per_env + spec_id + path = self._prim_paths[surface_id] + if label != path and not label.startswith(f"{path}/"): + continue + if shape_body[shape_id] >= 0: + raise ValueError(f"Conveyor shape must be static: {label}") + if surface_shapes[surface_id] >= 0: + raise RuntimeError(f"Multiple collision shapes match conveyor surface {path!r}.") + surface_shapes[surface_id] = shape_id + missing = [path for path, shape_id in zip(self._prim_paths, surface_shapes, strict=True) if shape_id < 0] + if missing: + raise RuntimeError(f"Missing {len(missing)} conveyor collision sections: {missing[:8]}.") + + body_is_tracked = np.zeros(model.body_count, dtype=np.int32) + tracked_counts = np.zeros(self._num_envs, dtype=np.int32) + body_world = model.body_world.numpy() + for body_id, label in enumerate(model.body_label): + if self._body_pattern.search(label) is None: + continue + world_id = int(body_world[body_id]) + if not 0 <= world_id < self._num_envs: + raise RuntimeError(f"Transported body {label!r} belongs to invalid world {world_id}.") + body_is_tracked[body_id] = 1 + tracked_counts[world_id] += 1 + if np.any(tracked_counts != self._body_count_per_env): + raise RuntimeError( + f"Body pattern {self._body_pattern.pattern!r} expected {self._body_count_per_env} bodies per world, " + f"got {tracked_counts.tolist()}." + ) + if not np.any(body_is_tracked): + raise RuntimeError(f"Body pattern {self._body_pattern.pattern!r} matched no Newton bodies.") + + conveyor = ConveyorForceModel(model, solver_type="mujoco") + transforms = model.shape_transform.numpy() + for surface_id, shape_id in enumerate(surface_shapes): + spec = self._surface_specs[surface_id % surfaces_per_env] + values = transforms[shape_id] + transform = wp.transform(wp.vec3(*values[:3]), wp.quat(*values[3:])) + direction = np.asarray(_world_vector(transform, spec.direction), dtype=np.float32) + speed = np.float32(spec.velocity) + parameters = { + "surface_normal": _world_vector(transform, spec.surface_normal), + "friction": spec.friction_coefficient, + "threshold": spec.contact_threshold, + } + if spec.curved: + conveyor.add_pivot_belt( + shape_id, + wp.transform_point(transform, wp.vec3(*spec.pivot_point)), + wp.vec3(*(direction * (speed / np.float32(spec.radius)))), + **parameters, + ) + else: + conveyor.add_constant_belt(shape_id, wp.vec3(*(direction * speed)), **parameters) + conveyor.finalize(contacts, observables) + conveyor.set_speed_scale(0.0) + + self._model = model + self._conveyor = conveyor + self._body_is_tracked = wp.array(body_is_tracked, dtype=wp.int32, device=model.device) + self._elapsed_time = wp.zeros(1, dtype=wp.float32, device=model.device) + self._world_mask_host = np.zeros(self._num_envs, dtype=np.bool_) + self._world_mask = wp.zeros(self._num_envs, dtype=wp.bool, device=model.device) + self._manager.register_state_force_callback(self._apply) + self._manager.register_post_solver_substep_callback(self._update) + + def _apply(self, state) -> None: + """Apply the wrench from the preceding solve.""" + self._conveyor.apply(state) + + def _update(self, contacts, observables, state, dt: float) -> None: + """Compute the next wrench using Newton's upstream traction model.""" + wp.launch( + _advance_startup_scale, + dim=1, + inputs=[dt, self._startup_duration_s], + outputs=[self._elapsed_time, self._conveyor.global_velocity_scale], + device=self._model.device, + ) + self._conveyor.update(contacts, observables, state, dt) + wp.launch( + _filter_body_forces, + dim=self._model.body_count, + inputs=[self._body_is_tracked], + outputs=[self._conveyor.conveyor_body_f], + device=self._model.device, + ) + + def reset(self, env_ids: Sequence[int] | torch.Tensor | None = None) -> None: + """Clear selected-world forces; a full reset also restarts the startup ramp.""" + if self._conveyor is None: + raise RuntimeError("Surface velocity is not bound to an initialized Newton solver.") + if env_ids is None: + self._conveyor.conveyor_body_f.zero_() + self._elapsed_time.zero_() + self._conveyor.global_velocity_scale.zero_() + return + ids = env_ids.detach().cpu().numpy() if isinstance(env_ids, torch.Tensor) else np.asarray(env_ids) + if not np.issubdtype(ids.dtype, np.integer): + raise IndexError(f"Surface reset environment indices must be integers, got {env_ids!r}.") + ids = ids.reshape(-1) + if np.any((ids < 0) | (ids >= self._num_envs)): + raise IndexError(f"Conveyor reset environment indices are out of range: {ids.tolist()}.") + self._world_mask_host.fill(False) + self._world_mask_host[ids] = True + if np.all(self._world_mask_host): + self.reset() + return + self._world_mask.assign(self._world_mask_host) + wp.launch( + _clear_selected_body_forces, + dim=self._model.body_count, + inputs=[self._model.body_world, self._world_mask], + outputs=[self._conveyor.conveyor_body_f], + device=self._model.device, + ) + + def close(self) -> None: + """Detach lifecycle and graph callbacks before the model is released.""" + if self._closed: + return + self._manager.unregister_solver_init_callback(self._bind_solver) + self._physics_ready_handle.deregister() + self._manager.unregister_state_force_callback(self._apply) + self._manager.unregister_post_solver_substep_callback(self._update) + self._conveyor = None + self._model = None + self._closed = True diff --git a/source/isaaclab_contrib/isaaclab_contrib/conveyors/surface_velocity.py b/source/isaaclab_contrib/isaaclab_contrib/conveyors/surface_velocity.py new file mode 100644 index 000000000000..639f08796363 --- /dev/null +++ b/source/isaaclab_contrib/isaaclab_contrib/conveyors/surface_velocity.py @@ -0,0 +1,113 @@ +# 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 + +"""Import-light descriptions of prescribed conveyor surface velocities.""" + +from __future__ import annotations + +import math +from collections.abc import Sequence +from dataclasses import dataclass + +from isaaclab.cloner.cloner_cfg import DEFAULT_ENV_TEMPLATE + +_ENV_REGEX_NS = "{ENV_REGEX_NS}" + + +@dataclass(frozen=True, slots=True) +class SurfaceVelocitySpec: + """Persistent intent for one collision surface with prescribed tangential velocity. + + Directions, surface normals, and the optional pivot point are expressed in the collision prim's + local frame. Applied forces and reset state belong to the adapter. + + ``prim_path`` starts with Isaac Lab's ``{ENV_REGEX_NS}/`` placeholder. The adapter + resolves it using the scene's clone template, preserving environment-major surface order. + + Args: + prim_path: Replicated collision-prim template with concrete descendant names. + velocity: Fixed signed centerline surface velocity [m/s]. + direction: Local travel direction, or local rotation axis for a curved belt. + curved: Whether the velocity field rotates about ``pivot_point``. + pivot_point: Local rotation center for a curved belt [m]. + radius: Centerline radius required for a curved belt [m]. + surface_normal: Local outward normal of the carrying surface. + contact_threshold: Minimum contact-normal alignment accepted for traction. + friction_coefficient: Coulomb limit for synthetic surface traction. + """ + + prim_path: str + velocity: float = 0.0 + direction: tuple[float, float, float] = (1.0, 0.0, 0.0) + curved: bool = False + pivot_point: tuple[float, float, float] = (0.0, 0.0, 0.0) + radius: float | None = None + surface_normal: tuple[float, float, float] = (0.0, 0.0, 1.0) + contact_threshold: float = 0.997 + friction_coefficient: float = 0.7 + + def __post_init__(self) -> None: + """Reject invalid authored geometry and contact parameters.""" + if not self.prim_path.startswith(f"{_ENV_REGEX_NS}/") or not all( + part.isidentifier() for part in self.prim_path[len(_ENV_REGEX_NS) + 1 :].split("/") + ): + raise ValueError(f"Surface prim_path must start with '{_ENV_REGEX_NS}/' and name concrete prims.") + for name, vector in ( + ("direction", self.direction), + ("pivot_point", self.pivot_point), + ("surface_normal", self.surface_normal), + ): + if len(vector) != 3 or not all(math.isfinite(value) for value in vector): + raise ValueError(f"Surface {name} must contain three finite values.") + if name != "pivot_point" and math.hypot(*vector) <= 1.0e-8: + raise ValueError(f"Surface {name} must be non-zero.") + if not math.isfinite(self.velocity): + raise ValueError("Surface velocity must be finite.") + if not 0.0 <= self.contact_threshold <= 1.0: + raise ValueError(f"Surface contact_threshold must be in [0, 1], got {self.contact_threshold!r}.") + if not math.isfinite(self.friction_coefficient) or self.friction_coefficient < 0.0: + raise ValueError("Surface friction_coefficient must be finite and non-negative.") + if self.curved and self.radius is None: + raise ValueError("Curved surfaces require a positive radius.") + if self.radius is not None and (not math.isfinite(self.radius) or self.radius <= 0.0): + raise ValueError("Surface radius must be finite and positive.") + + +def resolve_surface_velocity_paths( + num_envs: int, + surface_specs: Sequence[SurfaceVelocitySpec], + env_path_format: str = DEFAULT_ENV_TEMPLATE, +) -> tuple[str, ...]: + """Resolve conveyor templates to exact environment-major prim paths. + + Args: + num_envs: Number of replicated environments. + surface_specs: Within-environment surface descriptions. + env_path_format: Exact environment path format containing one ``{}`` field. + + Returns: + Exact paths ordered by environment, then by ``surface_specs`` order. + + Raises: + ValueError: If surface ownership is ambiguous or the clone template is invalid. + """ + if num_envs < 1: + raise ValueError("Conveyor num_envs must be positive.") + if not surface_specs: + raise ValueError("At least one SurfaceVelocitySpec is required.") + templates = tuple(spec.prim_path for spec in surface_specs) + ordered_templates = sorted(templates) + if any( + right == left or right.startswith(f"{left}/") for left, right in zip(ordered_templates, ordered_templates[1:]) + ): + raise ValueError("Conveyor prim paths must be unique and must not be ancestors of one another.") + if env_path_format.count("{}") != 1: + raise ValueError("Conveyor env_path_format must contain exactly one '{}'.") + env_paths = tuple(env_path_format.format(env_id) for env_id in range(num_envs)) + if any( + not path.startswith("/") or not all(part.isidentifier() for part in path[1:].split("/")) for path in env_paths + ): + raise ValueError("Conveyor env_path_format must produce concrete absolute prim paths.") + return tuple(template.replace(_ENV_REGEX_NS, env_path) for env_path in env_paths for template in templates) diff --git a/source/isaaclab_contrib/test/conveyors/test_surface_velocity_newton.py b/source/isaaclab_contrib/test/conveyors/test_surface_velocity_newton.py new file mode 100644 index 000000000000..dc76d21e55b6 --- /dev/null +++ b/source/isaaclab_contrib/test/conveyors/test_surface_velocity_newton.py @@ -0,0 +1,107 @@ +# 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 + +"""Native transport, reset, and callback-lifetime contracts for the Newton adapter.""" + +import numpy as np +import pytest +import warp as wp +from isaaclab_newton.physics import MJWarpSolverCfg, NewtonBuilderCfg, NewtonCfg + +from isaaclab.sim import SimulationCfg, build_simulation_context +from isaaclab.test.utils import DeviceScope, test_devices + +from isaaclab_contrib.conveyors.newton import SurfaceVelocity +from isaaclab_contrib.conveyors.surface_velocity import SurfaceVelocitySpec + + +@pytest.mark.parametrize("device", test_devices(DeviceScope.CPU_AND_DEFAULT_CUDA)) +def test_transport_reset_rebind_and_close(device): + """Real straight/pivot traction survives hard resets and leaves no eager or captured callbacks.""" + cfg = SimulationCfg( + device=device, + dt=1 / 120, + physics=NewtonCfg( + solver_cfg=MJWarpSolverCfg(use_mujoco_contacts=False, integrator="implicitfast", cone="elliptic"), + use_cuda_graph=device != "cpu", + ), + ) + with build_simulation_context(sim_cfg=cfg) as sim: + manager = sim.backend_manager + builder = sim.get_or_create_backend(NewtonBuilderCfg(physics_cfg=cfg.physics)) + for world in range(2): + builder.begin_world() + root = f"/World/envs/env_{world}" + for name, y in (("B", 0.5), ("A", -0.5)): + builder.add_shape_box( + -1, + xform=wp.transform((0, y, 0), wp.quat_identity()), + hx=1, + hy=0.35, + hz=0.1, + cfg=builder.ShapeConfig(mu=1.1e-5), + label=f"{root}/Belt{name}", + ) + for name, x, y in (("Cube", 0, -0.5 if world == 0 else 0.5), ("Untracked", -0.5, -0.5)): + body = builder.add_link(xform=wp.transform((x, y, 0.21), wp.quat_identity()), label=f"{root}/{name}") + builder.add_shape_box(body, hx=0.1, hy=0.1, hz=0.1, cfg=builder.ShapeConfig(mu=1.1e-5)) + builder.add_articulation([builder.add_joint_free(body)]) + builder.end_world() + callbacks_before = set(manager._callbacks) + driver = SurfaceVelocity( + 2, + ( + SurfaceVelocitySpec("{ENV_REGEX_NS}/BeltA", velocity=0.2), + SurfaceVelocitySpec( + "{ENV_REGEX_NS}/BeltB", + velocity=0.2, + curved=True, + direction=(0, 0, 1), + pivot_point=(-0.5, 0, 0), + radius=0.5, + ), + ), + manager=manager, + body_pattern="/Cube$", + body_count_per_env=1, + startup_duration_s=2.0, + ) + try: + previous_model = None + for _ in range(2): + sim.reset() + model = manager.get_model() + assert model is not previous_model + previous_model = model + before = manager.get_state_0().body_q.numpy().copy() + for _ in range(120): + sim.step(render=False) + after = manager.get_state_0().body_q.numpy() + assert after[0, 0] - before[0, 0] > 0.025 # Straight belt in world 0. + assert after[2, 1] - before[2, 1] > 0.025 # Pivot belt in world 1. + np.testing.assert_allclose(after[[1, 3], :2], before[[1, 3], :2], atol=0.002) + assert np.isfinite(after).all() and np.all(after[:, 2] > 0.15) + + forces = driver._conveyor.conveyor_body_f + forces.assign(np.ones((4, 6), dtype=np.float32)) + driver.reset([0]) + state = model.state() + driver._conveyor.apply(state) + np.testing.assert_array_equal(state.body_f.numpy(), [[0] * 6, [0] * 6, [1] * 6, [1] * 6]) + driver.reset([1, 0]) + np.testing.assert_array_equal(forces.numpy(), np.zeros((4, 6))) + assert driver._conveyor.global_velocity_scale.numpy()[0] == 0 + if device != "cpu": + assert manager._graph is not None + finally: + driver.close() + driver.close() + assert not manager._solver_init_callbacks + assert not manager._state_force_callbacks + assert not manager._post_solver_substep_callbacks + assert set(manager._callbacks) == callbacks_before + forces.fill_(1.0) + sim.step(render=False) + np.testing.assert_array_equal(forces.numpy(), np.ones((4, 6))) diff --git a/source/isaaclab_contrib/test/conveyors/test_surface_velocity_spec.py b/source/isaaclab_contrib/test/conveyors/test_surface_velocity_spec.py new file mode 100644 index 000000000000..9e01c56f442f --- /dev/null +++ b/source/isaaclab_contrib/test/conveyors/test_surface_velocity_spec.py @@ -0,0 +1,56 @@ +# 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 + +"""Focused tests for import-light surface-velocity descriptions.""" + +from __future__ import annotations + +import pytest + +from isaaclab_contrib.conveyors.surface_velocity import SurfaceVelocitySpec, resolve_surface_velocity_paths + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"prim_path": "relative/Belt"}, "prim_path"), + ({"prim_path": "{ENV_REGEX_NS}/Belt.*"}, "prim_path"), + ({"velocity": float("nan")}, "velocity"), + ({"direction": (0.0, 0.0, 0.0)}, "direction"), + ({"surface_normal": (0.0, 0.0)}, "surface_normal"), + ({"pivot_point": (0.0, float("inf"), 0.0)}, "pivot_point"), + ({"curved": True}, "radius"), + ({"radius": 0.0}, "radius"), + ({"contact_threshold": 1.1}, "contact_threshold"), + ({"friction_coefficient": -0.1}, "friction_coefficient"), + ], +) +def test_surface_velocity_spec_rejects_invalid_authored_values(kwargs: dict, message: str) -> None: + """Invalid persistent intent fails before any physics lifecycle is registered.""" + kwargs = {"prim_path": "{ENV_REGEX_NS}/Belt", **kwargs} + with pytest.raises(ValueError, match=message): + SurfaceVelocitySpec(**kwargs) + + +def test_paths_are_resolved_in_environment_major_order() -> None: + """Runtime rows stay deterministic across stage discovery ordering.""" + specs = (SurfaceVelocitySpec("{ENV_REGEX_NS}/Nested/BeltA"), SurfaceVelocitySpec("{ENV_REGEX_NS}/BeltB")) + + paths = resolve_surface_velocity_paths(2, specs, "/Factory/cell_{}") + + assert paths == ( + "/Factory/cell_0/Nested/BeltA", + "/Factory/cell_0/BeltB", + "/Factory/cell_1/Nested/BeltA", + "/Factory/cell_1/BeltB", + ) + + +@pytest.mark.parametrize("paths", [("Belt", "Belt"), ("Belt/Child", "Belt")]) +def test_paths_reject_ambiguous_surface_ownership(paths: tuple[str, str]) -> None: + """One collider cannot be assigned to duplicate or overlapping conveyor roots.""" + specs = tuple(SurfaceVelocitySpec(f"{{ENV_REGEX_NS}}/{path}") for path in paths) + with pytest.raises(ValueError, match="unique.*ancestors"): + resolve_surface_velocity_paths(1, specs) diff --git a/source/isaaclab_newton/changelog.d/maximiliank-conveyor-lifecycle.added.rst b/source/isaaclab_newton/changelog.d/maximiliank-conveyor-lifecycle.added.rst new file mode 100644 index 000000000000..fbe26da77e7e --- /dev/null +++ b/source/isaaclab_newton/changelog.d/maximiliank-conveyor-lifecycle.added.rst @@ -0,0 +1,2 @@ +* Added requested solver observables and post-substep hooks with callback cleanup for scene-owned force systems. +* Exposed collision filtering for static and kinematic shape pairs. diff --git a/source/isaaclab_newton/changelog.d/maximiliank-conveyor-lifecycle.minor b/source/isaaclab_newton/changelog.d/maximiliank-conveyor-lifecycle.minor new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab_newton/isaaclab_newton/physics/newton_collision_cfg.py b/source/isaaclab_newton/isaaclab_newton/physics/newton_collision_cfg.py index 8bb8b75cae0f..305671fa3fa8 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/newton_collision_cfg.py +++ b/source/isaaclab_newton/isaaclab_newton/physics/newton_collision_cfg.py @@ -106,6 +106,13 @@ class NewtonCollisionPipelineCfg: Defaults to ``"explicit"`` (same as Newton's default when ``broad_phase=None``). """ + include_static_kinematic_pairs: bool = True + """Whether to generate static-kinematic and kinematic-kinematic contacts. + + Static-static pairs are always excluded. Set to ``False`` to also exclude + static-kinematic and kinematic-kinematic pairs. Defaults to ``True``, matching Newton's default. + """ + reduce_contacts: bool = True """Whether to reduce contacts for mesh-mesh collisions. diff --git a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py index e522168a28bc..7a19428c66f3 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py +++ b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py @@ -524,6 +524,12 @@ def setup_deformable_body(cls, prim: Any, deformable_type: str, sim_mesh_prim: A _post_actuator_callbacks: list[Callable[[], None]] = [] # In-graph hooks invoked immediately before every solver substep. _state_force_callbacks: list[Callable[[State], None]] = [] + # Hooks run before state swapping and force clearing after each solver substep. + _post_solver_substep_callbacks: list[ + Callable[[Contacts | None, SolverBase.Observables | None, State, float], None] + ] = [] + _requested_solver_observables: set[str] = set() + _solver_observables: SolverBase.Observables | None = None # Body forces written to ``state_0`` before a step. The substep loop clears # the input state's forces, so every substep re-applies this copy. _staged_body_f: wp.array | None = None @@ -627,6 +633,7 @@ def reset(self, soft: bool = False) -> None: self._collision_pipeline = None self._contacts = None self._solver = None + self._solver_observables = None self._eval_fk = _eval_fk_unbound self._reset_solver_internals_delegate = _reset_solver_internals_unbound self._adapter = None @@ -842,6 +849,9 @@ def clear(self): self._pre_actuator_callbacks = [] self._post_actuator_callbacks = [] self._state_force_callbacks = [] + self._post_solver_substep_callbacks = [] + self._requested_solver_observables = set() + self._solver_observables = None self._staged_body_f = None self._post_step_callbacks = [] # Set by an articulation that took the ``use_newton_actuators=True`` @@ -1540,7 +1550,7 @@ def _step_solver( Default invokes :attr:`_solver` once. Subclasses can override to batch multiple solvers within a single substep. """ - self._solver.step(state_0, state_1, control, contacts, substep_dt) + self._solver.step(state_0, state_1, control, contacts, substep_dt, observables=self._solver_observables) def _solver_specific_clear(self) -> None: """Solver-specific cleanup hook called from :meth:`clear`. @@ -1617,6 +1627,11 @@ def initialize_solver(self) -> None: "self._supports_rigid_body_force_input." ) self._initialize_contacts() + self._solver_observables = ( + self._solver.observables(self._requested_solver_observables, contacts=self._contacts) + if self._requested_solver_observables + else None + ) state_0 = self.backend.state_0 self._staged_body_f = wp.zeros_like(state_0.body_f) if state_0.body_count else None @@ -1702,6 +1717,8 @@ def _run_solver_substeps(self, contacts) -> None: for callback in self._state_force_callbacks: callback(backend.state_0) self._step_solver(backend.state_0, backend.state_0, backend.control, contacts, self._solver_dt) + for callback in self._post_solver_substep_callbacks: + callback(contacts, self._solver_observables, backend.state_0, self._solver_dt) backend.state_0.clear_forces() if collide_mid_loop and (i + 1) % collide_every == 0 and i + 1 < self._num_substeps: self._collision_pipeline.collide(backend.state_0, contacts) @@ -1713,6 +1730,8 @@ def _run_solver_substeps(self, contacts) -> None: for callback in self._state_force_callbacks: callback(backend.state_0) self._step_solver(backend.state_0, backend.state_1, backend.control, contacts, self._solver_dt) + for callback in self._post_solver_substep_callbacks: + callback(contacts, self._solver_observables, backend.state_1, self._solver_dt) if need_copy_on_last and i == self._num_substeps - 1: backend.state_0.assign(backend.state_1) else: @@ -1817,6 +1836,14 @@ def get_contacts(self) -> Contacts | None: """Get the current Newton contact buffer, if the active solver exposes one.""" return self._contacts + def request_solver_observables(self, *kinds: str) -> None: + """Request solver outputs before initialization; requests survive hard resets.""" + self._requested_solver_observables.update(kinds) + + def get_solver_observables(self) -> SolverBase.Observables | None: + """Return the outputs populated by each solver substep, or ``None`` when unrequested.""" + return self._solver_observables + def get_num_envs(self) -> int: return self._num_envs @@ -1930,6 +1957,34 @@ def register_state_force_callback(self, callback: Callable[[State], None]) -> No return self._state_force_callbacks.append(callback) + def unregister_state_force_callback(self, callback: Callable[[State], None]) -> None: + """Detach a force callback and recapture an active graph without it.""" + if callback in self._state_force_callbacks: + self._state_force_callbacks.remove(callback) + if self._graph is not None: + self._invalidate_graph() + + def register_post_solver_substep_callback( + self, callback: Callable[[Contacts | None, SolverBase.Observables | None, State, float], None] + ) -> None: + """Register a graph-safe hook after each solve, before state swapping and force clearing.""" + if callback not in self._post_solver_substep_callbacks: + self._post_solver_substep_callbacks.append(callback) + + def unregister_post_solver_substep_callback( + self, callback: Callable[[Contacts | None, SolverBase.Observables | None, State, float], None] + ) -> None: + """Detach a post-solve callback and recapture an active graph without it.""" + if callback in self._post_solver_substep_callbacks: + self._post_solver_substep_callbacks.remove(callback) + if self._graph is not None: + self._invalidate_graph() + + def unregister_solver_init_callback(self, callback: Callable[[], None]) -> None: + """Remove a pending one-shot solver initialization hook, if still registered.""" + with contextlib.suppress(ValueError): + self._solver_init_callbacks.remove(callback) + def register_post_step_callback(self, callback: Callable[[], None]) -> None: """Append a hook to the list invoked after the last solver substep on every step. diff --git a/source/isaaclab_tasks/changelog.d/conveyor-success-guided-sampling.changed.rst b/source/isaaclab_tasks/changelog.d/conveyor-success-guided-sampling.changed.rst new file mode 100644 index 000000000000..420492605bb4 --- /dev/null +++ b/source/isaaclab_tasks/changelog.d/conveyor-success-guided-sampling.changed.rst @@ -0,0 +1,3 @@ +* Reused the NIST success-guided sampler for conveyor reset rows, with completed transfers as the + success signal and standard PPO. Removed the task-specific progress curriculum and checkpoint + extension; sampling history now starts fresh when an environment is recreated. diff --git a/source/isaaclab_tasks/changelog.d/maximiliank-conveyor-racetrack.added.rst b/source/isaaclab_tasks/changelog.d/maximiliank-conveyor-racetrack.added.rst new file mode 100644 index 000000000000..b7b736627707 --- /dev/null +++ b/source/isaaclab_tasks/changelog.d/maximiliank-conveyor-racetrack.added.rst @@ -0,0 +1 @@ +* Added Newton four-cube conveyor racetrack transfer with fixed belt speeds, a pretrained policy, reset curriculum, and seeded evaluation. diff --git a/source/isaaclab_tasks/changelog.d/maximiliank-conveyor-racetrack.minor b/source/isaaclab_tasks/changelog.d/maximiliank-conveyor-racetrack.minor new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/README.md b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/README.md new file mode 100644 index 000000000000..e9811f4614bf --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/README.md @@ -0,0 +1,60 @@ +# Conveyor racetrack transfer + +A Franka continuously transfers four numbered 40 mm cubes between two counter-rotating +racetrack conveyors. The trained policy uses 123 observations and eight actions at 60 Hz, +with 120 Hz physics and a 0.35 m/s belt speed. + +![Racetrack transfer](../../../../../docs/source/_static/tasks/conveyor/racetrack_transfer.jpg) + +Newton traction uses the upstream `ConveyorForceModel`, including its MuJoCo contact +reporting. `isaaclab_contrib.conveyors.newton` adapts replicated surfaces, parcel selection, +configured belt speeds, and reset lifecycle; the task owns its geometry, commands, and reset recipes. + +## Run + +Newton is kitless and supports the lightweight GL viewer: + +```bash +DISPLAY=:1 uv run isaaclab play --rl_library rsl_rl \ + --task IsaacContrib-Conveyor-Racetrack-Transfer-v0 \ + --checkpoint pretrained \ + --num_envs 8 --device cuda:0 --viz newton_gl --real-time +``` + +The `pretrained` selector downloads the RSL-RL policy published specifically for the Newton MJWarp +backend. The Newton checkpoint is published on Isaac Dev. While the public mirror is syncing, +use a local checkpoint path if the download is unavailable. + +Training uses the same task ID and defaults to 256 environments: + +```bash +uv run isaaclab train --rl_library rsl_rl \ + --task IsaacContrib-Conveyor-Racetrack-Transfer-v0 \ + --num_envs 256 --device cuda:0 +``` + +### Curriculum and evaluation + +Racetrack training samples physically calibrated resets from release through moving-belt pickup. +The shared `SuccessMonitor` records whether each episode completed a stable transfer from its +sampled reset row. NIST's reusable `Sampler` mixes uniform coverage with a Beta-shaped score +centered at 50% transfer success, giving partly learned rows more practice while keeping every +row eligible. Reset recipes generate the robot and cube states directly; the task needs no +cached physical-state bank or assembly-variant machinery. + +Training uses standard PPO with the existing arm-and-gripper action distribution. Checkpoints +save policy and optimizer state; sampling history starts fresh when the environment is recreated, +including when resuming training. Playback disables adaptive sampling. Evaluate from moving-belt +starts to measure pickup and transfer ability independently of the easier training resets. + +Evaluate from home on moving belts, with fixed reset settings and a chosen seed: + +```bash +uv run python -m isaaclab_tasks.contrib.conveyor_franka.evaluate \ + --checkpoint /path/to/model.pt --num_envs 8 --device cuda:0 \ + --steps 3600 --seed 0 --output /tmp/conveyor-evaluation.json +``` + +The report counts completed transfers in both directions, safety failures, non-finite observations, +and transfers per simulated environment minute. Repeat with different seeds when comparing policies; +intermediate-reset training reward alone does not establish deployment reliability. diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/__init__.py new file mode 100644 index 000000000000..d7d4cf571772 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/__init__.py @@ -0,0 +1,20 @@ +# 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 + +"""Four-cube racetrack transfer with Newton.""" + +import gymnasium as gym + +from . import agents + +gym.register( + id="IsaacContrib-Conveyor-Racetrack-Transfer-v0", + entry_point=f"{__name__}.conveyor_franka_env:ConveyorFrankaEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.conveyor_franka_env_cfg:ConveyorFrankaEnvCfg", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:ConveyorFrankaPPORunnerCfg", + }, +) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/agents/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/agents/__init__.py new file mode 100644 index 000000000000..0e6e80a26022 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/agents/__init__.py @@ -0,0 +1,6 @@ +# 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 + +"""Agent configurations for conveyor transfer.""" diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/agents/rsl_rl_ppo_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/agents/rsl_rl_ppo_cfg.py new file mode 100644 index 000000000000..c63d3574e8a0 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/agents/rsl_rl_ppo_cfg.py @@ -0,0 +1,183 @@ +# 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 + +"""RSL-RL PPO configuration for conveyor transfer.""" + +import torch +import torch.nn as nn +from rsl_rl.modules.distribution import GaussianDistribution +from torch.distributions import Bernoulli, Normal + +from isaaclab.utils.configclass import configclass + +from isaaclab_rl.rsl_rl import RslRlMLPModelCfg, RslRlOnPolicyRunnerCfg, RslRlPpoAlgorithmCfg + + +class _ConveyorDeterministicOutput(nn.Module): + """Threshold the gripper logit into the action used by the environment.""" + + def forward(self, output: torch.Tensor) -> torch.Tensor: + gripper = torch.where( + output[..., -1:] >= 0.0, + torch.ones_like(output[..., -1:]), + -torch.ones_like(output[..., -1:]), + ) + return torch.cat((output[..., :-1], gripper), dim=-1) + + +class ConveyorGaussianBernoulliDistribution(GaussianDistribution): + """Use Gaussian exploration for the arm and a Bernoulli gripper. + + The final policy output controls a physically binary action. Sampling it + from a Gaussian would assign different log probabilities to values that + become the same open or close command after thresholding. A Bernoulli + instead optimizes exactly the physical decision seen by the environment. + """ + + def __init__( + self, + output_dim: int, + init_std: float = 0.45, + std_range: tuple[float, float] = (0.15, 0.65), + std_type: str = "scalar", + **kwargs, + ) -> None: + if output_dim < 2: + raise ValueError("The conveyor distribution requires arm outputs followed by one gripper output.") + if len(std_range) != 2 or std_range[0] <= 0.0 or std_range[0] >= std_range[1]: + raise ValueError("std_range must contain positive, increasing bounds.") + if not std_range[0] < init_std < std_range[1]: + raise ValueError("init_std must lie strictly inside std_range.") + if std_type != "scalar": + raise ValueError("The conveyor distribution supports only scalar standard-deviation parameters.") + super().__init__(output_dim, init_std=init_std, std_type=std_type, **kwargs) + self.std_range = (float(std_range[0]), float(std_range[1])) + self._arm_distribution: Normal | None = None + self._gripper_distribution: Bernoulli | None = None + initial_fraction = (init_std - self.std_range[0]) / (self.std_range[1] - self.std_range[0]) + initial_logit = torch.logit(torch.tensor(initial_fraction, dtype=self.std_param.dtype)) + with torch.no_grad(): + self.std_param.fill_(initial_logit) + + def update(self, mlp_output: torch.Tensor) -> None: + """Update continuous-arm and binary-gripper distributions.""" + minimum_std, maximum_std = self.std_range + arm_std = minimum_std + (maximum_std - minimum_std) * torch.sigmoid(self.std_param[:-1]) + self._arm_distribution = Normal(mlp_output[..., :-1], arm_std) + self._gripper_distribution = Bernoulli(logits=mlp_output[..., -1:]) + + def sample(self) -> torch.Tensor: + """Sample seven residuals followed by an exact signed binary action.""" + gripper_open = self._gripper_distribution.sample() + return torch.cat((self._arm_distribution.sample(), 2.0 * gripper_open - 1.0), dim=-1) + + def deterministic_output(self, mlp_output: torch.Tensor) -> torch.Tensor: + """Return arm means and the most likely binary gripper command.""" + return _ConveyorDeterministicOutput()(mlp_output) + + def as_deterministic_output_module(self) -> nn.Module: + """Return an exportable deterministic-output transform.""" + return _ConveyorDeterministicOutput() + + @property + def mean(self) -> torch.Tensor: + """Return arm means and the expected signed gripper command.""" + gripper_mean = 2.0 * self._gripper_distribution.probs - 1.0 + return torch.cat((self._arm_distribution.mean, gripper_mean), dim=-1) + + @property + def std(self) -> torch.Tensor: + """Return arm standard deviations and signed-Bernoulli spread.""" + gripper_std = 2.0 * torch.sqrt(self._gripper_distribution.probs * (1.0 - self._gripper_distribution.probs)) + return torch.cat((self._arm_distribution.stddev, gripper_std), dim=-1) + + @property + def entropy(self) -> torch.Tensor: + """Return joint Gaussian-plus-Bernoulli entropy.""" + return self._arm_distribution.entropy().sum(dim=-1) + self._gripper_distribution.entropy().sum(dim=-1) + + @property + def params(self) -> tuple[torch.Tensor, ...]: + """Return parameters needed to evaluate the mixed KL divergence.""" + return self._arm_distribution.mean, self._arm_distribution.stddev, self._gripper_distribution.logits + + def log_prob(self, outputs: torch.Tensor) -> torch.Tensor: + """Evaluate the exact continuous/binary physical action.""" + arm_log_prob = self._arm_distribution.log_prob(outputs[..., :-1]).sum(dim=-1) + gripper_open = (outputs[..., -1:] >= 0.0).to(outputs.dtype) + return arm_log_prob + self._gripper_distribution.log_prob(gripper_open).sum(dim=-1) + + def kl_divergence( + self, + old_params: tuple[torch.Tensor, ...], + new_params: tuple[torch.Tensor, ...], + ) -> torch.Tensor: + """Return ``KL(old || new)`` for both action families.""" + old_arm_mean, old_arm_std, old_gripper_logits = old_params + new_arm_mean, new_arm_std, new_gripper_logits = new_params + arm_kl = torch.distributions.kl_divergence( + Normal(old_arm_mean, old_arm_std), + Normal(new_arm_mean, new_arm_std), + ).sum(dim=-1) + old_gripper_probability = torch.sigmoid(old_gripper_logits) + gripper_kl = ( + old_gripper_probability * (old_gripper_logits - new_gripper_logits) + - torch.nn.functional.softplus(old_gripper_logits) + + torch.nn.functional.softplus(new_gripper_logits) + ).sum(dim=-1) + return arm_kl + gripper_kl.clamp_min(0.0) + + +@configclass +class ConveyorGaussianBernoulliDistributionCfg(RslRlMLPModelCfg.GaussianDistributionCfg): + """Bounded Gaussian arm exploration with a Bernoulli gripper.""" + + class_name: str = ( + "isaaclab_tasks.contrib.conveyor_franka.agents.rsl_rl_ppo_cfg:ConveyorGaussianBernoulliDistribution" + ) + std_range: tuple[float, float] = (0.15, 0.65) + + +@configclass +class ConveyorFrankaPPORunnerCfg(RslRlOnPolicyRunnerCfg): + """PPO configuration for four-cube commanded transfer.""" + + num_steps_per_env = 32 + # RSL-RL's usual randomized episode counters would desynchronize the + # per-subgoal timeout clock before the first policy step. + init_at_random_ep_len = False + max_iterations = 4000 + save_interval = 50 + experiment_name = "conveyor_franka_transfer" + clip_actions = 1.0 + obs_groups = {"actor": ["policy"], "critic": ["policy"]} + actor = RslRlMLPModelCfg( + hidden_dims=[512, 256, 128], + activation="elu", + obs_normalization=True, + distribution_cfg=ConveyorGaussianBernoulliDistributionCfg(init_std=0.45), + ) + critic = RslRlMLPModelCfg( + hidden_dims=[512, 256, 128], + activation="elu", + obs_normalization=True, + ) + algorithm = RslRlPpoAlgorithmCfg( + class_name="PPO", + value_loss_coef=1.0, + use_clipped_value_loss=True, + clip_param=0.2, + entropy_coef=0.001, + num_learning_epochs=5, + num_mini_batches=16, + learning_rate=1.0e-4, + schedule="fixed", + # At 60 Hz, the ten-second pickup-to-placement horizon needs the same + # long-horizon discount used by the reset-driven Franka stack task. + gamma=0.999, + lam=0.95, + desired_kl=0.01, + max_grad_norm=1.0, + ) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_franka_env.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_franka_env.py new file mode 100644 index 000000000000..680c77692b42 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_franka_env.py @@ -0,0 +1,93 @@ +# 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 + +"""Manager-based conveyor environment with Newton contact forces.""" + +from __future__ import annotations + +from collections.abc import Sequence + +import torch + +from isaaclab.envs import ManagerBasedRLEnv + +from isaaclab_contrib.conveyors.newton import SurfaceVelocity + +from .conveyor_franka_env_cfg import ConveyorFrankaEnvCfg +from .conveyor_goal_selector import ConveyorGoalSelector + + +class ConveyorFrankaEnv(ManagerBasedRLEnv): + """Manager-based environment with Newton conveyor surfaces.""" + + cfg: ConveyorFrankaEnvCfg + + def __init__(self, cfg: ConveyorFrankaEnvCfg, render_mode: str | None = None, **kwargs) -> None: + self._conveyor_driver: SurfaceVelocity | None = None + super().__init__(cfg, render_mode=render_mode, **kwargs) + self._setup_goal_selector() + + def load_managers(self) -> None: + """Bind the fixed Panda tool frame before manager terms inspect it.""" + hand_ids, _ = self.scene["robot"].find_bodies("panda_hand") + if len(hand_ids) != 1: + raise ValueError(f"Expected one Panda hand body, found {len(hand_ids)}.") + self._conveyor_hand_body_id = hand_ids[0] + self._conveyor_tool_offset = torch.tensor((0.0, 0.0, 0.1034), dtype=torch.float32, device=self.device).expand( + self.num_envs, -1 + ) + super().load_managers() + + def close(self) -> None: + """Release conveyor callbacks before the physics scene is destroyed.""" + conveyor_driver = self._conveyor_driver + if conveyor_driver is not None: + conveyor_driver.close() + self._conveyor_driver = None + super().close() + + def _init_sim(self) -> None: + """Register solved-contact hooks before the first reset captures the solver.""" + driver = SurfaceVelocity( + manager=self.sim.backend_manager, + num_envs=self.cfg.scene.num_envs, + surface_specs=self.cfg.scene.build_conveyor_belt_specs( + velocity=self.cfg.conveyor_force.speed, + friction_coefficient=self.cfg.conveyor_force.friction, + contact_threshold=self.cfg.conveyor_force.normal_threshold, + ), + startup_duration_s=self.cfg.conveyor_force.startup_duration_s, + env_path_format=self.cfg.scene.clone_cfg.clone_template, + body_pattern=self.cfg.conveyor_force.transported_body_pattern, + body_count_per_env=self.cfg.conveyor_force.transported_body_count_per_env, + ) + self._conveyor_driver = driver + try: + super()._init_sim() + except Exception: + driver.close() + self._conveyor_driver = None + raise + + def _setup_goal_selector(self) -> None: + """Attach one task panel to the first GL visualizer.""" + for visualizer in self.sim.visualizers: + if visualizer.cfg.backend.visualizer_type != "newton_gl": + continue + visible_env_ids = visualizer.get_visualized_env_ids() + env_id = visible_env_ids[0] if visible_env_ids else 0 + selector = ConveyorGoalSelector(self, env_id) + visualizer.register_ui_callback(selector.render, position="panel") + return + + def _reset_idx(self, env_ids: Sequence[int] | torch.Tensor | slice) -> None: + """Reset selected environments and discard stale conveyor forces.""" + if isinstance(env_ids, slice): + env_ids = self.scene._ALL_INDICES[env_ids] + super()._reset_idx(env_ids) + + conveyor_driver = self._conveyor_driver + if conveyor_driver is not None: + conveyor_driver.reset(env_ids) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_franka_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_franka_env_cfg.py new file mode 100644 index 000000000000..1a2763e29147 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_franka_env_cfg.py @@ -0,0 +1,552 @@ +# 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 + +"""Configuration for the force-driven conveyor and Franka demonstration scene.""" + +from __future__ import annotations + +import math +import re + +from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg, NewtonCollisionPipelineCfg, NewtonShapeCfg +from isaaclab_newton.sim.schemas import MujocoCollisionCfg, NewtonCollisionCfg, NewtonMaterialPropertiesCfg + +import isaaclab.sim as sim_utils +from isaaclab.assets import ArticulationCfg, AssetBaseCfg, RigidObjectCfg +from isaaclab.envs import ManagerBasedRLEnvCfg +from isaaclab.managers import CurriculumTermCfg as CurrTerm +from isaaclab.managers import EventTermCfg as EventTerm +from isaaclab.managers import ObservationGroupCfg as ObsGroup +from isaaclab.managers import ObservationTermCfg as ObsTerm +from isaaclab.managers import RewardTermCfg as RewTerm +from isaaclab.managers import SceneEntityCfg +from isaaclab.managers import TerminationTermCfg as DoneTerm +from isaaclab.scene import InteractiveSceneCfg +from isaaclab.sim import SimulationCfg +from isaaclab.sim.schemas import CollisionFragment, UsdPhysicsCollisionCfg, UsdPhysicsMeshCollisionCfg +from isaaclab.sim.spawners.materials import RigidBodyMaterialBaseCfg +from isaaclab.utils.configclass import configclass +from isaaclab.visualizers import VisualizerCfg + +from isaaclab_contrib.conveyors.surface_velocity import SurfaceVelocitySpec + +from . import mdp +from .conveyor_geometry import ( + BELT_COLOR, + BELT_INNER_STRAIGHT_Y, + BELT_OUTER_STRAIGHT_Y, + CUBE_COLORS, + CUBE_INNER_SLOT_X, + CUBE_OUTER_SLOT_X, + GUARD_COLOR, + CuboidSpec, + MeshSpec, + belt_collision_geometry_specs, + belt_collision_section_specs, + belt_mesh_spec, + guard_mesh_specs, +) +from .franka_robot_cfg import FRANKA_PANDA_CONVEYOR_CFG +from .mdp.terminations import invalid_action as invalid_policy_action + +_CONTACT_GAP = 0.01 +_CUBE_CONTACT_MARGIN = 0.003 +_MUJOCO_SOLIMP = (0.9, 0.95, 0.001, 0.5, 2.0) +_MUJOCO_SOLREF = (0.02, 1.0) +_SUBGOAL_TIMEOUT_S = 20.0 +_TRANSFER_SEQUENCE_LENGTH = 8 +_ARM_JOINT_NAMES = tuple(f"panda_joint{joint_id}" for joint_id in range(1, 8)) +_FINGER_JOINT_NAMES = ("panda_finger_joint1", "panda_finger_joint2") +_ROBOT_INITIAL_STATE = ArticulationCfg.InitialStateCfg( + joint_pos={ + "panda_joint1": 0.0, + "panda_joint2": -0.35, + "panda_joint3": 0.0, + "panda_joint4": -2.35, + "panda_joint5": 0.0, + "panda_joint6": 2.0, + "panda_joint7": 0.78, + "panda_finger_joint.*": 0.04, + } +) + + +def _validate_common_config(cfg: ConveyorFrankaEnvCfg) -> None: + """Validate backend-independent timing and policy tensor contracts.""" + if not cfg.sim.use_newton_actuators: + raise ValueError("The conveyor Franka requires the shared Newton-actuator execution path.") + if not math.isfinite(cfg.sim.dt) or cfg.sim.dt <= 0.0 or cfg.decimation <= 0: + raise ValueError("Simulation dt and environment decimation must be positive.") + arm_action = cfg.actions.arm_action + if arm_action.joint_names != list(_ARM_JOINT_NAMES) or not arm_action.preserve_order: + raise ValueError("Arm actions must preserve the explicit panda_joint1-to-panda_joint7 ordering.") + + +def _collision_properties(contact_margin: float = 0.0, mujoco_priority: int = 0) -> list[CollisionFragment]: + """Build explicit Newton and MuJoCo contact properties for one collider.""" + return [ + UsdPhysicsCollisionCfg(collision_enabled=True), + NewtonCollisionCfg(contact_margin=contact_margin, contact_gap=_CONTACT_GAP), + MujocoCollisionCfg( + condim=3, + priority=mujoco_priority, + solimp=_MUJOCO_SOLIMP, + solmix=1.0, + solref=_MUJOCO_SOLREF, + ), + ] + + +@configclass +class ActionsCfg: + """Relative arm and binary gripper actions.""" + + arm_action = mdp.ConveyorRelativeJointPositionActionCfg( + asset_name="robot", + joint_names=list(_ARM_JOINT_NAMES), + preserve_order=True, + scale=0.12, + max_delta=0.12, + ) + gripper_action = mdp.ResetBufferedGripperActionCfg( + asset_name="robot", + joint_names=list(_FINGER_JOINT_NAMES), + open_command_expr={"panda_finger_joint.*": 0.04}, + close_command_expr={"panda_finger_joint.*": 0.0}, + force_close_steps=5, + ) + + +@configclass +class CommandsCfg: + """Success-driven cube-transfer command.""" + + transfer = mdp.ConveyorTransferCommandCfg( + reset_event_name="reset_from_state_table", + minimum_subgoal_steps=2, + hold_steps=3, + lateral_tolerance=0.055, + maximum_cube_speed=0.65, + minimum_finger_position=0.027, + minimum_tool_clearance=0.055, + ) + + +@configclass +class ObservationsCfg: + """Policy observations with stable cube identity and transfer commands.""" + + @configclass + class PolicyCfg(ObsGroup): + """Fully observed transfer policy input.""" + + joint_pos = ObsTerm( + func=mdp.joint_pos_rel, + params={"asset_cfg": SceneEntityCfg("robot", joint_names=list(_ARM_JOINT_NAMES), preserve_order=True)}, + ) + joint_vel = ObsTerm( + func=mdp.joint_vel_rel, + params={"asset_cfg": SceneEntityCfg("robot", joint_names=list(_ARM_JOINT_NAMES), preserve_order=True)}, + ) + gripper_pos = ObsTerm( + func=mdp.gripper_joint_positions, + params={"robot_cfg": SceneEntityCfg("robot", joint_names=list(_FINGER_JOINT_NAMES), preserve_order=True)}, + ) + objects = ObsTerm(func=mdp.transfer_object_observation) + active_transfer = ObsTerm(func=mdp.active_transfer_features) + target_cube = ObsTerm(func=mdp.target_cube_one_hot) + cube_conveyors = ObsTerm(func=mdp.cube_conveyor_state) + target_side = ObsTerm(func=mdp.target_side_one_hot) + eef_velocity = ObsTerm(func=mdp.end_effector_velocity) + eef_axes = ObsTerm(func=mdp.end_effector_axes) + last_action = ObsTerm(func=mdp.last_action) + + def __post_init__(self) -> None: + self.enable_corruption = False + self.concatenate_terms = True + + policy: PolicyCfg = PolicyCfg() + + +@configclass +class EventCfg: + """Restore validated physical reset states.""" + + reset_all = EventTerm(func=mdp.reset_scene_to_default, mode="reset") + reset_from_state_table = EventTerm( + func=mdp.ConveyorResetStateTable, + mode="reset", + params={ + "fixed_recipe": None, + "fixed_variant_id": None, + "fixed_target_cube_id": None, + "fixed_source_side_id": None, + "belt_start_x_range": (0.30, 0.82), + "cube_position_noise": 0.015, + "arm_joint_noise": 0.015, + }, + ) + + +@configclass +class RewardsCfg: + """Sparse transfer completion plus safety regularization rewards.""" + + success = RewTerm( + func=mdp.transfer_success_reward, + params={"command_name": "transfer"}, + weight=600.0, + ) + failure = RewTerm( + func=mdp.terminal_failure, + params={"command_name": "transfer"}, + weight=-60.0, + ) + arm_action_l2 = RewTerm( + func=mdp.action_term_l2, + params={"action_name": "arm_action"}, + weight=-1.0e-3, + ) + action_rate_l2 = RewTerm( + func=mdp.finite_action_rate_l2, + params={"action_names": ("arm_action", "gripper_action")}, + weight=-1.0e-3, + ) + joint_velocity_l2 = RewTerm( + func=mdp.finite_joint_velocity_l2, + params={"asset_cfg": SceneEntityCfg("robot", joint_names=list(_ARM_JOINT_NAMES), preserve_order=True)}, + weight=-1.0e-4, + ) + + +@configclass +class TerminationsCfg: + """Safety failures and bounded training sequences.""" + + cube_out_of_workspace = DoneTerm(func=mdp.cube_out_of_workspace) + invalid_action = DoneTerm( + func=invalid_policy_action, + params={"action_names": ("arm_action", "gripper_action")}, + ) + nonfinite_scene_state = DoneTerm(func=mdp.nonfinite_scene_state) + subgoal_time_out = DoneTerm( + func=mdp.subgoal_time_out, + params={"timeout_s": _SUBGOAL_TIMEOUT_S, "command_name": "transfer"}, + time_out=True, + ) + transfer_sequence_time_out = DoneTerm( + func=mdp.transfer_sequence_time_out, + params={"maximum_transfers": _TRANSFER_SEQUENCE_LENGTH, "command_name": "transfer"}, + time_out=True, + ) + + +@configclass +class CurriculumCfg: + """Practice reset states near the measured transfer-success frontier.""" + + reset_sampling = CurrTerm( + func=mdp.ConveyorResetCurriculum, + params={ + "command_name": "transfer", + "success_monitor": mdp.SuccessMonitorCfg(monitored_history_len=50), + "sampling": mdp.SamplerCfg( + strategies=[ + mdp.BetaSamplingStrategyCfg(target=0.5, kappa=2.0), + mdp.UniformSamplingStrategyCfg(weight=0.05), + ], + eps=1.0e-4, + ), + "fixed_source_side_id": None, + }, + ) + + +@configclass +class ConveyorForceCfg: + """Configuration for force-based conveyor traction.""" + + speed: float = 0.35 + """Tangential conveyor surface speed [m/s].""" + + friction: float = 0.5 + """Coulomb friction coefficient used to limit traction.""" + + normal_threshold: float = 0.997 + """Minimum upward contact-normal alignment in the range [0, 1].""" + + startup_duration_s: float = 1.0 + """Duration over which conveyor traction ramps to full speed [s].""" + + transported_body_pattern: str = r"(?:^|/)Cube_?[0-3](?:/|$)" + """Regular expression selecting rigid bodies that receive conveyor forces.""" + + transported_body_count_per_env: int = 4 + """Expected number of transported rigid bodies in each environment.""" + + def validate_config(self) -> None: + """Validate final conveyor-force values after overrides are applied.""" + if not math.isfinite(self.speed) or self.speed < 0.0: + raise ValueError(f"Conveyor speed must be non-negative, got {self.speed}.") + if not math.isfinite(self.friction) or self.friction < 0.0: + raise ValueError(f"Conveyor friction must be non-negative, got {self.friction}.") + if not math.isfinite(self.normal_threshold) or not 0.0 <= self.normal_threshold <= 1.0: + raise ValueError(f"Conveyor normal threshold must be in [0, 1], got {self.normal_threshold}.") + if not math.isfinite(self.startup_duration_s) or self.startup_duration_s <= 0.0: + raise ValueError(f"Conveyor startup duration must be positive, got {self.startup_duration_s}.") + if self.transported_body_count_per_env <= 0: + raise ValueError( + "Conveyor transported body count per environment must be positive, got " + f"{self.transported_body_count_per_env}." + ) + try: + re.compile(self.transported_body_pattern) + except re.error as exc: + raise ValueError(f"Invalid conveyor transported-body pattern: {self.transported_body_pattern!r}.") from exc + + +def _static_cuboid( + prim_path: str, + size: tuple[float, float, float], + pos: tuple[float, float, float], + color: tuple[float, float, float], +) -> AssetBaseCfg: + """Build a colliding table component.""" + return AssetBaseCfg( + prim_path=prim_path, + init_state=AssetBaseCfg.InitialStateCfg(pos=pos), + spawn=sim_utils.CuboidCfg( + size=size, + collision_props=_collision_properties(), + physics_material=RigidBodyMaterialBaseCfg(static_friction=0.7, dynamic_friction=0.7, restitution=0.0), + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=color, roughness=0.75, metallic=0.0), + ), + ) + + +def _conveyor_asset( + prim_path: str, + spec: MeshSpec | CuboidSpec, + *, + visual_material: sim_utils.PreviewSurfaceCfg | None = None, + colliding: bool = True, +) -> AssetBaseCfg: + """Build conveyor geometry with freely sliding contacts and optional rendering.""" + if isinstance(spec, CuboidSpec): + spawn = sim_utils.CuboidCfg(size=spec.size) + position = spec.position + else: + spawn = sim_utils.MeshFileCfg( + mesh=sim_utils.MeshFileCfg.TriangleMeshCfg(vertices=spec.vertices, faces=spec.faces), + mesh_collision_props=UsdPhysicsMeshCollisionCfg(mesh_approximation_name="none"), + ) + position = (0.0, 0.0, 0.0) + spawn.visible = visual_material is not None + spawn.visual_material = visual_material + if colliding: + spawn.collision_props = _collision_properties(mujoco_priority=1) + # Override cube/finger grasp friction at the belts and guide rails. MuJoCo + # requires positive friction; the force model supplies belt traction. + spawn.physics_material = NewtonMaterialPropertiesCfg( + static_friction=1.1e-5, dynamic_friction=1.1e-5, restitution=0.0 + ) + return AssetBaseCfg( + prim_path=prim_path, + init_state=AssetBaseCfg.InitialStateCfg(pos=position), + spawn=spawn, + ) + + +def _cube( + name: str, + color: tuple[float, float, float], + pos: tuple[float, float, float], +) -> RigidObjectCfg: + """Build one numbered dynamic transfer cube.""" + spawn = sim_utils.CuboidCfg( + size=(0.04, 0.04, 0.04), + rigid_props=sim_utils.UsdPhysicsRigidBodyCfg(), + mass_props=sim_utils.MassCfg(mass=0.05), + collision_props=_collision_properties(contact_margin=_CUBE_CONTACT_MARGIN), + # Belt priority overrides friction only for belt/cube contacts. + physics_material=NewtonMaterialPropertiesCfg(static_friction=0.8, dynamic_friction=0.6, restitution=0.0), + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=color, roughness=0.75), + ) + return RigidObjectCfg( + prim_path=f"{{ENV_REGEX_NS}}/{name}", + init_state=RigidObjectCfg.InitialStateCfg(pos=pos), + spawn=spawn, + ) + + +@configclass +class ConveyorFrankaSceneCfg(InteractiveSceneCfg): + """Scene with two counter-rotating racetrack conveyors around a table-mounted Franka.""" + + # Use the MuJoCo Menagerie-derived model with explicit manipulation gains. + robot = FRANKA_PANDA_CONVEYOR_CFG.replace( + prim_path="{ENV_REGEX_NS}/Robot", + init_state=_ROBOT_INITIAL_STATE, + ) + + tabletop = _static_cuboid( + prim_path="{ENV_REGEX_NS}/Tabletop", + size=(2.0, 1.9, 0.08), + pos=(0.50, 0.0, -0.04), + color=(0.32, 0.34, 0.37), + ) + table_pedestal = _static_cuboid( + prim_path="{ENV_REGEX_NS}/TablePedestal", + size=(0.75, 0.55, 0.76), + pos=(0.25, 0.0, -0.46), + color=(0.18, 0.20, 0.23), + ) + + cube_0 = _cube("Cube0", CUBE_COLORS[0], (CUBE_INNER_SLOT_X, BELT_INNER_STRAIGHT_Y, 0.06)) + cube_1 = _cube("Cube1", CUBE_COLORS[1], (CUBE_OUTER_SLOT_X, BELT_OUTER_STRAIGHT_Y, 0.06)) + cube_2 = _cube("Cube2", CUBE_COLORS[2], (CUBE_INNER_SLOT_X, -BELT_INNER_STRAIGHT_Y, 0.06)) + cube_3 = _cube("Cube3", CUBE_COLORS[3], (CUBE_OUTER_SLOT_X, -BELT_OUTER_STRAIGHT_Y, 0.06)) + + ground = AssetBaseCfg( + prim_path="/World/GroundPlane", + init_state=AssetBaseCfg.InitialStateCfg(pos=(0.0, 0.0, -0.85)), + spawn=sim_utils.GroundPlaneCfg(), + ) + dome_light = AssetBaseCfg( + prim_path="/World/DomeLight", + spawn=sim_utils.DomeLightCfg(color=(0.8, 0.8, 0.8), intensity=2500.0), + ) + + def __post_init__(self) -> None: + """Generate visual belts, velocity-field collision sections, and guardrails.""" + for side in ("Left", "Right"): + belt_spec = belt_mesh_spec(side) + setattr( + self, + f"conveyor_{side.lower()}_belt_visual", + _conveyor_asset( + prim_path=f"{{ENV_REGEX_NS}}/{belt_spec.name}", + spec=belt_spec, + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=BELT_COLOR, roughness=0.9, metallic=0.0), + colliding=False, + ), + ) + + section_keys = ("top_straight", "bottom_straight", "right_turn", "left_turn") + for section_key, spec in zip(section_keys, belt_collision_geometry_specs(side), strict=True): + setattr( + self, + f"conveyor_{side.lower()}_{section_key}_collision", + _conveyor_asset( + prim_path=f"{{ENV_REGEX_NS}}/{spec.name}", + spec=spec, + ), + ) + + for spec in guard_mesh_specs(side): + boundary = "inner" if spec.name.endswith("Inner") else "outer" + setattr( + self, + f"guard_{side.lower()}_{boundary}", + _conveyor_asset( + prim_path=f"{{ENV_REGEX_NS}}/{spec.name}Collision", + spec=spec, + visual_material=sim_utils.PreviewSurfaceCfg( + diffuse_color=GUARD_COLOR, roughness=0.3, metallic=0.8 + ), + ), + ) + + def build_conveyor_belt_specs(self, **kwargs: float) -> tuple[SurfaceVelocitySpec, ...]: + """Describe the velocity field on each authored collision section.""" + return tuple( + section.belt for side in ("Left", "Right") for section in belt_collision_section_specs(side, **kwargs) + ) + + +@configclass +class ConveyorFrankaEnvCfg(ManagerBasedRLEnvCfg): + """Manager-based RL task for commanded conveyor-to-conveyor cube transfer.""" + + class_type: type | str = "{DIR}.conveyor_franka_env:ConveyorFrankaEnv" + scene: ConveyorFrankaSceneCfg = ConveyorFrankaSceneCfg(num_envs=256, env_spacing=3.0) + conveyor_force: ConveyorForceCfg = ConveyorForceCfg() + actions: ActionsCfg = ActionsCfg() + commands: CommandsCfg = CommandsCfg() + observations: ObservationsCfg = ObservationsCfg() + events: EventCfg = EventCfg() + rewards: RewardsCfg = RewardsCfg() + terminations: TerminationsCfg = TerminationsCfg() + curriculum: CurriculumCfg = CurriculumCfg() + decimation: int = 2 + episode_length_s: float = _SUBGOAL_TIMEOUT_S * _TRANSFER_SEQUENCE_LENGTH + + sim: SimulationCfg = SimulationCfg( + dt=1.0 / 120.0, + render_interval=2, + physics=NewtonCfg( + solver_cfg=MJWarpSolverCfg( + solver="newton", + integrator="implicitfast", + # Per-environment capacities include headroom for all four cubes, + # the gripper, table, belts, and guard contacts. + njmax=300, + nconmax=200, + impratio=1.0, + cone="elliptic", + update_data_interval=1, + iterations=100, + ls_iterations=50, + use_mujoco_contacts=False, + ccd_iterations=35, + ), + collision_cfg=NewtonCollisionPipelineCfg(), + # Manager decimation supplies the reference's two 120 Hz solves + # per 60 Hz policy step, with contacts refreshed before each solve. + collision_decimation=0, + default_shape_cfg=NewtonShapeCfg(margin=0.0, gap=_CONTACT_GAP, ke=2.5e3, kd=100.0), + num_substeps=1, + use_cuda_graph=True, + # Import render-only geometry only when a visualizer or camera needs it. + load_visual_shapes=None, + ), + use_newton_actuators=True, + ) + + def __post_init__(self) -> None: + # Any visualizer selected at runtime receives these shared camera hints. + # Newton camera pose: position (2.13, 0.0, 1.0), pitch -23.9 degrees, + # yaw 180 degrees. The look-at point is one unit along that view ray. + self.sim.visualizer_cfgs = VisualizerCfg( + eye=(2.13, 0.0, 1.0), + lookat=(1.2157460448, 0.0, 0.5948584132), + max_visible_envs=1, + randomly_sample_visible_envs=False, + ) + + def validate_config(self) -> None: + """Validate the final task configuration after command-line overrides are applied.""" + _validate_common_config(self) + physics = self.sim.physics + if not isinstance(physics, NewtonCfg) or not isinstance(physics.solver_cfg, MJWarpSolverCfg): + raise ValueError("The conveyor force driver requires the Newton MJWarp backend.") + if physics.solver_cfg.use_mujoco_contacts: + raise ValueError("The conveyor force driver requires the Newton collision-pipeline contact path.") + if physics.collision_cfg is None: + raise ValueError("The conveyor force driver requires an explicit Newton collision pipeline.") + + def play_mode(self) -> None: + """Run continuing transfers from evenly distributed moving-belt starts.""" + super().play_mode() + self.scene.num_envs = min(self.scene.num_envs, 8) + self.events.reset_from_state_table.params["fixed_recipe"] = int(mdp.ConveyorResetRecipe.BELT) + self.events.reset_from_state_table.params["fixed_variant_id"] = mdp.BELT_DEPLOYMENT_VARIANT + self.events.reset_from_state_table.params["cube_position_noise"] = 0.0 + # Successful placements already transition to a new commanded cube. + # Playback removes training-only refreshes and runs until physics leaves + # the recoverable workspace. + self.terminations.subgoal_time_out = None + self.terminations.transfer_sequence_time_out = None + self.curriculum = None diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_geometry.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_geometry.py new file mode 100644 index 000000000000..1217e449364d --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_geometry.py @@ -0,0 +1,320 @@ +# 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 + +"""Shared racetrack geometry and velocity-field descriptions.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass + +from isaaclab_contrib.conveyors.surface_velocity import SurfaceVelocitySpec + +BELT_COLOR = (0.09, 0.09, 0.09) +"""Dark-rubber color used by Newton's conveyor example.""" + +GUARD_COLOR = (0.66, 0.69, 0.74) +"""Brushed-metal color used by Newton's conveyor example.""" + +PARCEL_COLOR = (0.72, 0.55, 0.35) +"""Cardboard color used by Newton's conveyor example.""" + +CUBE_COLORS = ( + (0.15, 0.35, 0.90), + (0.90, 0.20, 0.15), + (0.15, 0.75, 0.25), + PARCEL_COLOR, +) +"""Stable sRGB colors for the four numbered transfer cubes.""" + +BELT_CENTER_X = 0.58 +BELT_CENTER_Y = 0.51 +BELT_TURN_RADIUS = 0.24 +BELT_HALF_STRAIGHT = 0.44 +BELT_WIDTH = 0.15 +BELT_THICKNESS = 0.04 +BELT_TOP_Z = 0.04 +TURN_SEGMENT_COUNT = 96 + +# Four deployment slots place one cube on each straight run of the two +# racetracks. The x coordinates are mirrored about the racetrack center, so +# the inner/outer pair on a belt is separated by exactly half a lap. +BELT_INNER_STRAIGHT_Y = BELT_CENTER_Y - BELT_TURN_RADIUS +BELT_OUTER_STRAIGHT_Y = BELT_CENTER_Y + BELT_TURN_RADIUS +CUBE_INNER_SLOT_X = BELT_CENTER_X - BELT_HALF_STRAIGHT / 3.0 +CUBE_OUTER_SLOT_X = BELT_CENTER_X + BELT_HALF_STRAIGHT / 3.0 + +# Collision surfaces extend underneath the rails and overlap at section seams. +# This keeps the dynamic parcels on a continuous +Z-facing surface without +# exposing the belt prism's vertical side faces to the contact solver. +BELT_COLLISION_OVERHANG = 0.02 +BELT_COLLISION_SEAM_OVERLAP = 0.004 + +GUARD_THICKNESS = 0.018 +# Keep the rails below the parcel tops so both lanes remain easy to read from +# the default oblique camera. +GUARD_HEIGHT = 0.02 +GUARD_BASE_OVERLAP = 0.005 + + +@dataclass(frozen=True) +class MeshSpec: + """Triangle mesh and semantic name for one static racetrack component.""" + + name: str + vertices: tuple[tuple[float, float, float], ...] + faces: tuple[tuple[int, int, int], ...] + + +@dataclass(frozen=True) +class CuboidSpec: + """Native cuboid and semantic name for one static racetrack component.""" + + name: str + size: tuple[float, float, float] + position: tuple[float, float, float] + + +@dataclass(frozen=True) +class ConveyorSectionSpec: + """Task geometry paired with its backend-neutral conveyor description.""" + + geometry: MeshSpec | CuboidSpec + belt: SurfaceVelocitySpec + + +def belt_direction(side: str) -> float: + """Return ``1`` for clockwise motion and ``-1`` for counter-clockwise motion.""" + if side == "Left": + return 1.0 + if side == "Right": + return -1.0 + raise ValueError(f"Unknown conveyor side: {side!r}.") + + +def _racetrack_centerline(center_y: float) -> tuple[tuple[float, float, float, float], ...]: + """Sample a clockwise racetrack centerline with an outward normal at every point.""" + left_x = BELT_CENTER_X - BELT_HALF_STRAIGHT + right_x = BELT_CENTER_X + BELT_HALF_STRAIGHT + radius = BELT_TURN_RADIUS + points: list[tuple[float, float, float, float]] = [ + (left_x, center_y + radius, 0.0, 1.0), + (right_x, center_y + radius, 0.0, 1.0), + ] + + for index in range(1, TURN_SEGMENT_COUNT + 1): + angle = 0.5 * math.pi - index * math.pi / TURN_SEGMENT_COUNT + normal_x = math.cos(angle) + normal_y = math.sin(angle) + points.append((right_x + radius * normal_x, center_y + radius * normal_y, normal_x, normal_y)) + + points.append((left_x, center_y - radius, 0.0, -1.0)) + for index in range(1, TURN_SEGMENT_COUNT): + angle = -0.5 * math.pi - index * math.pi / TURN_SEGMENT_COUNT + normal_x = math.cos(angle) + normal_y = math.sin(angle) + points.append((left_x + radius * normal_x, center_y + radius * normal_y, normal_x, normal_y)) + return tuple(points) + + +def _prism_faces(count: int, *, closed: bool) -> tuple[tuple[int, int, int], ...]: + """Triangulate top, bottom, and side walls for inner/outer vertex rings. + + Vertices are grouped as inner top, outer top, inner bottom, outer bottom. + Counter-clockwise centerlines produce outward faces; open arcs also receive end caps. + """ + quads = [] + for i in range(count if closed else count - 1): + j = (i + 1) % count + quads.extend( + ( + (i, count + i, count + j, j), + (2 * count + i, 2 * count + j, 3 * count + j, 3 * count + i), + (3 * count + i, 3 * count + j, count + j, count + i), + (2 * count + i, i, j, 2 * count + j), + ) + ) + if not closed: + end = count - 1 + quads.extend(((0, 2 * count, 3 * count, count), (end, count + end, 3 * count + end, 2 * count + end))) + return tuple(triangle for a, b, c, d in quads for triangle in ((a, b, c), (a, c, d))) + + +def _racetrack_prism_mesh( + name: str, + center_y: float, + lateral_offset: float, + width: float, + z_min: float, + z_max: float, +) -> MeshSpec: + """Build one closed prism following a racetrack centerline.""" + centerline = _racetrack_centerline(center_y) + half_width = 0.5 * width + outer_offset = lateral_offset + half_width + inner_offset = lateral_offset - half_width + + points = tuple( + (x + nx * offset, y + ny * offset, z) + for z in (z_max, z_min) + for offset in (inner_offset, outer_offset) + for x, y, nx, ny in centerline + ) + + # Racetrack centerlines are clockwise; reverse the shared counter-clockwise winding. + faces = tuple((a, c, b) for a, b, c in _prism_faces(len(centerline), closed=True)) + return MeshSpec(name=name, vertices=points, faces=faces) + + +def belt_mesh_spec(side: str) -> MeshSpec: + """Build the seamless, watertight visual mesh for one conveyor belt.""" + center_y = BELT_CENTER_Y if side == "Left" else -BELT_CENTER_Y + return _racetrack_prism_mesh( + name=f"Conveyor{side}BeltVisual", + center_y=center_y, + lateral_offset=0.0, + width=BELT_WIDTH, + z_min=BELT_TOP_Z - BELT_THICKNESS, + z_max=BELT_TOP_Z, + ) + + +def _straight_collision_cuboid(name: str, center_y: float) -> CuboidSpec: + """Build one solid native cuboid for a straight conveyor section.""" + half_width = 0.5 * BELT_WIDTH + BELT_COLLISION_OVERHANG + x_min = BELT_CENTER_X - BELT_HALF_STRAIGHT - BELT_COLLISION_SEAM_OVERLAP + x_max = BELT_CENTER_X + BELT_HALF_STRAIGHT + BELT_COLLISION_SEAM_OVERLAP + return CuboidSpec( + name=name, + size=(x_max - x_min, 2.0 * half_width, BELT_THICKNESS), + position=(0.5 * (x_min + x_max), center_y, BELT_TOP_Z - 0.5 * BELT_THICKNESS), + ) + + +def _turn_collision_mesh( + name: str, pivot_x: float, center_y: float, start_angle: float, angle_span: float = math.pi +) -> MeshSpec: + """Build a closed annular turn with a +Z top surface and an angular span [rad].""" + half_width = 0.5 * BELT_WIDTH + BELT_COLLISION_OVERHANG + inner_radius = BELT_TURN_RADIUS - half_width + outer_radius = BELT_TURN_RADIUS + half_width + angle_overlap = BELT_COLLISION_SEAM_OVERLAP / BELT_TURN_RADIUS + angle_start = start_angle - angle_overlap + segment_count = round(TURN_SEGMENT_COUNT * angle_span / math.pi) + angle_step = (angle_span + 2.0 * angle_overlap) / segment_count + angles = tuple(angle_start + index * angle_step for index in range(segment_count + 1)) + + inner_top = tuple( + (pivot_x + inner_radius * math.cos(angle), center_y + inner_radius * math.sin(angle), BELT_TOP_Z) + for angle in angles + ) + outer_top = tuple( + (pivot_x + outer_radius * math.cos(angle), center_y + outer_radius * math.sin(angle), BELT_TOP_Z) + for angle in angles + ) + bottom_z = BELT_TOP_Z - BELT_THICKNESS + inner_bottom = tuple((x, y, bottom_z) for x, y, _ in inner_top) + outer_bottom = tuple((x, y, bottom_z) for x, y, _ in outer_top) + + return MeshSpec( + name=name, + vertices=inner_top + outer_top + inner_bottom + outer_bottom, + faces=_prism_faces(len(inner_top), closed=False), + ) + + +def belt_collision_geometry_specs(side: str) -> tuple[CuboidSpec | MeshSpec, ...]: + """Build native straight and closed-mesh turn collision geometry for one racetrack.""" + center_y = BELT_CENTER_Y if side == "Left" else -BELT_CENTER_Y + left_x = BELT_CENTER_X - BELT_HALF_STRAIGHT + right_x = BELT_CENTER_X + BELT_HALF_STRAIGHT + return ( + _straight_collision_cuboid(f"Conveyor{side}TopStraightCollision", center_y + BELT_TURN_RADIUS), + _straight_collision_cuboid(f"Conveyor{side}BottomStraightCollision", center_y - BELT_TURN_RADIUS), + _turn_collision_mesh(f"Conveyor{side}RightTurnCollision", right_x, center_y, -0.5 * math.pi), + _turn_collision_mesh(f"Conveyor{side}LeftTurnCollision", left_x, center_y, 0.5 * math.pi), + ) + + +def belt_collision_section_specs( + side: str, + *, + velocity: float = 0.0, + friction_coefficient: float = 0.7, + contact_threshold: float = 0.997, +) -> tuple[ConveyorSectionSpec, ConveyorSectionSpec, ConveyorSectionSpec, ConveyorSectionSpec]: + """Build collision geometry and authored conveyor intent for one racetrack.""" + center_y = BELT_CENTER_Y if side == "Left" else -BELT_CENTER_Y + left_x = BELT_CENTER_X - BELT_HALF_STRAIGHT + right_x = BELT_CENTER_X + BELT_HALF_STRAIGHT + direction = belt_direction(side) + top_straight, bottom_straight, right_turn, left_turn = belt_collision_geometry_specs(side) + + def belt( + geometry: MeshSpec | CuboidSpec, + travel_direction: tuple[float, float, float], + *, + curved: bool = False, + pivot_point: tuple[float, float, float] = (0.0, 0.0, 0.0), + radius: float | None = None, + ) -> ConveyorSectionSpec: + return ConveyorSectionSpec( + geometry=geometry, + belt=SurfaceVelocitySpec( + prim_path=f"{{ENV_REGEX_NS}}/{geometry.name}", + velocity=velocity, + direction=travel_direction, + curved=curved, + pivot_point=pivot_point, + radius=radius, + contact_threshold=contact_threshold, + friction_coefficient=friction_coefficient, + ), + ) + + return ( + belt(top_straight, (direction, 0.0, 0.0)), + belt(bottom_straight, (-direction, 0.0, 0.0)), + belt( + right_turn, + (0.0, 0.0, -direction), + curved=True, + pivot_point=(right_x, center_y, 0.0), + radius=BELT_TURN_RADIUS, + ), + belt( + left_turn, + (0.0, 0.0, -direction), + curved=True, + pivot_point=(left_x, center_y, 0.0), + radius=BELT_TURN_RADIUS, + ), + ) + + +def guard_mesh_specs(side: str) -> tuple[MeshSpec, MeshSpec]: + """Build seamless inner and outer guardrail meshes for one racetrack.""" + center_y = BELT_CENTER_Y if side == "Left" else -BELT_CENTER_Y + rail_offset = 0.5 * (BELT_WIDTH + GUARD_THICKNESS) + z_min = BELT_TOP_Z - GUARD_BASE_OVERLAP + z_max = BELT_TOP_Z + GUARD_HEIGHT + return ( + _racetrack_prism_mesh( + name=f"Guard{side}Inner", + center_y=center_y, + lateral_offset=-rail_offset, + width=GUARD_THICKNESS, + z_min=z_min, + z_max=z_max, + ), + _racetrack_prism_mesh( + name=f"Guard{side}Outer", + center_y=center_y, + lateral_offset=rail_offset, + width=GUARD_THICKNESS, + z_min=z_min, + z_max=z_max, + ), + ) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_goal_selector.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_goal_selector.py new file mode 100644 index 000000000000..093dcdbd0a16 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_goal_selector.py @@ -0,0 +1,102 @@ +# 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 + +"""Minimal Newton-viewer selector for conveyor transfer goals.""" + +from __future__ import annotations + +import time +from typing import TYPE_CHECKING, Any + +from .conveyor_geometry import CUBE_COLORS +from .mdp.reset_events import LEFT_SIDE + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedRLEnv + + +def _mix_color( + color: tuple[float, float, float], + target: tuple[float, float, float], + amount: float, +) -> tuple[float, float, float]: + """Linearly mix two RGB colors.""" + return tuple(value + amount * (target_value - value) for value, target_value in zip(color, target, strict=True)) + + +class ConveyorGoalSelector: + """Render four color-matched cube buttons for one displayed environment.""" + + def __init__(self, env: ManagerBasedRLEnv, env_id: int) -> None: + """Create a selector bound to one vectorized environment index.""" + if not 0 <= env_id < env.num_envs: + raise IndexError(f"Conveyor selector environment {env_id} is out of range.") + self._env = env + self._env_id = env_id + self._command = env.command_manager.get_term("transfer") + self._target_cube_id = 0 + self._source_side_id = LEFT_SIDE + self._last_refresh_time = float("-inf") + + def render(self, imgui: Any) -> None: + """Draw the selector and publish a new transfer command when clicked.""" + imgui.set_next_item_open(True, imgui.Cond_.appearing) + if not imgui.collapsing_header(f"Transfer Goal · Env {self._env_id}"): + return + imgui.separator() + + self._refresh_command() + + style = imgui.get_style() + available_width = float(imgui.get_content_region_avail().x) + spacing = float(style.item_spacing.x) + button_width = max(36.0, (available_width - spacing * (len(CUBE_COLORS) - 1)) / len(CUBE_COLORS)) + button_height = max(34.0, min(46.0, button_width * 0.72)) + + selected_cube_id: int | None = None + for cube_id, color in enumerate(CUBE_COLORS): + selected = cube_id == self._target_cube_id + text_color = (0.05, 0.05, 0.05) if sum(color) > 1.45 else (1.0, 1.0, 1.0) + border_color = (1.0, 0.88, 0.20) if selected else (0.12, 0.12, 0.12) + label = f"{cube_id + 1}##conveyor_goal_{cube_id}" + + imgui.push_style_color(imgui.Col_.button, imgui.ImVec4(*color, 1.0)) + imgui.push_style_color( + imgui.Col_.button_hovered, + imgui.ImVec4(*_mix_color(color, (1.0, 1.0, 1.0), 0.18), 1.0), + ) + imgui.push_style_color( + imgui.Col_.button_active, + imgui.ImVec4(*_mix_color(color, (0.0, 0.0, 0.0), 0.12), 1.0), + ) + imgui.push_style_color(imgui.Col_.border, imgui.ImVec4(*border_color, 1.0)) + imgui.push_style_color(imgui.Col_.text, imgui.ImVec4(*text_color, 1.0)) + imgui.push_style_var(imgui.StyleVar_.frame_border_size, 3.0 if selected else 1.0) + imgui.push_style_var(imgui.StyleVar_.frame_rounding, 5.0) + + if imgui.button(label, imgui.ImVec2(button_width, button_height)): + selected_cube_id = cube_id + + imgui.pop_style_var(2) + imgui.pop_style_color(5) + if cube_id + 1 < len(CUBE_COLORS): + imgui.same_line() + + if selected_cube_id is not None and selected_cube_id != self._target_cube_id: + self._command.set_goal(selected_cube_id, env_ids=(self._env_id,)) + self._refresh_command(force=True) + + source_name = "Left" if self._source_side_id == LEFT_SIDE else "Right" + target_name = "Right" if self._source_side_id == LEFT_SIDE else "Left" + imgui.text(f"Cube {self._target_cube_id + 1}: {source_name} -> {target_name}") + imgui.text_disabled("Click a color to change the next transfer.") + + def _refresh_command(self, force: bool = False) -> None: + """Refresh the small host-side UI cache at most ten times per second.""" + current_time = time.monotonic() + if force or current_time - self._last_refresh_time >= 0.1: + self._target_cube_id = int(self._command.target_cube_ids[self._env_id].item()) + self._source_side_id = int(self._command.source_side_ids[self._env_id].item()) + self._last_refresh_time = current_time diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/evaluate.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/evaluate.py new file mode 100644 index 000000000000..6ae9ba7a36e0 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/evaluate.py @@ -0,0 +1,106 @@ +# 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 + +"""Evaluate a conveyor checkpoint from reproducible, full moving-belt starts.""" + +from __future__ import annotations + +import argparse +import json +from contextlib import closing +from pathlib import Path + +import gymnasium as gym +import torch +from rsl_rl.runners import OnPolicyRunner + +from isaaclab.app import add_launcher_args, launch_simulation +from isaaclab.utils import to_dict +from isaaclab.utils.assets import retrieve_file_path + +from isaaclab_rl.entrypoints.common import resolve_published_checkpoint +from isaaclab_rl.rsl_rl import RslRlVecEnvWrapper + +import isaaclab_tasks # noqa: F401 +from isaaclab_tasks.utils import load_cfg_from_registry, parse_env_cfg + +_TASK = "IsaacContrib-Conveyor-Racetrack-Transfer-v0" + + +def main() -> None: + """Run a bounded evaluation and report transfers and failures instead of training reward.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--checkpoint", default="pretrained") + parser.add_argument("--num_envs", type=int, default=8) + parser.add_argument("--steps", type=int, default=3600) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--output", type=Path) + add_launcher_args(parser) + parser.set_defaults(device="cpu", headless=True) + args = parser.parse_args() + if args.num_envs < 1 or args.steps < 1: + parser.error("--num_envs and --steps must be positive.") + cfg = parse_env_cfg(_TASK, device=args.device, num_envs=args.num_envs) + cfg.play_mode() + cfg.scene.num_envs = args.num_envs + cfg.seed = args.seed + cfg.events.reset_from_state_table.params["arm_joint_noise"] = 0.0 + cfg.sim.use_fabric = False + cfg.sim.save_logs_to_file = False + cfg.sim.physics.use_cuda_graph = torch.device(args.device).type == "cuda" + agent = load_cfg_from_registry(_TASK, "rsl_rl_cfg_entry_point") + agent.device = args.device + checkpoint = ( + resolve_published_checkpoint("rsl_rl", _TASK, cfg) + if args.checkpoint == "pretrained" + else retrieve_file_path(args.checkpoint) + ) + if checkpoint is None: + raise FileNotFoundError("No published conveyor checkpoint was found; pass --checkpoint /path/to/model.pt.") + with launch_simulation(cfg, args), closing(gym.make(_TASK, cfg=cfg).unwrapped) as env: + wrapped = RslRlVecEnvWrapper(env, clip_actions=agent.clip_actions) + runner = OnPolicyRunner(wrapped, to_dict(agent), device=agent.device) + runner.load(checkpoint, map_location=args.device) + policy = runner.get_inference_policy(device=args.device) + env.reset(seed=args.seed) + observations = wrapped.get_observations() + transfers = torch.zeros(2, device=env.device) + failures = 0 + nonfinite_observations = 0 + with torch.inference_mode(): + for _ in range(args.steps): + observations, _, dones, info = wrapped.step(policy(observations)) + policy.reset(dones) + ended = int(dones.sum()) + failures += ended + nonfinite_observations += int((~torch.isfinite(observations["policy"])).any(dim=1).sum()) + if ended: + for side, direction in enumerate(("left_to_right", "right_to_left")): + transfers[side] += info["log"][f"Metrics/transfer/{direction}_transfers"] * ended + transfers += env.command_manager.get_term("transfer").direction_transfer_counts.sum(dim=0) + seconds = args.steps * env.step_dt * args.num_envs + result = { + "task": _TASK, + "checkpoint": checkpoint, + "seed": args.seed, + "num_envs": args.num_envs, + "steps": args.steps, + "simulated_environment_seconds": seconds, + "completed_transfers": int(transfers.sum().round()), + "left_to_right_transfers": int(transfers[0].round()), + "right_to_left_transfers": int(transfers[1].round()), + "failures": failures, + "nonfinite_observations": nonfinite_observations, + "transfers_per_environment_minute": float(transfers.sum()) * 60.0 / seconds, + } + report = json.dumps(result, indent=2) + print(report, flush=True) + if args.output is not None: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(report + "\n") + + +if __name__ == "__main__": + main() diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/franka_robot_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/franka_robot_cfg.py new file mode 100644 index 000000000000..aaa4e4b01ce5 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/franka_robot_cfg.py @@ -0,0 +1,53 @@ +# 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 + +"""Task-calibrated Franka configuration for conveyor manipulation.""" + +from isaaclab_newton.sim.schemas import MujocoJointCfg + +from isaaclab.actuators import ImplicitActuatorCfg + +from isaaclab_assets.robots.franka import FRANKA_PANDA_MENAGERIE_CFG + +FRANKA_PANDA_CONVEYOR_CFG = FRANKA_PANDA_MENAGERIE_CFG.copy() +FRANKA_PANDA_CONVEYOR_CFG.spawn.rigid_props.disable_gravity = False +# Route gravity compensation through MuJoCo's actuator channel so effort limits +# and the solver apply it consistently with the configured implicit drives. +FRANKA_PANDA_CONVEYOR_CFG.spawn.joint_drive_props = [MujocoJointCfg(actuatorgravcomp=True)] +FRANKA_PANDA_CONVEYOR_CFG.actuators = { + "panda_arm": ImplicitActuatorCfg( + joint_names_expr=["panda_joint[1-7]"], + effort_limit_sim={"panda_joint[1-4]": 87.0, "panda_joint[5-7]": 12.0}, + velocity_limit_sim={"panda_joint[1-4]": 20.0, "panda_joint[5-7]": 25.0}, + stiffness={ + "panda_joint[1-4]": 600.0, + "panda_joint5": 250.0, + "panda_joint6": 150.0, + "panda_joint7": 50.0, + }, + damping={ + "panda_joint[1-4]": 50.0, + "panda_joint5": 30.0, + "panda_joint6": 25.0, + "panda_joint7": 15.0, + }, + armature={ + "panda_joint[1-2]": 0.6057, + "panda_joint[3-4]": 0.4625, + "panda_joint[5-7]": 0.2055, + }, + ), + "panda_hand": ImplicitActuatorCfg( + joint_names_expr=["panda_finger_joint[1-2]"], + effort_limit_sim=70.0, + velocity_limit_sim=2.0, + stiffness=350.0, + # Keep the 0.1 kg m^2 armature close to critical damping so the + # fingers establish contact within a few 50 Hz policy steps. + damping=10.0, + armature=0.1, + ), +} +"""Menagerie Franka with explicit manipulation gains and solver-native gravity compensation.""" diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/__init__.py new file mode 100644 index 000000000000..56b9c4494587 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/__init__.py @@ -0,0 +1,10 @@ +# 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 + +"""MDP terms for the conveyor-to-conveyor Franka transfer task.""" + +from isaaclab.utils.module import lazy_export + +lazy_export() diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/__init__.pyi b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/__init__.pyi new file mode 100644 index 000000000000..5fc03c8498f2 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/__init__.pyi @@ -0,0 +1,85 @@ +# 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 + +__all__ = [ + "BELT_DEPLOYMENT_VARIANT", + "BetaSamplingStrategyCfg", + "ConveyorRelativeJointPositionAction", + "ConveyorRelativeJointPositionActionCfg", + "ConveyorResetCurriculum", + "ConveyorResetRecipe", + "ConveyorResetStateTable", + "ConveyorTransferCommand", + "ConveyorTransferCommandCfg", + "ResetBufferedGripperAction", + "ResetBufferedGripperActionCfg", + "SamplerCfg", + "SuccessMonitorCfg", + "UniformSamplingStrategyCfg", + "action_term_l2", + "active_transfer_features", + "build_reset_rows", + "cube_conveyor_state", + "cube_out_of_workspace", + "end_effector_axes", + "end_effector_velocity", + "finite_action_rate_l2", + "finite_joint_velocity_l2", + "gripper_joint_positions", + "invalid_action", + "nonfinite_scene_state", + "select_next_transfer_cube", + "subgoal_time_out", + "target_cube_one_hot", + "target_side_one_hot", + "terminal_failure", + "transfer_object_observation", + "transfer_sequence_time_out", + "transfer_success_mask", + "transfer_success_reward", +] + +from .actions import ConveyorRelativeJointPositionAction, ResetBufferedGripperAction +from .actions_cfg import ConveyorRelativeJointPositionActionCfg, ResetBufferedGripperActionCfg +from .commands import ConveyorTransferCommand, ConveyorTransferCommandCfg, transfer_success_mask +from .curriculums import ConveyorResetCurriculum +from .observations import ( + active_transfer_features, + cube_conveyor_state, + end_effector_axes, + end_effector_velocity, + gripper_joint_positions, + target_cube_one_hot, + target_side_one_hot, + transfer_object_observation, +) +from .reset_events import ( + BELT_DEPLOYMENT_VARIANT, + ConveyorResetRecipe, + ConveyorResetStateTable, + build_reset_rows, + select_next_transfer_cube, +) +from .rewards import ( + action_term_l2, + finite_action_rate_l2, + finite_joint_velocity_l2, + terminal_failure, + transfer_success_reward, +) +from .terminations import ( + cube_out_of_workspace, + invalid_action, + nonfinite_scene_state, + subgoal_time_out, + transfer_sequence_time_out, +) +from isaaclab.envs.mdp import * +from isaaclab_tasks.contrib.nist.utils.sampling import ( + BetaSamplingStrategyCfg, + SamplerCfg, + UniformSamplingStrategyCfg, +) +from isaaclab_tasks.utils.success_monitor import SuccessMonitorCfg diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/actions.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/actions.py new file mode 100644 index 000000000000..5ca9f7c14a4b --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/actions.py @@ -0,0 +1,132 @@ +# 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 + +"""Reset-safe relative joint actions for conveyor transfer.""" + +from __future__ import annotations + +import math +from collections.abc import Sequence +from typing import TYPE_CHECKING + +import torch + +from isaaclab.envs.mdp.actions.binary_joint_actions import BinaryJointPositionAction +from isaaclab.envs.mdp.actions.joint_actions import JointAction + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedEnv + + from .actions_cfg import ConveyorRelativeJointPositionActionCfg, ResetBufferedGripperActionCfg + + +class ConveyorRelativeJointPositionAction(JointAction): + """Apply one measured-state-relative target per policy step. + + Isaac Lab's generic relative action adds the residual during every physics + substep. This term computes the target once in :meth:`process_actions`, so + simulation decimation does not multiply the requested displacement. + """ + + cfg: ConveyorRelativeJointPositionActionCfg + + def __init__(self, cfg: ConveyorRelativeJointPositionActionCfg, env: ManagerBasedEnv) -> None: + super().__init__(cfg, env) + if not math.isfinite(cfg.max_delta) or cfg.max_delta <= 0.0: + raise ValueError("max_delta must be finite and positive.") + if not math.isfinite(cfg.joint_limit_margin) or cfg.joint_limit_margin < 0.0: + raise ValueError("joint_limit_margin must be finite and non-negative.") + self._workspace_lower = torch.tensor(cfg.workspace_lower, dtype=torch.float32, device=self.device) + self._workspace_upper = torch.tensor(cfg.workspace_upper, dtype=torch.float32, device=self.device) + if self._workspace_lower.shape != (self.action_dim,) or self._workspace_upper.shape != (self.action_dim,): + raise ValueError("workspace bounds must contain one value per controlled joint.") + if torch.any(self._workspace_lower >= self._workspace_upper): + raise ValueError("Every lower workspace bound must be less than its upper bound.") + if not torch.all(torch.isfinite(self._workspace_lower)) or not torch.all(torch.isfinite(self._workspace_upper)): + raise ValueError("workspace bounds must be finite.") + self._previous_actions = torch.zeros_like(self._raw_actions) + self._invalid_actions = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device) + self._processed_actions = self._asset.data.joint_pos.torch[:, self._joint_ids].clone() + + @property + def previous_actions(self) -> torch.Tensor: + """Previous finite policy actions, shape ``(num_envs, action_dim)``.""" + return self._previous_actions + + @property + def invalid_actions(self) -> torch.Tensor: + """Whether the latest policy action contained a non-finite component.""" + return self._invalid_actions + + def process_actions(self, actions: torch.Tensor) -> None: + """Convert normalized residuals into bounded position targets [rad].""" + self._previous_actions.copy_(self._raw_actions) + self._invalid_actions.copy_(~torch.isfinite(actions).all(dim=1)) + finite_actions = torch.nan_to_num(actions, nan=0.0, posinf=1.0, neginf=-1.0).clamp(-1.0, 1.0) + super().process_actions(finite_actions) + delta = torch.clamp(self._processed_actions, min=-self.cfg.max_delta, max=self.cfg.max_delta) + positions = self._asset.data.joint_pos.torch[:, self._joint_ids] + limits = self._asset.data.soft_joint_pos_limits.torch[:, self._joint_ids] + lower = torch.maximum(limits[..., 0] + self.cfg.joint_limit_margin, self._workspace_lower) + upper = torch.minimum(limits[..., 1] - self.cfg.joint_limit_margin, self._workspace_upper) + self._processed_actions = torch.clamp(positions + delta, min=lower, max=upper) + + def apply_actions(self) -> None: + """Hold the policy-step target through all physics substeps.""" + self._asset.actuators.target_command.set_position_index( + value=self._processed_actions, joint_ids=self._joint_ids + ) + + def reset(self, env_ids: Sequence[int] | torch.Tensor | None = None) -> None: + """Initialize targets from the sampled reset pose.""" + super().reset(env_ids) + env_ids = slice(None) if env_ids is None else env_ids + positions = self._asset.data.joint_pos.torch[:, self._joint_ids] + self._processed_actions[env_ids] = positions[env_ids] + self._previous_actions[env_ids] = 0.0 + self._invalid_actions[env_ids] = False + + +class ResetBufferedGripperAction(BinaryJointPositionAction): + """Keep reset-authored grasps closed during a short settling window.""" + + cfg: ResetBufferedGripperActionCfg + + def __init__(self, cfg: ResetBufferedGripperActionCfg, env: ManagerBasedEnv) -> None: + super().__init__(cfg, env) + self._previous_actions = torch.zeros_like(self._raw_actions) + self._invalid_actions = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device) + + @property + def previous_actions(self) -> torch.Tensor: + """Previous finite policy actions, shape ``(num_envs, action_dim)``.""" + return self._previous_actions + + @property + def invalid_actions(self) -> torch.Tensor: + """Whether the latest gripper command contained a non-finite value.""" + return self._invalid_actions + + def process_actions(self, actions: torch.Tensor) -> None: + """Map finite binary commands and preserve initially held cubes.""" + self._previous_actions.copy_(self._raw_actions) + if actions.dtype == torch.bool: + self._invalid_actions.zero_() + finite_actions = actions + else: + self._invalid_actions.copy_(~torch.isfinite(actions).all(dim=1)) + # A non-finite gripper command closes the fingers for the final safe + # step before the invalid-action termination is evaluated. + finite_actions = torch.nan_to_num(actions, nan=-1.0, posinf=1.0, neginf=-1.0).clamp(-1.0, 1.0) + super().process_actions(finite_actions) + command = self._env.command_manager.get_term(self.cfg.command_name) + force_close = (command.held_cube_ids >= 0) & (self._env.episode_length_buf < self.cfg.force_close_steps) + self._processed_actions[force_close] = self._close_command + + def reset(self, env_ids: Sequence[int] | None = None) -> None: + """Clear buffered invalid-command state for selected environments.""" + super().reset(env_ids) + self._previous_actions[env_ids] = 0.0 + self._invalid_actions[env_ids] = False diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/actions_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/actions_cfg.py new file mode 100644 index 000000000000..40bc923f779e --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/actions_cfg.py @@ -0,0 +1,48 @@ +# 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 + +"""Action configurations for conveyor transfer.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from isaaclab.envs.mdp.actions.actions_cfg import BinaryJointPositionActionCfg, JointActionCfg +from isaaclab.utils.configclass import configclass + +if TYPE_CHECKING: + from .actions import ConveyorRelativeJointPositionAction, ResetBufferedGripperAction + + +@configclass +class ConveyorRelativeJointPositionActionCfg(JointActionCfg): + """Configuration for measured-state relative Franka joint control.""" + + class_type: type[ConveyorRelativeJointPositionAction] | str = "{DIR}.actions:ConveyorRelativeJointPositionAction" + + joint_limit_margin: float = 0.02 + """Distance kept from each soft joint limit [rad].""" + + max_delta: float = 0.12 + """Maximum target change per policy step [rad].""" + + workspace_lower: tuple[float, ...] = (-0.75, -0.45, -0.55, -2.75, -0.45, 1.85, -0.10) + """Lower boundary of the validated transfer workspace [rad].""" + + workspace_upper: tuple[float, ...] = (0.85, 0.85, 0.35, -1.75, 0.45, 3.05, 1.65) + """Upper boundary of the validated transfer workspace [rad].""" + + +@configclass +class ResetBufferedGripperActionCfg(BinaryJointPositionActionCfg): + """Configuration for reset-grasp protection.""" + + class_type: type[ResetBufferedGripperAction] | str = "{DIR}.actions:ResetBufferedGripperAction" + + force_close_steps: int = 5 + """Initial policy steps that preserve a reset-authored grasp.""" + + command_name: str = "transfer" + """Command term that owns reset-authored held-cube state.""" diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/commands.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/commands.py new file mode 100644 index 000000000000..3b4d859d1405 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/commands.py @@ -0,0 +1,286 @@ +# 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 + +"""Success-driven transfer commands for the conveyor Franka task.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import TYPE_CHECKING + +import torch + +from isaaclab.managers import CommandTerm, CommandTermCfg +from isaaclab.utils.configclass import configclass + +from ..conveyor_geometry import BELT_CENTER_X, BELT_HALF_STRAIGHT +from .cube_state import cube_values +from .kinematics import end_effector_pose +from .reset_events import CUBE_COUNT, select_next_transfer_cube, side_inner_y + +if TYPE_CHECKING: + from isaaclab.assets import Articulation + from isaaclab.envs import ManagerBasedRLEnv + + +def transfer_success_mask( + cube_positions: torch.Tensor, + cube_linear_velocities: torch.Tensor, + tool_positions: torch.Tensor, + finger_positions: torch.Tensor, + target_side_ids: torch.Tensor, + lateral_tolerance: float = 0.055, + maximum_cube_speed: float = 0.65, + minimum_finger_position: float = 0.027, + minimum_tool_clearance: float = 0.055, +) -> torch.Tensor: + """Return whether the active cube is released on its destination belt.""" + target_y = side_inner_y(target_side_ids) + on_straight = torch.abs(cube_positions[:, 0] - BELT_CENTER_X) < BELT_HALF_STRAIGHT + on_lane = torch.abs(cube_positions[:, 1] - target_y) < lateral_tolerance + supported_height = (cube_positions[:, 2] > 0.045) & (cube_positions[:, 2] < 0.095) + moving_safely = torch.linalg.vector_norm(cube_linear_velocities, dim=1) < maximum_cube_speed + released = torch.amin(finger_positions, dim=1) > minimum_finger_position + hand_clear = torch.linalg.vector_norm(tool_positions - cube_positions, dim=1) > minimum_tool_clearance + return on_straight & on_lane & supported_height & moving_safely & released & hand_clear + + +class ConveyorTransferCommand(CommandTerm): + """Command one numbered cube to the opposite belt and redraw on success.""" + + cfg: ConveyorTransferCommandCfg + + def __init__(self, cfg: ConveyorTransferCommandCfg, env: ManagerBasedRLEnv) -> None: + self._validate_cfg(cfg) + super().__init__(cfg, env) + + reset_term = env.event_manager.get_term_cfg(cfg.reset_event_name).func + required_reset_fields = ("row_ids", "target_cube_ids", "source_side_ids", "held_rows") + if not all(hasattr(reset_term, name) for name in required_reset_fields): + raise RuntimeError("ConveyorTransferCommand requires ConveyorResetStateTable reset metadata.") + self._reset_term = reset_term + self._robot: Articulation = env.scene["robot"] + self._finger_joint_ids = self._robot.find_joints("panda_finger_joint[1-2]", preserve_order=True)[0] + + self.target_cube_ids = torch.zeros(self.num_envs, dtype=torch.long, device=self.device) + self.source_side_ids = torch.zeros_like(self.target_cube_ids) + self.held_cube_ids = torch.full_like(self.target_cube_ids, -1) + self.subgoal_start_steps = torch.zeros_like(self.target_cube_ids) + self.transfer_counts = torch.zeros_like(self.target_cube_ids) + self.direction_transfer_counts = torch.zeros((self.num_envs, 2), dtype=torch.long, device=self.device) + + self._stable_steps = torch.zeros_like(self.target_cube_ids) + self.is_success = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device) + self.new_success = torch.zeros_like(self.is_success) + self.pending_success = torch.zeros_like(self.is_success) + self.ever_success = torch.zeros_like(self.is_success) + self._last_evaluation_steps = torch.full_like(self.target_cube_ids, -1) + self._resampling_from_reset = False + + self.metrics["success_rate"] = torch.zeros(self.num_envs, dtype=torch.float32, device=self.device) + self.metrics["transfer_count"] = torch.zeros_like(self.metrics["success_rate"]) + self.metrics["left_to_right_transfers"] = torch.zeros_like(self.metrics["success_rate"]) + self.metrics["right_to_left_transfers"] = torch.zeros_like(self.metrics["success_rate"]) + + @staticmethod + def _validate_cfg(cfg: ConveyorTransferCommandCfg) -> None: + """Validate stable transfer-completion thresholds.""" + if cfg.minimum_subgoal_steps < 0 or cfg.hold_steps < 1: + raise ValueError("minimum_subgoal_steps must be non-negative and hold_steps must be positive.") + if cfg.lateral_tolerance <= 0.0 or cfg.maximum_cube_speed <= 0.0: + raise ValueError("Conveyor placement tolerances and speed limits must be positive.") + if cfg.minimum_finger_position <= 0.0 or cfg.minimum_tool_clearance <= 0.0: + raise ValueError("Conveyor release thresholds must be positive.") + + @property + def command(self) -> torch.Tensor: + """Return target-cube and destination-belt one-hot commands.""" + cube = torch.nn.functional.one_hot(self.target_cube_ids, num_classes=CUBE_COUNT) + destination = torch.nn.functional.one_hot(1 - self.source_side_ids, num_classes=2) + return torch.cat((cube, destination), dim=1).float() + + def reset(self, env_ids: Sequence[int] | slice | None = None) -> dict[str, float]: + """Log per-command completion and initialize commands from sampled reset rows.""" + ids = self._resolve_env_ids(env_ids) + valid = self.command_counter[ids] > 0 + attempts = self.command_counter[ids].clamp_min(1) + self.metrics["success_rate"][ids] = torch.where( + valid, + self.transfer_counts[ids].float() / attempts.float(), + 0.0, + ) + self.metrics["transfer_count"][ids] = self.transfer_counts[ids].float() + self.metrics["left_to_right_transfers"][ids] = self.direction_transfer_counts[ids, 0].float() + self.metrics["right_to_left_transfers"][ids] = self.direction_transfer_counts[ids, 1].float() + + self._resampling_from_reset = True + try: + extras = super().reset(ids) + finally: + self._resampling_from_reset = False + + self._last_evaluation_steps[ids] = -1 + self._env.extras.setdefault("log", {})["Metrics/success_rate"] = extras.pop("success_rate") + return extras + + def evaluate(self) -> None: + """Evaluate stable transfer completion once per policy step.""" + evaluation_steps = self._env.episode_length_buf + evaluate_mask = (self.command_counter > 0) & (self._last_evaluation_steps != evaluation_steps) + if not bool(torch.any(evaluate_mask)): + return + + positions = cube_values(self._env, "root_pos_w") + velocities = cube_values(self._env, "root_lin_vel_w") + index = self.target_cube_ids.view(self.num_envs, 1, 1).expand(-1, 1, 3) + active_position = torch.gather(positions, 1, index).squeeze(1) - self._env.scene.env_origins + active_velocity = torch.gather(velocities, 1, index).squeeze(1) + tool_position, _ = end_effector_pose(self._env) + tool_position = tool_position - self._env.scene.env_origins + finger_positions = self._robot.data.joint_pos.torch[:, self._finger_joint_ids] + successful = transfer_success_mask( + active_position, + active_velocity, + tool_position, + finger_positions, + 1 - self.source_side_ids, + lateral_tolerance=self.cfg.lateral_tolerance, + maximum_cube_speed=self.cfg.maximum_cube_speed, + minimum_finger_position=self.cfg.minimum_finger_position, + minimum_tool_clearance=self.cfg.minimum_tool_clearance, + ) + subgoal_steps = evaluation_steps - self.subgoal_start_steps + successful &= subgoal_steps >= self.cfg.minimum_subgoal_steps + + next_stable_steps = torch.where(successful, self._stable_steps + 1, torch.zeros_like(self._stable_steps)) + self._stable_steps[evaluate_mask] = next_stable_steps[evaluate_mask] + stable = self._stable_steps >= self.cfg.hold_steps + new_success = evaluate_mask & stable & ~self.is_success & ~self.pending_success + self.is_success[evaluate_mask] = stable[evaluate_mask] + self.new_success.copy_(new_success) + self.pending_success |= new_success + self.ever_success |= new_success + success_ids = new_success.nonzero(as_tuple=False).squeeze(-1) + if success_ids.numel(): + source_sides = self.source_side_ids[success_ids] + self.transfer_counts[success_ids] += 1 + self.direction_transfer_counts[success_ids, source_sides] += 1 + + self._last_evaluation_steps[evaluate_mask] = evaluation_steps[evaluate_mask] + self._env.extras["successes"] = self.ever_success.clone() + + def set_goal(self, target_cube_id: int, env_ids: Sequence[int] | torch.Tensor | None = None) -> None: + """Replace the active command using the selected cube's current conveyor.""" + if not isinstance(target_cube_id, int) or isinstance(target_cube_id, bool): + raise TypeError("target_cube_id must be an integer.") + if not 0 <= target_cube_id < CUBE_COUNT: + raise ValueError(f"target_cube_id must lie in [0, {CUBE_COUNT - 1}].") + ids = self._resolve_env_ids(env_ids) + if ids.numel() == 0: + return + local_y = cube_values(self._env, "root_pos_w")[ids, target_cube_id, 1] - self._env.scene.env_origins[ids, 1] + source_side_ids = (local_y < 0.0).long() + target_cube_ids = torch.full_like(ids, target_cube_id) + self._assign_goal(ids, target_cube_ids, source_side_ids) + + def _update_metrics(self) -> None: + self.evaluate() + + def _resample_command(self, env_ids: Sequence[int]) -> None: + ids = self._resolve_env_ids(env_ids) + if ids.numel() == 0: + return + if self._resampling_from_reset: + rows = self._reset_term.row_ids[ids] + target_cube_ids = self._reset_term.target_cube_ids[rows] + source_side_ids = self._reset_term.source_side_ids[rows] + held_rows = self._reset_term.held_rows[rows] + self.transfer_counts[ids] = 0 + self.direction_transfer_counts[ids] = 0 + self.ever_success[ids] = False + self._assign_goal(ids, target_cube_ids, source_side_ids) + self.held_cube_ids[ids] = torch.where(held_rows, target_cube_ids, -1) + self.subgoal_start_steps[ids] = 0 + return + + positions = cube_values(self._env, "root_pos_w")[ids] + positions -= self._env.scene.env_origins[ids].unsqueeze(1) + next_source_side_ids = 1 - self.source_side_ids[ids] + next_cube_ids = select_next_transfer_cube( + positions, + self.target_cube_ids[ids], + next_source_side_ids, + transit_half_width=self.cfg.transit_half_width, + ) + self._assign_goal(ids, next_cube_ids, next_source_side_ids) + + def _update_command(self) -> None: + completed_ids = self.pending_success.nonzero(as_tuple=False).squeeze(-1) + if completed_ids.numel(): + self._resample(completed_ids) + + def _assign_goal( + self, + env_ids: torch.Tensor, + target_cube_ids: torch.Tensor, + source_side_ids: torch.Tensor, + ) -> None: + """Publish one command and clear completion state from the previous command.""" + self.target_cube_ids[env_ids] = target_cube_ids + self.source_side_ids[env_ids] = source_side_ids + self.held_cube_ids[env_ids] = -1 + self.subgoal_start_steps[env_ids] = self._env.episode_length_buf[env_ids] + self._stable_steps[env_ids] = 0 + self.is_success[env_ids] = False + self.new_success[env_ids] = False + self.pending_success[env_ids] = False + self._last_evaluation_steps[env_ids] = -1 + + def _resolve_env_ids(self, env_ids: Sequence[int] | torch.Tensor | slice | None) -> torch.Tensor: + """Return validated environment indices on the command device.""" + if env_ids is None: + ids = torch.arange(self.num_envs, dtype=torch.long, device=self.device) + elif isinstance(env_ids, slice): + ids = torch.arange(self.num_envs, dtype=torch.long, device=self.device)[env_ids] + else: + ids = torch.as_tensor(env_ids, dtype=torch.long, device=self.device).flatten() + if bool(torch.any((ids < 0) | (ids >= self.num_envs))): + raise IndexError("Conveyor command environment indices are out of range.") + return ids + + +@configclass +class ConveyorTransferCommandCfg(CommandTermCfg): + """Configuration for success-driven, bidirectional cube-transfer commands.""" + + class_type: type[ConveyorTransferCommand] | str = "{DIR}.commands:ConveyorTransferCommand" + """Command-term implementation resolved lazily by the manager.""" + + resampling_time_range: tuple[float, float] = (1.0e6, 1.0e6) + """Time-based resampling interval [s]; success normally resamples first.""" + + reset_event_name: str = "reset_from_state_table" + """Reset event containing the immutable physical-state table.""" + + transit_half_width: float = 0.14 + """Half-width of the corridor excluded when selecting a source-belt cube [m].""" + + minimum_subgoal_steps: int = 2 + """Minimum policy steps before a placement can succeed.""" + + hold_steps: int = 3 + """Consecutive policy steps for which a placement must remain valid.""" + + lateral_tolerance: float = 0.055 + """Maximum lateral error from the destination belt center [m].""" + + maximum_cube_speed: float = 0.65 + """Maximum cube speed for a completed placement [m/s].""" + + minimum_finger_position: float = 0.027 + """Minimum position of each finger for the cube to count as released [m].""" + + minimum_tool_clearance: float = 0.055 + """Minimum tool-to-cube distance for the cube to count as released [m].""" diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/cube_state.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/cube_state.py new file mode 100644 index 000000000000..9ce81862d05b --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/cube_state.py @@ -0,0 +1,22 @@ +# 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 + +"""Rigid-object data in the policy's four fixed cube identities.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from .reset_events import CUBE_COUNT + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedRLEnv + + +def cube_values(env: ManagerBasedRLEnv, attribute: str, *, all_cubes: bool = False) -> torch.Tensor: + """Gather cube data in identity order, with shape (num_envs, four, attribute_size).""" + return torch.stack([getattr(env.scene[f"cube_{i}"].data, attribute).torch for i in range(CUBE_COUNT)], dim=1) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/curriculums.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/curriculums.py new file mode 100644 index 000000000000..7824f81bbee7 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/curriculums.py @@ -0,0 +1,99 @@ +# 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 + +"""Success-guided reset sampling for conveyor transfer.""" + +from __future__ import annotations + +import math +from collections.abc import Sequence +from typing import TYPE_CHECKING + +import torch + +from isaaclab.managers import CurriculumTermCfg, ManagerTermBase +from isaaclab.utils import instantiate + +from isaaclab_tasks.utils.success_monitor import SuccessMonitorCfg + +from ...nist.utils.sampling import Sampler, SamplerCfg +from .reset_events import ConveyorResetRecipe + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedRLEnv + + +class ConveyorResetCurriculum(ManagerTermBase): + """Sample reset rows using the shared NIST sampler and completed-transfer outcomes. + + Reset geometry and command metadata belong to the reset term. Only recent success + evidence lives here; standard PPO checkpoints start a fresh sampling history on resume. + """ + + def __init__(self, cfg: CurriculumTermCfg, env: ManagerBasedRLEnv) -> None: + super().__init__(cfg, env) + self._reset_term = env.event_manager.get_term_cfg("reset_from_state_table").func + self.success_monitor = instantiate(cfg.params["success_monitor"], 1, self._reset_term.row_count, env.device) + self.sampler: Sampler = instantiate(cfg.params["sampling"], self.success_monitor.success_rate) + + def __call__( + self, + env: ManagerBasedRLEnv, + env_ids: Sequence[int] | torch.Tensor | slice, + success_monitor: SuccessMonitorCfg, + sampling: SamplerCfg, + command_name: str = "transfer", + fixed_source_side_id: int | None = None, + ) -> dict[str, torch.Tensor]: + """Record finished episodes and draw rows for only the resetting environments.""" + del success_monitor, sampling + ids = ( + env.scene._ALL_INDICES[env_ids] + if isinstance(env_ids, slice) + else torch.as_tensor(env_ids, dtype=torch.long, device=env.device).flatten() + ) + command = env.command_manager.get_term(command_name) + completed = (command.command_counter[ids] > 0) & (env.episode_length_buf[ids] > 0) + completed_ids = ids[completed] + succeeded = command.ever_success[completed_ids] + self.success_monitor.success_update(self._reset_term.row_ids[completed_ids], succeeded) + + probabilities = self.sampler.probabilities() + if fixed_source_side_id is not None: + if fixed_source_side_id not in (0, 1): + raise ValueError("fixed_source_side_id must be 0 (left) or 1 (right).") + probabilities.mul_(self._reset_term.source_side_ids == fixed_source_side_id) + probabilities.div_(probabilities.sum()) + if ids.numel(): + self._reset_term.row_ids[ids] = self.sampler.sample(probabilities, ids.numel()) + + history_size = self.success_monitor.success_size + history_successes = self.success_monitor.success_buf.sum(dim=1) + entropy = -(probabilities * probabilities.clamp_min(torch.finfo(probabilities.dtype).tiny).log()).sum() + metrics: dict[str, torch.Tensor] = { + "batch_success_rate": succeeded.float().sum() / max(completed_ids.numel(), 1), + "batch_transfer_count": command.transfer_counts[completed_ids].float().sum() + / max(completed_ids.numel(), 1), + "batch_left_to_right_transfers": command.direction_transfer_counts[completed_ids, 0].float().sum() + / max(completed_ids.numel(), 1), + "batch_right_to_left_transfers": command.direction_transfer_counts[completed_ids, 1].float().sum() + / max(completed_ids.numel(), 1), + "row_coverage": (history_size > 0).float().mean(), + "overall_success_rate": history_successes.sum() / history_size.sum().clamp_min(1), + "sampling_entropy": entropy / math.log(probabilities.numel()), + } + for recipe in ConveyorResetRecipe: + mask = self._reset_term.recipe_ids == int(recipe) + metrics[f"recipe_{recipe.name.lower()}_probability"] = probabilities[mask].sum() + metrics[f"recipe_{recipe.name.lower()}_success_rate"] = history_successes[mask].sum() / history_size[ + mask + ].sum().clamp_min(1) + for side_id, side_name in ((0, "left_to_right"), (1, "right_to_left")): + mask = self._reset_term.source_side_ids == side_id + metrics[f"direction_{side_name}_probability"] = probabilities[mask].sum() + metrics[f"direction_{side_name}_success_rate"] = history_successes[mask].sum() / history_size[ + mask + ].sum().clamp_min(1) + return metrics diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/kinematics.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/kinematics.py new file mode 100644 index 000000000000..9380bc345a46 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/kinematics.py @@ -0,0 +1,40 @@ +# 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 + +"""Backend-independent Franka tool-state helpers.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from isaaclab.utils import math as math_utils + +if TYPE_CHECKING: + from isaaclab.assets import Articulation + + from ..conveyor_franka_env import ConveyorFrankaEnv + + +def end_effector_pose(env: ConveyorFrankaEnv) -> tuple[torch.Tensor, torch.Tensor]: + """Return the Franka tool-center position [m] and orientation.""" + robot: Articulation = env.scene["robot"] + body_id = env._conveyor_hand_body_id + orientation = robot.data.body_quat_w.torch[:, body_id] + position = robot.data.body_pos_w.torch[:, body_id] + position = position + math_utils.quat_apply(orientation, env._conveyor_tool_offset) + return position, orientation + + +def tool_velocity(env: ConveyorFrankaEnv) -> torch.Tensor: + """Return tool-center linear and angular velocity [m/s, rad/s].""" + robot: Articulation = env.scene["robot"] + body_id = env._conveyor_hand_body_id + orientation = robot.data.body_quat_w.torch[:, body_id] + body_velocity = robot.data.body_link_vel_w.torch[:, body_id] + offset_world = math_utils.quat_apply(orientation, env._conveyor_tool_offset) + linear_velocity = body_velocity[:, :3] + torch.linalg.cross(body_velocity[:, 3:], offset_world) + return torch.cat((linear_velocity, body_velocity[:, 3:]), dim=1) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/observations.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/observations.py new file mode 100644 index 000000000000..9e690ff58ecf --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/observations.py @@ -0,0 +1,133 @@ +# 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 + +"""Task-conditioned observations for conveyor transfer.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from isaaclab.managers import SceneEntityCfg +from isaaclab.utils import math as math_utils + +from .cube_state import cube_values +from .kinematics import end_effector_pose, tool_velocity +from .reset_events import CUBE_COUNT, TRANSFER_X, side_inner_y + +if TYPE_CHECKING: + from isaaclab.assets import Articulation + from isaaclab.envs import ManagerBasedRLEnv + + +def _cube_state(env: ManagerBasedRLEnv) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Stack cube world positions, orientations, and spatial velocities.""" + state = ( + cube_values(env, "root_pos_w"), + cube_values(env, "root_quat_w"), + cube_values(env, "root_vel_w"), + ) + return state + + +def target_cube_one_hot(env: ManagerBasedRLEnv, command_name: str = "transfer") -> torch.Tensor: + """Encode which numbered cube the policy must transfer.""" + command = env.command_manager.get_term(command_name) + return torch.nn.functional.one_hot(command.target_cube_ids.long(), num_classes=CUBE_COUNT).float() + + +def target_side_one_hot(env: ManagerBasedRLEnv, command_name: str = "transfer") -> torch.Tensor: + """Encode the destination conveyor, opposite the reset source side.""" + command = env.command_manager.get_term(command_name) + return torch.nn.functional.one_hot(1 - command.source_side_ids.long(), num_classes=2).float() + + +def classify_cube_conveyors(local_positions: torch.Tensor, transit_half_width: float = 0.14) -> torch.Tensor: + """Classify positions as left conveyor, in transit, or right conveyor.""" + if local_positions.shape[-1] != 3: + raise ValueError("Cube positions must end in xyz coordinates.") + side_ids = torch.full(local_positions.shape[:-1], 1, dtype=torch.long, device=local_positions.device) + side_ids[local_positions[..., 1] > transit_half_width] = 0 + side_ids[local_positions[..., 1] < -transit_half_width] = 2 + return torch.nn.functional.one_hot(side_ids, num_classes=3).float() + + +def cube_conveyor_state(env: ManagerBasedRLEnv) -> torch.Tensor: + """Return left/transit/right one-hot state for every numbered cube.""" + positions, _, _ = _cube_state(env) + local_positions = positions - env.scene.env_origins.unsqueeze(1) + return classify_cube_conveyors(local_positions).flatten(start_dim=1) + + +def transfer_object_observation(env: ManagerBasedRLEnv) -> torch.Tensor: + """Describe all four cubes in stable identity slots. + + The observation contains local positions, tool-relative positions, local + up axes, and linear/angular velocities in four fixed identity slots. + :func:`target_cube_one_hot` selects the active slot. + """ + positions, quaternions, velocities = _cube_state(env) + local_positions = positions - env.scene.env_origins.unsqueeze(1) + tool_position, _ = end_effector_pose(env) + tool_relative = positions - tool_position.unsqueeze(1) + rotations = math_utils.matrix_from_quat(quaternions.flatten(end_dim=1)).view(env.num_envs, CUBE_COUNT, 3, 3) + up_axes = rotations[..., 2] + return torch.cat( + ( + local_positions.flatten(start_dim=1), + tool_relative.flatten(start_dim=1), + up_axes.flatten(start_dim=1), + velocities.flatten(start_dim=1), + ), + dim=1, + ) + + +def active_transfer_features(env: ManagerBasedRLEnv, command_name: str = "transfer") -> torch.Tensor: + """Return active-cube and destination-relative position features [m].""" + command = env.command_manager.get_term(command_name) + positions, _, _ = _cube_state(env) + active_position = positions[torch.arange(env.num_envs, device=env.device), command.target_cube_ids.long()] + local_active_position = active_position - env.scene.env_origins + tool_position, _ = end_effector_pose(env) + target_side_ids = 1 - command.source_side_ids.long() + target_position = torch.stack( + ( + torch.full_like(target_side_ids, TRANSFER_X, dtype=active_position.dtype), + side_inner_y(target_side_ids), + torch.full_like(target_side_ids, 0.06, dtype=active_position.dtype), + ), + dim=1, + ) + return torch.cat( + ( + local_active_position, + active_position - tool_position, + target_position - local_active_position, + ), + dim=1, + ) + + +def gripper_joint_positions( + env: ManagerBasedRLEnv, + robot_cfg: SceneEntityCfg = SceneEntityCfg("robot", joint_names=["panda_finger_joint[1-2]"]), +) -> torch.Tensor: + """Return the two Franka finger positions [m].""" + robot: Articulation = env.scene[robot_cfg.name] + return robot.data.joint_pos.torch[:, robot_cfg.joint_ids] + + +def end_effector_axes(env: ManagerBasedRLEnv) -> torch.Tensor: + """Return continuous tool-frame x and z axes.""" + _, orientation = end_effector_pose(env) + rotation = math_utils.matrix_from_quat(orientation) + return torch.cat((rotation[:, :, 0], rotation[:, :, 2]), dim=1) + + +def end_effector_velocity(env: ManagerBasedRLEnv) -> torch.Tensor: + """Return tool-center linear and angular velocity [m/s, rad/s].""" + return tool_velocity(env) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/reset_events.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/reset_events.py new file mode 100644 index 000000000000..e0b41ba43354 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/reset_events.py @@ -0,0 +1,482 @@ +# 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 + +"""Reset-state sampling for conveyor transfer.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from enum import IntEnum +from typing import TYPE_CHECKING + +import torch + +from isaaclab.managers import EventTermCfg, ManagerTermBase + +from ..conveyor_geometry import ( + BELT_CENTER_X, + BELT_INNER_STRAIGHT_Y, + BELT_OUTER_STRAIGHT_Y, + BELT_TOP_Z, + CUBE_INNER_SLOT_X, + CUBE_OUTER_SLOT_X, +) + +if TYPE_CHECKING: + from isaaclab.assets import Articulation, RigidObject + from isaaclab.envs import ManagerBasedRLEnv + + +CUBE_COUNT = 4 +CUBE_SIZE = 0.04 +CUBE_REST_Z = BELT_TOP_Z + 0.5 * CUBE_SIZE +TRANSFER_X = 0.52 +LEFT_SIDE = 0 +RIGHT_SIDE = 1 +_BELT_CURRICULUM_FRACTIONS = (0.0, 0.15, 0.30, 0.45, 0.60, 0.80, 1.0) +_GRASP_CLOSURE_FRACTIONS = (0.25, 0.50, 0.75, 0.90, 1.0) +_OPEN_FINGER_POSITION = 0.040 +_CLOSED_FINGER_POSITION = 0.019 +_MINIMUM_CUBE_SEPARATION = 0.055 +BELT_DEPLOYMENT_VARIANT = len(_BELT_CURRICULUM_FRACTIONS) - 1 + + +class ConveyorResetRecipe(IntEnum): + """Reset phases ordered from easiest to complete task start.""" + + GOAL = 0 + PLACE = 1 + CARRY = 2 + LIFT = 3 + GRASP = 4 + PREGRASP = 5 + BELT = 6 + + +@dataclass(frozen=True) +class ConveyorResetRow: + """One physical reset state and transfer command.""" + + recipe: ConveyorResetRecipe + variant_id: int + target_cube_id: int + source_side_id: int + arm_positions: tuple[float, ...] + finger_position: float + held: bool + belt_range_fraction: float + + +_HOME_ARM = (0.0, -0.35, 0.0, -2.35, 0.0, 2.0, 0.78) + +# Constrained IK anchors target x=0.52 m with a downward-facing tool. Side 0 +# is the positive-y conveyor and side 1 is the negative-y conveyor. +_SOURCE_GRASP_ARM = ( + (0.5144946, 0.5427007, -0.0340481, -2.0731085, 0.0350549, 2.6153000, 1.2350098), + (-0.3400387, 0.5475255, -0.1330232, -2.0714802, 0.1368175, 2.6109921, 0.2080209), +) +_SOURCE_PREGRASP_ARM = ( + (0.5190744, 0.3742786, -0.0406565, -2.0872751, 0.0236238, 2.4611441, 1.2428656), + (-0.3220402, 0.3787346, -0.1588824, -2.0865274, 0.0928583, 2.4588233, 0.2380420), +) +_SOURCE_LIFT_ARM = ( + (0.5212608, 0.2372272, -0.0439391, -2.0514939, 0.0137043, 2.2884652, 1.2495379), + (-0.3132103, 0.2404478, -0.1720021, -2.0512362, 0.0540920, 2.2875542, 0.2640870), +) +_TARGET_PLACE_ARM = ( + # Source left, target right. + (-0.2921792, 0.4521663, -0.1853937, -2.0854943, 0.1398595, 2.5262992, 0.2062777), + # Source right, target left. + (0.4780173, 0.4442472, 0.0008774, -2.0873070, -0.0005178, 2.5315716, 1.2591928), +) +_CARRY_ARM = (0.0, 0.0137827, 0.0, -2.2661237, 0.0, 2.2798989, 0.78) + + +def _interpolate_arm( + start: tuple[float, ...], + end: tuple[float, ...], + fractions: tuple[float, ...], +) -> tuple[tuple[float, ...], ...]: + """Linearly interpolate validated joint-space anchors.""" + return tuple( + tuple( + start_value + fraction * (end_value - start_value) + for start_value, end_value in zip(start, end, strict=True) + ) + for fraction in fractions + ) + + +def _arm_position_variants( + recipe: ConveyorResetRecipe, + source_side_id: int, +) -> tuple[tuple[float, ...], ...]: + """Return dense reset states along the nominal transfer trajectory.""" + source_grasp = _SOURCE_GRASP_ARM[source_side_id] + source_pregrasp = _SOURCE_PREGRASP_ARM[source_side_id] + source_lift = _SOURCE_LIFT_ARM[source_side_id] + target_lift = _SOURCE_LIFT_ARM[1 - source_side_id] + if recipe == ConveyorResetRecipe.GOAL: + return (_SOURCE_PREGRASP_ARM[1 - source_side_id],) + if recipe == ConveyorResetRecipe.PLACE: + return _interpolate_arm(target_lift, _TARGET_PLACE_ARM[source_side_id], (0.25, 0.50, 0.75, 1.0)) + if recipe == ConveyorResetRecipe.CARRY: + return ( + *_interpolate_arm(source_lift, _CARRY_ARM, (0.33, 0.66, 1.0)), + *_interpolate_arm(_CARRY_ARM, target_lift, (0.33, 0.66, 1.0)), + ) + if recipe == ConveyorResetRecipe.LIFT: + return _interpolate_arm(source_grasp, source_lift, (0.25, 0.50, 0.75, 1.0)) + if recipe == ConveyorResetRecipe.GRASP: + return (source_grasp,) * len(_GRASP_CLOSURE_FRACTIONS) + if recipe == ConveyorResetRecipe.PREGRASP: + # Include the open-gripper contact pose before the closure trajectory. + return _interpolate_arm(source_pregrasp, source_grasp, (0.0, 0.50, 0.75, 0.92, 1.0)) + if recipe == ConveyorResetRecipe.BELT: + return _interpolate_arm(source_pregrasp, _HOME_ARM, _BELT_CURRICULUM_FRACTIONS) + raise ValueError(f"Unsupported reset recipe: {recipe}.") + + +def build_reset_rows() -> tuple[ConveyorResetRow, ...]: + """Build the complete identity, direction, and dense-trajectory cross product.""" + held_recipes = (ConveyorResetRecipe.LIFT, ConveyorResetRecipe.CARRY, ConveyorResetRecipe.PLACE) + return tuple( + ConveyorResetRow( + recipe=recipe, + variant_id=variant_id, + target_cube_id=cube_id, + source_side_id=source_side, + arm_positions=arm_positions, + finger_position=( + _OPEN_FINGER_POSITION + + _GRASP_CLOSURE_FRACTIONS[variant_id] * (_CLOSED_FINGER_POSITION - _OPEN_FINGER_POSITION) + if recipe == ConveyorResetRecipe.GRASP + else _CLOSED_FINGER_POSITION + if recipe in held_recipes + else _OPEN_FINGER_POSITION + ), + held=( + recipe in held_recipes + or (recipe == ConveyorResetRecipe.GRASP and variant_id == len(_GRASP_CLOSURE_FRACTIONS) - 1) + ), + belt_range_fraction=_BELT_CURRICULUM_FRACTIONS[variant_id] if recipe == ConveyorResetRecipe.BELT else 0.0, + ) + for recipe in ConveyorResetRecipe + for cube_id in range(CUBE_COUNT) + for source_side in (LEFT_SIDE, RIGHT_SIDE) + for variant_id, arm_positions in enumerate(_arm_position_variants(recipe, source_side)) + ) + + +_FRANKA_JOINT_ORIGINS = ( + ((0.0, 0.0, 0.333), (0.0, 0.0, 0.0)), + ((0.0, 0.0, 0.0), (-math.pi / 2.0, 0.0, 0.0)), + ((0.0, -0.316, 0.0), (math.pi / 2.0, 0.0, 0.0)), + ((0.0825, 0.0, 0.0), (math.pi / 2.0, 0.0, 0.0)), + ((-0.0825, 0.384, 0.0), (-math.pi / 2.0, 0.0, 0.0)), + ((0.0, 0.0, 0.0), (math.pi / 2.0, 0.0, 0.0)), + ((0.088, 0.0, 0.0), (math.pi / 2.0, 0.0, 0.0)), +) + + +def _rotation_matrix_from_rpy(roll: float, pitch: float, yaw: float, reference: torch.Tensor) -> torch.Tensor: + """Return a fixed XYZ roll-pitch-yaw rotation matrix.""" + cr, sr = math.cos(roll), math.sin(roll) + cp, sp = math.cos(pitch), math.sin(pitch) + cy, sy = math.cos(yaw), math.sin(yaw) + return reference.new_tensor( + ( + (cy * cp, cy * sp * sr - sy * cr, cy * sp * cr + sy * sr), + (sy * cp, sy * sp * sr + cy * cr, sy * sp * cr - cy * sr), + (-sp, cp * sr, cp * cr), + ) + ) + + +def franka_tool_position(joint_positions: torch.Tensor) -> torch.Tensor: + """Compute Panda tool-center positions [m] from seven joints [rad].""" + if joint_positions.shape[-1] != 7: + raise ValueError("Franka reset forward kinematics expects seven joint positions.") + shape = joint_positions.shape[:-1] + joints = joint_positions.reshape(-1, 7) + count = joints.shape[0] + rotation = torch.eye(3, dtype=joints.dtype, device=joints.device).expand(count, -1, -1).clone() + position = torch.zeros((count, 3), dtype=joints.dtype, device=joints.device) + reference = joints[0] if count else joint_positions.new_zeros(7) + for joint_id, (origin_position, origin_rpy) in enumerate(_FRANKA_JOINT_ORIGINS): + origin = joint_positions.new_tensor(origin_position).expand(count, -1) + position += torch.bmm(rotation, origin.unsqueeze(-1)).squeeze(-1) + rotation = torch.matmul(rotation, _rotation_matrix_from_rpy(*origin_rpy, reference=reference)) + angle = joints[:, joint_id] + cosine, sine = torch.cos(angle), torch.sin(angle) + zeros, ones = torch.zeros_like(angle), torch.ones_like(angle) + joint_rotation = torch.stack( + ( + torch.stack((cosine, -sine, zeros), dim=1), + torch.stack((sine, cosine, zeros), dim=1), + torch.stack((zeros, zeros, ones), dim=1), + ), + dim=1, + ) + rotation = torch.bmm(rotation, joint_rotation) + tool_offset = joint_positions.new_tensor((0.0, 0.0, 0.2104)).expand(count, -1) + position += torch.bmm(rotation, tool_offset.unsqueeze(-1)).squeeze(-1) + return position.reshape(*shape, 3) + + +def side_inner_y(side_ids: torch.Tensor) -> torch.Tensor: + """Return the reachable inner-straight y coordinate [m] for each side.""" + return torch.where(side_ids == LEFT_SIDE, BELT_INNER_STRAIGHT_Y, -BELT_INNER_STRAIGHT_Y) + + +def _balanced_cube_slots( + target_cube_ids: torch.Tensor, + source_side_ids: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Assign one cube to every inner/outer racetrack run. + + Slots are ordered ``left inner``, ``left outer``, ``right inner``, and + ``right outer``. The commanded cube swaps with the canonical occupant of + its source-side inner slot, keeping it reachable without duplicating or + emptying any deployment slot. + """ + count = target_cube_ids.numel() + cube_slots = torch.arange(CUBE_COUNT, device=target_cube_ids.device).expand(count, -1).clone() + source_inner_slots = 2 * source_side_ids + cube_slots.scatter_(1, target_cube_ids.unsqueeze(1), source_inner_slots.unsqueeze(1)) + cube_slots.scatter_(1, source_inner_slots.unsqueeze(1), target_cube_ids.unsqueeze(1)) + + cube_sides = torch.div(cube_slots, 2, rounding_mode="floor") + on_outer_run = torch.remainder(cube_slots, 2).bool() + cube_x = torch.where(on_outer_run, CUBE_OUTER_SLOT_X, CUBE_INNER_SLOT_X) + y_magnitude = torch.where(on_outer_run, BELT_OUTER_STRAIGHT_Y, BELT_INNER_STRAIGHT_Y) + cube_y = torch.where(cube_sides == LEFT_SIDE, y_magnitude, -y_magnitude) + return cube_slots, cube_sides, cube_x, cube_y + + +def _sample_collision_free_active_x( + base_x: torch.Tensor, + cube_sides: torch.Tensor, + target_cube_ids: torch.Tensor, + source_side_ids: torch.Tensor, + lower: torch.Tensor, + upper: torch.Tensor, +) -> torch.Tensor: + """Sample active-cube positions without overlapping an inactive cube.""" + cube_ids = torch.arange(CUBE_COUNT, device=base_x.device) + inactive_on_source = (cube_sides == source_side_ids.unsqueeze(1)) & (cube_ids != target_cube_ids.unsqueeze(1)) + sampled = lower + torch.rand_like(lower) * (upper - lower) + for _ in range(16): + conflicts = torch.any( + (torch.abs(sampled.unsqueeze(1) - base_x) < _MINIMUM_CUBE_SEPARATION) & inactive_on_source, dim=1 + ) + replacement = lower + torch.rand_like(lower) * (upper - lower) + sampled = torch.where(conflicts, replacement, sampled) + + conflicts = torch.any( + (torch.abs(sampled.unsqueeze(1) - base_x) < _MINIMUM_CUBE_SEPARATION) & inactive_on_source, dim=1 + ) + fallback = torch.full_like(sampled, TRANSFER_X) + return torch.where(conflicts, fallback, sampled) + + +class ConveyorResetStateTable(ManagerTermBase): + """Restore validated transfer states spanning released goal to moving start.""" + + def __init__(self, cfg: EventTermCfg, env: ManagerBasedRLEnv): + super().__init__(cfg, env) + rows = build_reset_rows() + self.recipe_ids = torch.tensor([row.recipe for row in rows], dtype=torch.long, device=env.device) + self.target_cube_ids = torch.tensor([row.target_cube_id for row in rows], dtype=torch.long, device=env.device) + self.variant_ids = torch.tensor([row.variant_id for row in rows], dtype=torch.long, device=env.device) + self.source_side_ids = torch.tensor([row.source_side_id for row in rows], dtype=torch.long, device=env.device) + self._arm_positions = torch.tensor([row.arm_positions for row in rows], dtype=torch.float32, device=env.device) + self._finger_positions = torch.tensor( + [row.finger_position for row in rows], dtype=torch.float32, device=env.device + ) + self.held_rows = torch.tensor([row.held for row in rows], dtype=torch.bool, device=env.device) + self._belt_range_fractions = torch.tensor( + [row.belt_range_fraction for row in rows], dtype=torch.float32, device=env.device + ) + self._robot: Articulation = env.scene["robot"] + self._cubes: tuple[RigidObject, ...] = tuple(env.scene[f"cube_{cube_id}"] for cube_id in range(CUBE_COUNT)) + self._arm_joint_ids = self._robot.find_joints("panda_joint[1-7]", preserve_order=True)[0] + self._finger_joint_ids = self._robot.find_joints("panda_finger_joint[1-2]", preserve_order=True)[0] + if len(self._arm_joint_ids) != 7 or len(self._finger_joint_ids) != 2: + raise ValueError("Conveyor transfer requires seven Panda arm joints and two finger joints.") + self.row_ids = torch.randint(self.row_count, (env.num_envs,), dtype=torch.long, device=env.device) + + @property + def row_count(self) -> int: + """Number of immutable physical reset rows.""" + return self.recipe_ids.numel() + + def _filtered_rows( + self, + fixed_recipe: int | None, + fixed_variant_id: int | None, + fixed_target_cube_id: int | None, + fixed_source_side_id: int | None, + ) -> torch.Tensor: + """Return rows matching optional deterministic evaluation controls.""" + mask = torch.ones(self.row_count, dtype=torch.bool, device=self.device) + if fixed_recipe is not None: + if not 0 <= fixed_recipe < len(ConveyorResetRecipe): + raise ValueError(f"fixed_recipe must lie in [0, {len(ConveyorResetRecipe) - 1}].") + mask &= self.recipe_ids == fixed_recipe + if fixed_variant_id is not None: + if fixed_variant_id < 0: + raise ValueError("fixed_variant_id must be non-negative.") + mask &= self.variant_ids == fixed_variant_id + if fixed_target_cube_id is not None: + if not 0 <= fixed_target_cube_id < CUBE_COUNT: + raise ValueError(f"fixed_target_cube_id must lie in [0, {CUBE_COUNT - 1}].") + mask &= self.target_cube_ids == fixed_target_cube_id + if fixed_source_side_id is not None: + if fixed_source_side_id not in (LEFT_SIDE, RIGHT_SIDE): + raise ValueError("fixed_source_side_id must be 0 (left) or 1 (right).") + mask &= self.source_side_ids == fixed_source_side_id + return torch.nonzero(mask, as_tuple=False).flatten() + + def __call__( + self, + env: ManagerBasedRLEnv, + env_ids: torch.Tensor, + fixed_recipe: int | None = None, + fixed_variant_id: int | None = None, + fixed_target_cube_id: int | None = None, + fixed_source_side_id: int | None = None, + belt_start_x_range: tuple[float, float] = (0.30, 0.82), + cube_position_noise: float = 0.015, + arm_joint_noise: float = 0.015, + ) -> None: + """Write sampled robot and four-cube states directly into simulation.""" + if env_ids is None or env_ids.numel() == 0: + return + if belt_start_x_range[0] >= belt_start_x_range[1]: + raise ValueError("belt_start_x_range must be strictly increasing.") + if cube_position_noise < 0.0 or arm_joint_noise < 0.0: + raise ValueError("Reset randomization ranges must be non-negative.") + + fixed_values = (fixed_recipe, fixed_variant_id, fixed_target_cube_id, fixed_source_side_id) + if any(value is not None for value in fixed_values): + candidates = self._filtered_rows(*fixed_values) + if candidates.numel() == 0: + raise RuntimeError("No conveyor reset rows match the fixed reset controls.") + row_ids = candidates[torch.randint(candidates.numel(), (env_ids.numel(),), device=self.device)] + self.row_ids[env_ids] = row_ids + else: + row_ids = self.row_ids[env_ids] + + recipes = self.recipe_ids[row_ids] + target_cube_ids = self.target_cube_ids[row_ids] + source_side_ids = self.source_side_ids[row_ids] + held_rows = self.held_rows[row_ids] + + arm_positions = self._arm_positions[row_ids].clone() + if arm_joint_noise > 0.0: + noise = (2.0 * torch.rand_like(arm_positions) - 1.0) * arm_joint_noise + # Preserve the exact approach, closure, and held-object manifold. + # Only deployment-like BELT starts receive robot randomization. + noise[recipes != int(ConveyorResetRecipe.BELT)] = 0.0 + arm_positions += noise + joint_positions = self._robot.data.default_joint_pos.torch[env_ids].clone() + joint_velocities = torch.zeros_like(joint_positions) + joint_positions[:, self._arm_joint_ids] = arm_positions + joint_positions[:, self._finger_joint_ids] = self._finger_positions[row_ids, None] + self._robot.actuators.target_command.set_position_index(value=joint_positions, env_ids=env_ids) + self._robot.actuators.target_command.set_velocity_index(value=joint_velocities, env_ids=env_ids) + self._robot.write_joint_position_to_sim_index(position=joint_positions, env_ids=env_ids) + self._robot.write_joint_velocity_to_sim_index(velocity=joint_velocities, env_ids=env_ids) + + count = env_ids.numel() + cube_slots, cube_sides, base_x, cube_y = _balanced_cube_slots(target_cube_ids, source_side_ids) + base_x = base_x.to(dtype=arm_positions.dtype) + cube_y = cube_y.to(dtype=arm_positions.dtype) + if cube_position_noise > 0.0: + base_x += (2.0 * torch.rand_like(base_x) - 1.0) * cube_position_noise + + belt_rows = recipes == int(ConveyorResetRecipe.BELT) + range_fraction = self._belt_range_fractions[row_ids] + active_lower = TRANSFER_X + range_fraction * (belt_start_x_range[0] - TRANSFER_X) + active_upper = TRANSFER_X + range_fraction * (belt_start_x_range[1] - TRANSFER_X) + active_x = _sample_collision_free_active_x( + base_x, + cube_sides, + target_cube_ids, + source_side_ids, + active_lower, + active_upper, + ) + + # On a full deployment reset, mirror the other cube on the source + # conveyor across the racetrack center. Opposite straight runs then + # differ by exactly half a lap even when the active start is sampled. + source_outer_slots = 2 * source_side_ids + 1 + source_outer_cube = cube_slots == source_outer_slots.unsqueeze(1) + mirrored_outer_x = (2.0 * BELT_CENTER_X - active_x).unsqueeze(1) + base_x = torch.where(belt_rows.unsqueeze(1) & source_outer_cube, mirrored_outer_x, base_x) + cube_positions = torch.stack( + (base_x, cube_y, torch.full_like(base_x, CUBE_REST_Z)), + dim=2, + ) + + active_positions = torch.stack( + (active_x, side_inner_y(source_side_ids), torch.full_like(active_x, CUBE_REST_Z)), + dim=1, + ) + goal_rows = recipes == int(ConveyorResetRecipe.GOAL) + active_positions[goal_rows, 1] = side_inner_y(1 - source_side_ids[goal_rows]) + if bool(torch.any(held_rows)): + active_positions[held_rows] = franka_tool_position(arm_positions[held_rows]) + cube_positions.scatter_( + 1, + target_cube_ids.view(-1, 1, 1).expand(-1, 1, 3), + active_positions.unsqueeze(1), + ) + + identity_quaternion = arm_positions.new_tensor((0.0, 0.0, 0.0, 1.0)).expand(count, -1) + for cube_id, cube in enumerate(self._cubes): + root_pose = cube.data.default_root_pose.torch[env_ids].clone() + root_pose[:, :3] = cube_positions[:, cube_id] + env.scene.env_origins[env_ids] + root_pose[:, 3:7] = identity_quaternion + root_velocity = torch.zeros((count, 6), dtype=root_pose.dtype, device=self.device) + cube.write_root_pose_to_sim_index(root_pose=root_pose, env_ids=env_ids) + cube.write_root_velocity_to_sim_index(root_velocity=root_velocity, env_ids=env_ids) + + +def select_next_transfer_cube( + cube_positions: torch.Tensor, + current_cube_ids: torch.Tensor, + source_side_ids: torch.Tensor, + transit_half_width: float = 0.14, +) -> torch.Tensor: + """Sample the next numbered cube already located on each source belt. + + A different eligible cube is sampled uniformly. The just-placed cube is + the fallback, so a valid goal remains available even when it is temporarily + the only parcel on that conveyor. + """ + if cube_positions.ndim != 3 or cube_positions.shape[1:] != (CUBE_COUNT, 3): + raise ValueError(f"cube_positions must have shape (N, {CUBE_COUNT}, 3).") + count = cube_positions.shape[0] + if current_cube_ids.shape != (count,) or source_side_ids.shape != (count,): + raise ValueError("Current cube and source-side ids must match the position batch.") + if transit_half_width <= 0.0: + raise ValueError("transit_half_width must be positive.") + + cube_ids = torch.arange(CUBE_COUNT, device=cube_positions.device).expand(count, -1) + on_left = cube_positions[:, :, 1] > transit_half_width + on_right = cube_positions[:, :, 1] < -transit_half_width + candidates = torch.where(source_side_ids.unsqueeze(1) == LEFT_SIDE, on_left, on_right) + alternatives = candidates & (cube_ids != current_cube_ids.unsqueeze(1)) + candidates = torch.where(torch.any(alternatives, dim=1, keepdim=True), alternatives, candidates) + + has_candidates = torch.any(candidates, dim=1) + fallback = torch.zeros_like(candidates) + fallback.scatter_(1, current_cube_ids.unsqueeze(1), True) + candidates = torch.where(has_candidates.unsqueeze(1), candidates, fallback) + return torch.multinomial(candidates.float(), 1).squeeze(1) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/rewards.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/rewards.py new file mode 100644 index 000000000000..17e70c37d94a --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/rewards.py @@ -0,0 +1,71 @@ +# 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 + +"""Sparse completion rewards and finite action penalties for conveyor transfer.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from isaaclab.managers import SceneEntityCfg + +if TYPE_CHECKING: + from isaaclab.assets import Articulation + from isaaclab.envs import ManagerBasedRLEnv + + +def transfer_success_reward( + env: ManagerBasedRLEnv, + command_name: str = "transfer", +) -> torch.Tensor: + """Return one on each stable transfer-completion transition.""" + command = env.command_manager.get_term(command_name) + command.evaluate() + return command.new_success.float() + + +def terminal_failure( + env: ManagerBasedRLEnv, + command_name: str = "transfer", +) -> torch.Tensor: + """Return one for non-timeout terminal failures.""" + command = env.command_manager.get_term(command_name) + command.evaluate() + return (env.reset_terminated & ~command.pending_success).float() + + +def action_term_l2(env: ManagerBasedRLEnv, action_name: str) -> torch.Tensor: + """Penalize one named raw action term.""" + action = env.action_manager.get_term(action_name).raw_actions + return torch.sum(torch.square(action), dim=1) + + +def finite_action_rate_l2( + env: ManagerBasedRLEnv, + action_names: tuple[str, ...] = ("arm_action", "gripper_action"), +) -> torch.Tensor: + """Penalize changes between the finite policy commands accepted by action terms.""" + if not action_names: + raise ValueError("At least one action term is required for the action-rate reward.") + terms = tuple(env.action_manager.get_term(name) for name in action_names) + action = torch.cat(tuple(term.raw_actions for term in terms), dim=1) + previous_action = torch.cat(tuple(term.previous_actions for term in terms), dim=1) + return torch.sum(torch.square(action - previous_action), dim=1) + + +def finite_joint_velocity_l2( + env: ManagerBasedRLEnv, + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot", joint_names=["panda_joint[1-7]"]), + maximum_velocity: float = 3.0, +) -> torch.Tensor: + """Penalize bounded arm velocity while sanitizing divergent states.""" + if maximum_velocity <= 0.0: + raise ValueError("maximum_velocity must be positive.") + robot: Articulation = env.scene[asset_cfg.name] + velocity = robot.data.joint_vel.torch[:, asset_cfg.joint_ids] + velocity = torch.nan_to_num(velocity, nan=0.0, posinf=maximum_velocity, neginf=-maximum_velocity) + return torch.sum(torch.square(torch.clamp(velocity, -maximum_velocity, maximum_velocity)), dim=1) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/terminations.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/terminations.py new file mode 100644 index 000000000000..eee38a695683 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/terminations.py @@ -0,0 +1,114 @@ +# 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 + +"""Actual episode termination and truncation terms for conveyor transfer.""" + +from __future__ import annotations + +import math +from collections.abc import Sequence +from typing import TYPE_CHECKING + +import torch + +from isaaclab.managers import SceneEntityCfg + +from ..conveyor_geometry import ( + BELT_CENTER_X, + BELT_HALF_STRAIGHT, + BELT_TURN_RADIUS, + BELT_WIDTH, + GUARD_THICKNESS, +) +from .cube_state import cube_values +from .reset_events import CUBE_SIZE + +if TYPE_CHECKING: + from isaaclab.assets import Articulation + from isaaclab.envs import ManagerBasedRLEnv + + +_TRACK_X_CLEARANCE = BELT_TURN_RADIUS + 0.5 * BELT_WIDTH + GUARD_THICKNESS + CUBE_SIZE + + +def invalid_action( + env: ManagerBasedRLEnv, + action_names: Sequence[str] = ("arm_action", "gripper_action"), +) -> torch.Tensor: + """Terminate environments whose latest policy action contained a non-finite value. + + Args: + env: Manager-based conveyor environment. + action_names: Action terms exposing an ``invalid_actions`` Boolean tensor. + + Returns: + Per-environment invalid-action mask. + """ + if not action_names: + raise ValueError("At least one action term is required for invalid-action termination.") + masks = tuple(env.action_manager.get_term(name).invalid_actions for name in action_names) + return torch.stack(masks, dim=0).any(dim=0) + + +def subgoal_time_out( + env: ManagerBasedRLEnv, + timeout_s: float = 20.0, + command_name: str = "transfer", +) -> torch.Tensor: + """Truncate environments that make no transfer within one subgoal timeout [s].""" + if timeout_s <= 0.0: + raise ValueError("timeout_s must be positive.") + command = env.command_manager.get_term(command_name) + command.evaluate() + timeout_steps = math.ceil(timeout_s / env.step_dt) + elapsed = env.episode_length_buf - command.subgoal_start_steps + return (elapsed >= timeout_steps) & ~command.pending_success + + +def transfer_sequence_time_out( + env: ManagerBasedRLEnv, + maximum_transfers: int = 8, + command_name: str = "transfer", +) -> torch.Tensor: + """Truncate completed sequences so reset-state coverage remains fresh.""" + if maximum_transfers < 1: + raise ValueError("maximum_transfers must be positive.") + command = env.command_manager.get_term(command_name) + command.evaluate() + return command.transfer_counts >= maximum_transfers + + +def cube_out_of_workspace( + env: ManagerBasedRLEnv, + minimum: tuple[float, float, float] = ( + BELT_CENTER_X - BELT_HALF_STRAIGHT - _TRACK_X_CLEARANCE, + -1.05, + -0.05, + ), + maximum: tuple[float, float, float] = ( + BELT_CENTER_X + BELT_HALF_STRAIGHT + _TRACK_X_CLEARANCE, + 1.05, + 0.80, + ), +) -> torch.Tensor: + """Terminate when any cube leaves the complete guarded racetrack workspace.""" + positions = cube_values(env, "root_pos_w", all_cubes=True) + positions -= env.scene.env_origins.unsqueeze(1) + lower = positions.new_tensor(minimum) + upper = positions.new_tensor(maximum) + return torch.any((positions < lower) | (positions > upper), dim=(1, 2)) + + +def nonfinite_scene_state( + env: ManagerBasedRLEnv, + robot_cfg: SceneEntityCfg = SceneEntityCfg("robot"), +) -> torch.Tensor: + """Terminate environments containing nonfinite robot or cube state.""" + robot: Articulation = env.scene[robot_cfg.name] + invalid = ~torch.all(torch.isfinite(robot.data.joint_pos.torch), dim=1) + invalid |= ~torch.all(torch.isfinite(robot.data.joint_vel.torch), dim=1) + for attribute in ("root_link_pose_w", "root_com_vel_w"): + invalid |= ~torch.isfinite(cube_values(env, attribute, all_cubes=True)).all(dim=(1, 2)) + return invalid diff --git a/source/isaaclab_tasks/pyproject.toml b/source/isaaclab_tasks/pyproject.toml index abaa241175d1..0c5b883dc231 100644 --- a/source/isaaclab_tasks/pyproject.toml +++ b/source/isaaclab_tasks/pyproject.toml @@ -20,11 +20,15 @@ requires-python = ">=3.12" dependencies = [ "isaaclab", "isaaclab_assets", + "isaaclab_contrib", + "isaaclab_newton", ] [tool.uv.sources] isaaclab = { path = "../isaaclab", editable = true } isaaclab_assets = { path = "../isaaclab_assets", editable = true } +isaaclab_contrib = { path = "../isaaclab_contrib", editable = true } +isaaclab_newton = { path = "../isaaclab_newton", editable = true } [project.urls] Homepage = "https://github.com/isaac-sim/IsaacLab" diff --git a/source/isaaclab_tasks/test/contrib/test_conveyor_franka_geometry.py b/source/isaaclab_tasks/test/contrib/test_conveyor_franka_geometry.py new file mode 100644 index 000000000000..79ecd533f33b --- /dev/null +++ b/source/isaaclab_tasks/test/contrib/test_conveyor_franka_geometry.py @@ -0,0 +1,57 @@ +# 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 + +"""Durable geometry checks for the contributed conveyor Franka task.""" + +from collections import Counter + +from isaaclab_tasks.contrib.conveyor_franka.conveyor_franka_env_cfg import _collision_properties, _cube +from isaaclab_tasks.contrib.conveyor_franka.conveyor_geometry import ( + MeshSpec, + belt_collision_geometry_specs, + belt_mesh_spec, + guard_mesh_specs, +) + + +def _edge_use_counts(spec: MeshSpec) -> Counter[tuple[int, int]]: + """Count triangle uses of every undirected mesh edge.""" + edges: Counter[tuple[int, int]] = Counter() + for triangle in spec.faces: + for start, end in zip(triangle, triangle[1:] + triangle[:1], strict=True): + edges[tuple(sorted((start, end)))] += 1 + return edges + + +def test_belt_top_faces_point_upward(): + """One-sided triangle-mesh surfaces support parcels from above.""" + for side in ("Left", "Right"): + specs = [belt_mesh_spec(side), *guard_mesh_specs(side)] + specs.extend(spec for spec in belt_collision_geometry_specs(side) if isinstance(spec, MeshSpec)) + for spec in specs: + assert set(_edge_use_counts(spec).values()) == {2} + _assert_top_faces_point_upward(spec) + + +def _assert_top_faces_point_upward(spec): + top_z = max(vertex[2] for vertex in spec.vertices) + for face in spec.faces: + vertices = tuple(spec.vertices[index] for index in face) + if all(vertex[2] == top_z for vertex in vertices): + a, b, c = vertices + cross_z = (b[0] - a[0]) * (c[1] - a[1]) - (b[1] - a[1]) * (c[0] - a[0]) + assert cross_z > 0.0 + + +def test_contact_configuration_uses_one_mujoco_parameterization(): + """Raw MuJoCo solref must not be combined with shadowed Newton force-space gains.""" + mujoco_cfg = _collision_properties()[-1] + cube_material = _cube("TestCube", (1.0, 0.0, 0.0), (0.0, 0.0, 0.0)).spawn.physics_material + + assert mujoco_cfg.solref is not None + assert cube_material.contact_stiffness is None + assert cube_material.contact_damping is None + assert cube_material.torsional_friction is None + assert cube_material.rolling_friction is None diff --git a/source/isaaclab_tasks/test/contrib/test_conveyor_franka_mdp.py b/source/isaaclab_tasks/test/contrib/test_conveyor_franka_mdp.py new file mode 100644 index 000000000000..8ad1fef8cb30 --- /dev/null +++ b/source/isaaclab_tasks/test/contrib/test_conveyor_franka_mdp.py @@ -0,0 +1,308 @@ +# 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 + +"""Racetrack policy contracts exercised through real CPU managers and assets.""" + +from contextlib import closing + +import gymnasium as gym +import pytest +import torch + +import isaaclab_tasks # noqa: F401 +from isaaclab_tasks.contrib.conveyor_franka import mdp +from isaaclab_tasks.contrib.conveyor_franka.agents.rsl_rl_ppo_cfg import ConveyorGaussianBernoulliDistribution +from isaaclab_tasks.contrib.conveyor_franka.conveyor_geometry import ( + BELT_CENTER_X, + BELT_INNER_STRAIGHT_Y, + BELT_OUTER_STRAIGHT_Y, +) +from isaaclab_tasks.contrib.conveyor_franka.mdp.commands import select_next_transfer_cube +from isaaclab_tasks.contrib.conveyor_franka.mdp.reset_events import ( + CUBE_COUNT, + ConveyorResetRecipe, + _balanced_cube_slots, + _sample_collision_free_active_x, + build_reset_rows, + franka_tool_position, +) +from isaaclab_tasks.utils.parse_cfg import parse_env_cfg + +_BASE = "IsaacContrib-Conveyor-Racetrack-Transfer-v0" + + +def _environment(task): + cfg = parse_env_cfg(task, device="cpu", num_envs=2) + cfg.sim.physics.use_cuda_graph = False + cfg.sim.visualizer_cfgs = [] + cfg.sim.use_fabric = False + cfg.sim.save_logs_to_file = False + cfg.seed = 41 + return gym.make(task, cfg=cfg).unwrapped + + +def _parcels(env): + return tuple(env.scene[f"cube_{i}"] for i in range(env.cfg.conveyor_force.transported_body_count_per_env)) + + +def _waiting_positions(env): + positions = torch.tensor([[[2.5, 0.8, 0.5]] * len(_parcels(env))] * 2) + ids = torch.arange(positions.shape[1]) + positions[:, :, 0] = 2.4 + 0.06 * (ids % 6) + positions[:, :, 1] += 0.06 * (ids // 6) + return positions + + +def _place(env, positions): + for i, parcel in enumerate(_parcels(env)): + pose = parcel.data.default_root_pose.torch.clone() + pose[:, :3] = positions[:, i] + env.scene.env_origins + parcel.write_root_pose_to_sim_index(root_pose=pose, skip_forward=True) + parcel.write_root_velocity_to_sim_index(root_velocity=torch.zeros(2, 6), skip_forward=True) + env.sim.forward() + env.scene.update(env.step_dt) + + +def test_checkpoint_interface_and_invalid_action_recovery(): + """The real racetrack preserves learned feature order, finite rewards and independent resets.""" + with closing(_environment(_BASE)) as env: + obs, _ = env.reset() + assert obs["policy"].shape == (2, 123) + assert env.action_manager.total_action_dim == 8 + assert env.physics_dt == 1 / 120 and env.step_dt == 1 / 60 + result = env.step(torch.zeros(2, 8)) + assert torch.isfinite(result[0]["policy"]).all() and torch.isfinite(result[1]).all() + positions = _waiting_positions(env) + positions[:, :4] = torch.tensor([[0.4, 0.27, 0.06], [0.6, -0.27, 0.06], [0.8, 0, 0.2], [0.9, -0.27, 0.06]]) + positions[1, 2] = torch.tensor([0.8, -0.27, 0.2]) + _place(env, positions) + command = env.command_manager.get_term("transfer") + env.episode_length_buf[:] = torch.tensor([11, 19]) + command.set_goal(2) + policy = env.observation_manager.compute()["policy"] + expected = positions[:, :4] + torch.testing.assert_close(policy[:, 16:28], expected.flatten(1)) + torch.testing.assert_close(policy[:, 76:79], positions[:, 2]) + torch.testing.assert_close(policy[:, 85:89], torch.tensor([[0, 0, 1, 0]] * 2).float()) + torch.testing.assert_close(policy[:, 101:103], torch.tensor([[0, 1], [1, 0]]).float()) + torch.testing.assert_close( + policy[:, 89:101], + torch.tensor([[1, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 1], [1, 0, 0, 0, 0, 1, 0, 0, 1, 0, 0, 1]]).float(), + ) + assert command.held_cube_ids.tolist() == [-1, -1] + assert command.subgoal_start_steps.tolist() == [11, 19] + command.subgoal_start_steps[:] = env.episode_length_buf - torch.tensor([1199, 1200]) + assert mdp.subgoal_time_out(env, timeout_s=20).tolist() == [False, True] + command.pending_success[1] = True + assert not mdp.subgoal_time_out(env, timeout_s=20).any() + command.transfer_counts[:] = torch.tensor([7, 8]) + assert mdp.transfer_sequence_time_out(env, maximum_transfers=8).tolist() == [False, True] + command.pending_success.zero_() + arm = env.action_manager.get_term("arm_action") + gripper = env.action_manager.get_term("gripper_action") + previous = torch.tensor([[0.1] * 7 + [-1.0], [-0.2] * 7 + [1.0]]) + current = torch.tensor([[-0.3] * 7 + [1.0], [0.4] * 7 + [-1.0]]) + env.action_manager.process_action(previous) + env.action_manager.process_action(current) + torch.testing.assert_close(mdp.finite_action_rate_l2(env), (current - previous).square().sum(1)) + bad = torch.tensor( + [[float("nan"), float("inf"), -float("inf"), 5, -5, 0.5, -0.5, float("nan")], [0.5] * 7 + [float("inf")]] + ) + robot = env.scene["robot"] + joint_ids, _ = robot.find_joints("panda_joint[1-7]", preserve_order=True) + joints = robot.data.joint_pos.torch.clone() + initial = torch.tensor([0.05, 0.2, -0.1, -2.25, 0, 2.45, 0.775]).repeat(2, 1) + joints[:, joint_ids] = initial + robot.write_joint_position_to_sim_index(position=joints) + env.sim.forward() + env.scene.update(env.step_dt) + if env.spec.id == _BASE: + # Differentiate the independent reset FK rather than repeating the velocity shift. + rates = torch.tensor([0.5, -0.4, 0.3, -0.2, 0.4, -0.5, 0.2]).repeat(2, 1) + velocity = torch.zeros_like(robot.data.joint_vel.torch) + velocity[:, joint_ids] = rates + robot.write_joint_velocity_to_sim_index(velocity=velocity) + env.sim.forward() + env.scene.update(env.step_dt) + expected_velocity = ( + franka_tool_position(initial + 1e-3 * rates) - franka_tool_position(initial - 1e-3 * rates) + ) / 2e-3 + torch.testing.assert_close(mdp.end_effector_velocity(env)[:, :3], expected_velocity, atol=1e-4, rtol=1e-3) + env.action_manager.process_action(bad) + torch.testing.assert_close( + arm.processed_actions, initial + torch.tensor([[0, 0.12, -0.12, 0.12, -0.12, 0.06, -0.06], [0.06] * 7]) + ) + torch.testing.assert_close(arm.raw_actions, torch.tensor([[0, 1, -1, 1, -1, 0.5, -0.5], [0.5] * 7])) + torch.testing.assert_close(gripper.raw_actions, torch.tensor([[-1.0], [1.0]])) + torch.testing.assert_close(gripper.processed_actions, torch.tensor([[0.0, 0.0], [0.04, 0.04]])) + assert torch.isfinite(arm.processed_actions).all() and torch.isfinite(gripper.processed_actions).all() + assert torch.isfinite(mdp.finite_action_rate_l2(env)).all() + bad[0, 7] = -1 + env.action_manager.process_action(bad) + assert mdp.invalid_action(env).tolist() == [True, True] + env.action_manager.reset(env_ids=[0]) + assert mdp.invalid_action(env).tolist() == [False, True] + env.action_manager.reset() + assert not mdp.invalid_action(env).any() + terminated = env.step(bad) + assert terminated[2].tolist() == [True, True] + assert torch.isfinite(terminated[0]["policy"]).all() and torch.isfinite(terminated[1]).all() + cfg = env.cfg.copy() + cfg.actions.arm_action.preserve_order = False + with pytest.raises(ValueError, match="preserve"): + cfg.validate() + + +def test_success_guided_sampling_records_completed_selected_episodes(): + """Real resets credit the previous row once and practice intermediate-success starts more often.""" + with closing(_environment(_BASE)) as env: + env.reset() + command = env.command_manager.get_term("transfer") + curriculum = env.curriculum_manager.cfg.reset_sampling.func + monitor = curriculum.success_monitor + table = env.event_manager.get_term_cfg("reset_from_state_table").func + # Repeated startup resets have no episode to credit. + env.reset() + assert monitor.success_size.sum() == 0 + for succeeded in (True, False): + table.row_ids[:] = torch.tensor([0, 1]) + env.episode_length_buf[:] = 5 + command.ever_success[:] = torch.tensor([succeeded, True]) + untouched_row = table.row_ids[1].clone() + env.reset(env_ids=torch.tensor([0])) + assert table.row_ids[1] == untouched_row + assert monitor.success_size[1] == 0 + assert monitor.success_size.sum() == monitor.success_size[0] == 2 + assert monitor.success_rate[0] == 0.5 + probabilities = curriculum.sampler.probabilities() + assert torch.isfinite(probabilities).all() and (probabilities > 0).all() + torch.testing.assert_close(probabilities.sum(), torch.tensor(1.0)) + assert probabilities[0] > probabilities[1] + # A genuine unsuccessful episode is counted; a zero-step manual reset is not. + table.row_ids[:] = torch.tensor([0, 1]) + command.ever_success[1] = False + env.reset(env_ids=torch.tensor([1])) + assert monitor.success_size[1] == 1 and monitor.success_rate[1] == 0 + env.reset(env_ids=torch.tensor([1])) + assert monitor.success_size.sum() == 3 + + +def test_reset_bank_is_complete_and_physically_calibrated(): + """Every command/phase is represented once; tool anchors and held metadata agree with physical poses.""" + rows = build_reset_rows() + counts = (1, 4, 6, 4, 5, 5, 7) + keys = {(r.recipe, r.variant_id, r.target_cube_id, r.source_side_id) for r in rows} + assert len(keys) == len(rows) == sum(counts) * 8 + assert all( + recipe in ConveyorResetRecipe and 0 <= variant < counts[recipe] and 0 <= cube < 4 and side in (0, 1) + for recipe, variant, cube, side in keys + ) + for row in rows: + held = row.recipe in (ConveyorResetRecipe.LIFT, ConveyorResetRecipe.CARRY, ConveyorResetRecipe.PLACE) + held |= row.recipe == ConveyorResetRecipe.GRASP and row.variant_id == counts[ConveyorResetRecipe.GRASP] - 1 + assert row.held == held + closure = [ + r.finger_position + for r in rows + if r.recipe == ConveyorResetRecipe.GRASP and r.target_cube_id == 0 and r.source_side_id == 0 + ] + assert all(a > b for a, b in zip(closure, closure[1:])) + # Independently calibrated approach, lift, transit and release tool positions. + anchors = { + ConveyorResetRecipe.GOAL: (0, 0.14), + ConveyorResetRecipe.PLACE: (3, 0.105), + ConveyorResetRecipe.CARRY: (2, 0.25), + ConveyorResetRecipe.LIFT: (3, 0.22), + ConveyorResetRecipe.GRASP: (0, 0.06), + ConveyorResetRecipe.PREGRASP: (0, 0.14), + } + for row in rows: + if row.target_cube_id or row.recipe not in anchors or row.variant_id != anchors[row.recipe][0]: + continue + y = 0.27 * (1 - 2 * row.source_side_id) + if row.recipe in (ConveyorResetRecipe.GOAL, ConveyorResetRecipe.PLACE): + y = -y + elif row.recipe == ConveyorResetRecipe.CARRY: + y = 0 + actual = franka_tool_position(torch.tensor([row.arm_positions], dtype=torch.float64))[0] + torch.testing.assert_close( + actual, torch.tensor([0.52, y, anchors[row.recipe][1]], dtype=actual.dtype), atol=3e-4, rtol=0 + ) + + +def test_binary_policy_likelihoods(): + """The mixed distribution samples exact binary gripper actions with finite likelihoods and KL.""" + distribution = ConveyorGaussianBernoulliDistribution(output_dim=8) + distribution.update(torch.zeros(4096, 8)) + samples, previous = distribution.sample(), tuple(p.clone() for p in distribution.params) + distribution.update(torch.full((4096, 8), 0.2)) + divergence = distribution.kl_divergence(previous, distribution.params) + assert set(samples[:, -1].unique().tolist()) == {-1, 1} + assert torch.isfinite(distribution.log_prob(samples)).all() + assert torch.isfinite(divergence).all() and (divergence >= 0).all() + + +def test_next_transfer_cube_is_random_among_eligible_alternatives(): + """Continuing commands use the one-hot target instead of a cyclic identity shortcut.""" + count = 4096 + positions = torch.zeros((count, CUBE_COUNT, 3)) + positions[:, :, 1] = torch.tensor((-0.27, -0.27, -0.27, 0.27)) + current_cube_ids = torch.zeros(count, dtype=torch.long) + source_side_ids = torch.ones(count, dtype=torch.long) + torch.manual_seed(7) + + selected = select_next_transfer_cube(positions, current_cube_ids, source_side_ids) + + assert set(selected.tolist()) == {1, 2} + frequencies = torch.bincount(selected, minlength=CUBE_COUNT).float() / count + assert torch.all(torch.abs(frequencies[1:3] - 0.5) < 0.05) + + +def test_active_cube_sampling_avoids_inactive_source_lane_cubes(): + """Random deployment starts cannot begin with interpenetrating parcels.""" + count = 2048 + base_x = torch.tensor((0.26, 0.42, 0.72, 0.88)).expand(count, -1).clone() + cube_sides = torch.tensor((0, 0, 1, 1)).expand(count, -1).clone() + target_cube_ids = torch.arange(count) % CUBE_COUNT + source_side_ids = torch.arange(count) % 2 + cube_sides.scatter_(1, target_cube_ids.unsqueeze(1), source_side_ids.unsqueeze(1)) + + sampled = _sample_collision_free_active_x( + base_x, + cube_sides, + target_cube_ids, + source_side_ids, + torch.full((count,), 0.30), + torch.full((count,), 0.82), + ) + + cube_ids = torch.arange(CUBE_COUNT).expand(count, -1) + inactive_on_source = (cube_sides == source_side_ids.unsqueeze(1)) & (cube_ids != target_cube_ids.unsqueeze(1)) + separation = torch.abs(sampled.unsqueeze(1) - base_x) + assert torch.all(separation[inactive_on_source] >= 0.055) + + +def test_deployment_layout_uses_each_racetrack_straight_run_once(): + """Every command keeps its target reachable and distributes all four cubes evenly.""" + target_cube_ids = torch.arange(CUBE_COUNT).repeat_interleave(2) + source_side_ids = torch.arange(2).repeat(CUBE_COUNT) + + slots, cube_sides, cube_x, cube_y = _balanced_cube_slots(target_cube_ids, source_side_ids) + + expected_slots = torch.arange(CUBE_COUNT).expand(target_cube_ids.numel(), -1) + torch.testing.assert_close(torch.sort(slots, dim=1).values, expected_slots) + active_slots = slots.gather(1, target_cube_ids.unsqueeze(1)).squeeze(1) + torch.testing.assert_close(active_slots, 2 * source_side_ids) + assert torch.all(torch.sum(cube_sides == 0, dim=1) == 2) + assert torch.all(torch.sum(cube_sides == 1, dim=1) == 2) + + expected_y = torch.tensor( + (-BELT_OUTER_STRAIGHT_Y, -BELT_INNER_STRAIGHT_Y, BELT_INNER_STRAIGHT_Y, BELT_OUTER_STRAIGHT_Y) + ) + torch.testing.assert_close(torch.sort(cube_y, dim=1).values, expected_y.expand_as(cube_y)) + for side_id in (0, 1): + side_x = torch.where(cube_sides == side_id, cube_x, torch.nan) + expected_sum = torch.full((source_side_ids.numel(),), 2 * BELT_CENTER_X, dtype=cube_x.dtype) + torch.testing.assert_close(torch.nansum(side_x, dim=1), expected_sum) diff --git a/uv.lock b/uv.lock index 0d48f7d187cf..cb84f120eeb4 100644 --- a/uv.lock +++ b/uv.lock @@ -2224,12 +2224,16 @@ source = { editable = "source/isaaclab_tasks" } dependencies = [ { name = "isaaclab" }, { name = "isaaclab-assets" }, + { name = "isaaclab-contrib" }, + { name = "isaaclab-newton" }, ] [package.metadata] requires-dist = [ { name = "isaaclab", editable = "source/isaaclab" }, { name = "isaaclab-assets", editable = "source/isaaclab_assets" }, + { name = "isaaclab-contrib", editable = "source/isaaclab_contrib" }, + { name = "isaaclab-newton", editable = "source/isaaclab_newton" }, ] [[package]]