diff --git a/.gitattributes b/.gitattributes index 388f98ea091f..8c45aa4d5b70 100644 --- a/.gitattributes +++ b/.gitattributes @@ -1,5 +1,7 @@ *.usd filter=lfs diff=lfs merge=lfs -text *.usda filter=lfs diff=lfs merge=lfs -text +source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/conveyor_*_supported.usda -filter diff merge text +source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/parcel*.usda -filter diff merge text *.psd filter=lfs diff=lfs merge=lfs -text *.hdr filter=lfs diff=lfs merge=lfs -text *.dae filter=lfs diff=lfs merge=lfs -text diff --git a/.gitignore b/.gitignore index 4a74400b8bef..483c436e1cac 100644 --- a/.gitignore +++ b/.gitignore @@ -14,6 +14,7 @@ # No USD files allowed in the repo **/*.usd **/*.usda +!source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/*.usda **/*.usdc **/*.usdz diff --git a/docs/source/_static/css/environment-browser.js b/docs/source/_static/css/environment-browser.js index edb297ec0bf1..cd0b70e3517b 100644 --- a/docs/source/_static/css/environment-browser.js +++ b/docs/source/_static/css/environment-browser.js @@ -61,6 +61,8 @@ ["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-Conveyor-Warehouse-Sorting-v0", "rsl_rl", "", "", "", "./tasks/conveyor/warehouse_sorting.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/docs/source/_static/tasks/conveyor/warehouse_sorting.jpg b/docs/source/_static/tasks/conveyor/warehouse_sorting.jpg new file mode 100644 index 000000000000..0345c19ea368 Binary files /dev/null and b/docs/source/_static/tasks/conveyor/warehouse_sorting.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/conveyor-warehouse-reset-frame.fixed.rst b/source/isaaclab_tasks/changelog.d/conveyor-warehouse-reset-frame.fixed.rst new file mode 100644 index 000000000000..f810496457cd --- /dev/null +++ b/source/isaaclab_tasks/changelog.d/conveyor-warehouse-reset-frame.fixed.rst @@ -0,0 +1,2 @@ +* Fixed warehouse robot resets by keeping physics assets at their clone origins and + translating the presentation instead of applying the workspace height twice. 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/changelog.d/maximiliank-conveyor-warehouse.added.rst b/source/isaaclab_tasks/changelog.d/maximiliank-conveyor-warehouse.added.rst new file mode 100644 index 000000000000..0e1f45b16423 --- /dev/null +++ b/source/isaaclab_tasks/changelog.d/maximiliank-conveyor-warehouse.added.rst @@ -0,0 +1 @@ +* Added warehouse parcel sorting with four dynamically assigned policy slots, colored batches, circulating conveyors, and referenced USD presentation assets. Background parcel motion follows simulation steps in both interactive Kit playback and headless recordings. diff --git a/source/isaaclab_tasks/changelog.d/maximiliank-conveyor-warehouse.minor b/source/isaaclab_tasks/changelog.d/maximiliank-conveyor-warehouse.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..978c7bdb4873 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/README.md @@ -0,0 +1,148 @@ +# Conveyor tasks + +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. + +## Warehouse sorting task + +The workcell contains 24 physical 40 mm, 50 g parcels: six each in blue, orange, green, +and purple. Blue and green belong on the positive-Y conveyor; orange and purple belong +on the negative-Y conveyor. Resets shuffle a mixed batch across the supply feeds using +seeded randomness. Slow feeds deliver parcels through 12 cm gravity drops; misplaced +arrivals recirculate until transferred. Colors and physical identities remain fixed. + +The two manipulation straights and adjoining 90-degree bends preserve the trained +positions, dimensions, and 0.35 m/s speed. Compact elevated returns connect them to +asymmetric routes. Incline normals follow their panels, with 0.95 traction friction; +the original manipulation sections retain their 0.5 setting. + +Four policy slots map to the larger physical pool. Local parcels and the active grasp +keep their assignments; the command manager assigns misplaced arrivals to available +slots before observations are computed. Reassignment changes identity mapping, never +physical poses. Commands, rewards, placement checks, and safety checks use the same +mapping. All 24 parcels receive physical belt forces. + +Remote inventory is represented through canonical waiting-slot observations because +the checkpoint learned compact, flat returns. Local states and physical interactions +remain exact. Class dispatch uses supervisory metadata; the state-based policy does +not recognize color. Playback parks the arm between transfers and after sorting. +Invalid actions still reach sanitization and termination. The shared racetrack +checkpoint has not demonstrated successful warehouse sorting: a two-environment, +2,400-step check completed no transfers despite active arrivals, with no safety +resets or non-finite observations. Evaluate warehouse training with physical +transfer and complete-batch sorting metrics. + +`commands.transfer.parcel_colors` sets appearances and `parcel_destinations` sets +conveyor IDs; each color must have one destination. Set `randomize_arrivals=False` +for authored arrival ordering. `sorted_parcels` counts settled, correctly routed parcels; +`batch_complete` indicates that all parcels have reached their assigned loops. + +Use Kit/RTX for the USD materials, warehouse lighting, and twelve animated background +cartons. Background inventory is render-only and adds no contacts or policy observations. + +```bash +DISPLAY=:1 uv run --extra isaacsim isaaclab play --rl_library rsl_rl \ + --task IsaacContrib-Conveyor-Warehouse-Sorting-v0 \ + --checkpoint pretrained --num_envs 1 --device cuda:0 --viz kit --real-time \ + --kit_args=--/UJITSO/geometry=false +``` + +Kit renders USD directly; `sim.physics.load_visual_shapes=False` excludes dressing +from the Newton model. For a static approximation with `--viz newton_gl`, pass +`env.sim.physics.load_visual_shapes=True`. The Kit override disables experimental +geometry streaming to prevent disappearing meshes with Fabric transforms. Assets +and textures are downloaded and cached on first use. Presentation defaults to one environment. + +### Fine-tuning + +Warehouse training starts at home with parcels on the actual feeds, disables the +racetrack reset curriculum, and passes every sampled policy action through. +Idle parking applies only during playback. Warm-start with the shared checkpoint: + +```bash +uv run isaaclab train --rl_library rsl_rl \ + --task IsaacContrib-Conveyor-Warehouse-Sorting-v0 \ + --checkpoint /path/to/model.pt --num_envs 256 --device cuda:0 +``` + +### Editing the presentation + +The authored USDA layers are versioned with the task because their route metadata, +collision paths, and animation must match its implementation. NVIDIA meshes and +textures stay at their original remote locations under their original licenses. + +- `assets/warehouse.usda` owns layout, props, materials, lights, cameras, and looping + background transforms. `Cameras/Workcell` and `Cameras/Overview` provide two views. +- `assets/conveyor_routes.usda` owns route centerlines, frame meshes, traction, + feed speeds, and `conveyor:parcelSpawnPositions`. Python builds collision proxies + from these paths after application startup; task discovery does not import USD. +- `assets/parcel.usda` normalizes the referenced SimReady carton to the original + centered 40 mm collider; `parcel_{blue,orange,green,purple}.usda` add color bands. +- `assets/conveyor_*_supported.usda` reference source conveyors and describe support + extension with two scalar attributes. The spawner derives lower-frame points and normals. +- `env.conveyor_cube_pool.slot_ids` maps policy slots to physical IDs; + `transfer_counts` records placements. Selected resets clear that batch's counters. + +The spawner caches external assets, removes imported physics and action graphs, +and samples background animation using simulation time. See the +[SimReady Explorer documentation](https://docs.omniverse.nvidia.com/extensions/latest/ext_core/ext_browser-extensions/simready-explorer.html) +for catalog access. Layout references include +[Dematic modular conveyors](https://www.dematic.com/content/dam/dematic/downloads/brochures/NA_BR_1039_MCS.pdf), +[Dematic sortation](https://www.dematic.com/content/dam/dematic/downloads/whitepapers/NA_WP-1015_Sorting-Out-Sortation.pdf), +and [KION warehouse installations](https://www.kiongroup.com/en/News-Stories/Stories/Innovation/Warehousing-is-being-transformed.html). 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..6f25470f4768 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/__init__.py @@ -0,0 +1,30 @@ +# 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 and warehouse sorting tasks sharing a Franka policy.""" + +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", + }, +) + +gym.register( + id="IsaacContrib-Conveyor-Warehouse-Sorting-v0", + entry_point=f"{__name__}.conveyor_franka_warehouse_env:ConveyorFrankaWarehouseEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.conveyor_franka_asset_env_cfg:ConveyorFrankaA09A12EnvCfg", + "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/assets/conveyor_quarter_supported.usda b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/conveyor_quarter_supported.usda new file mode 100644 index 000000000000..ac9c4bd584a2 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/conveyor_quarter_supported.usda @@ -0,0 +1,25 @@ +#usda 1.0 +# 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 +( + defaultPrim = "Conveyor" +) + +def Xform "Conveyor" ( + prepend references = @https://omniverse-content-production.s3-us-west-2.amazonaws.com/Assets/Isaac/6.1/Isaac/Props/Conveyors/ConveyorBelt_A03.usd@ +) +{ + over "SM_ConveyorBelt_A03" + { + custom double conveyor:supportAnchor = 24.452552917378103 + custom double conveyor:supportScale = 3.448973206624863 + } + + over "SM_ConveyorBelt_A03_Decal" + { + custom double conveyor:supportAnchor = 140.0000003367605 + custom double conveyor:supportScale = 3.448973193849789 + } +} diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/conveyor_routes.usda b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/conveyor_routes.usda new file mode 100644 index 000000000000..c35390b0218f --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/conveyor_routes.usda @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5ffa5fa3f772265e0d6e563b5ffdf3ea315602dd45ce4a2205a30477aee6e572 +size 1045928 diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/conveyor_straight_supported.usda b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/conveyor_straight_supported.usda new file mode 100644 index 000000000000..4f06ee40ba80 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/conveyor_straight_supported.usda @@ -0,0 +1,25 @@ +#usda 1.0 +# 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 +( + defaultPrim = "Conveyor" +) + +def Xform "Conveyor" ( + prepend references = @https://omniverse-content-production.s3-us-west-2.amazonaws.com/Assets/Isaac/6.1/Isaac/Props/Conveyors/ConveyorBelt_A09.usd@ +) +{ + over "SM_ConveyorBelt_A09_02" + { + custom double conveyor:supportAnchor = 24.45254509724512 + custom double conveyor:supportScale = 3.448973216496531 + } + + over "SM_ConveyorBelt_A09_Decal_02" + { + custom double conveyor:supportAnchor = -28.255675946768882 + custom double conveyor:supportScale = 3.448973111935675 + } +} diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/parcel.usda b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/parcel.usda new file mode 100644 index 000000000000..cdd38bc17061 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/parcel.usda @@ -0,0 +1,23 @@ +#usda 1.0 +# 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 +( + defaultPrim = "Parcel" + doc = "Conveyor Franka presentation. NVIDIA assets remain external references; this layer owns only layout and appearance. SPDX-License-Identifier: BSD-3-Clause" + metersPerUnit = 1 + upAxis = "Z" +) + +def Xform "Parcel" +{ + def Xform "Cardboard" ( + prepend references = @https://omniverse-content-production.s3-us-west-2.amazonaws.com/Assets/simready_content/common_assets/props/cardbox_a1/cardbox_a1.usd@ + ) + { + float3 xformOp:scale = (0.057189006, 0.07675307, 0.07844259) + double3 xformOp:translate = (-0, 8.463453207854634e-8, -0.02) + uniform token[] xformOpOrder = ["xformOp:translate", "xformOp:scale"] + } +} diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/parcel_blue.usda b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/parcel_blue.usda new file mode 100644 index 000000000000..684b6cbe54aa --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/parcel_blue.usda @@ -0,0 +1,60 @@ +#usda 1.0 +# 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 +( + defaultPrim = "Parcel" + metersPerUnit = 1 + upAxis = "Z" +) + +def Xform "Parcel" ( + prepend references = @parcel.usda@ +) +{ + over "Cardboard" + { + over "cardbox_a1_inst" + { + over "Looks" + { + over "opaque__cardboard__cardbox" + { + over "Shader" + { + custom color3f inputs:diffuse_tint = (0.22, 0.5, 1) + } + } + } + } + } + + def "Looks" + { + def Material "ClassBand" + { + token outputs:surface.connect = + + def Shader "Shader" + { + uniform token info:id = "UsdPreviewSurface" + color3f inputs:diffuseColor = (0.035, 0.25, 0.85) + float inputs:roughness = 0.58 + token outputs:surface + } + } + } + + def Mesh "ClassBand" ( + prepend apiSchemas = ["MaterialBindingAPI"] + ) + { + uniform bool doubleSided = 1 + int[] faceVertexCounts = [4, 4, 4, 4] + int[] faceVertexIndices = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15] + rel material:binding = + point3f[] points = [(-0.009, -0.02, 0.02), (0.009, -0.02, 0.02), (0.009, 0.02, 0.02), (-0.009, 0.02, 0.02), (-0.009, -0.02, -0.02), (-0.009, 0.02, -0.02), (0.009, 0.02, -0.02), (0.009, -0.02, -0.02), (-0.009, -0.02, -0.02), (0.009, -0.02, -0.02), (0.009, -0.02, 0.02), (-0.009, -0.02, 0.02), (-0.009, 0.02, -0.02), (-0.009, 0.02, 0.02), (0.009, 0.02, 0.02), (0.009, 0.02, -0.02)] + uniform token subdivisionScheme = "none" + } +} diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/parcel_green.usda b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/parcel_green.usda new file mode 100644 index 000000000000..2360ff252987 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/parcel_green.usda @@ -0,0 +1,60 @@ +#usda 1.0 +# 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 +( + defaultPrim = "Parcel" + metersPerUnit = 1 + upAxis = "Z" +) + +def Xform "Parcel" ( + prepend references = @parcel.usda@ +) +{ + over "Cardboard" + { + over "cardbox_a1_inst" + { + over "Looks" + { + over "opaque__cardboard__cardbox" + { + over "Shader" + { + custom color3f inputs:diffuse_tint = (0.32, 0.9, 0.36) + } + } + } + } + } + + def "Looks" + { + def Material "ClassBand" + { + token outputs:surface.connect = + + def Shader "Shader" + { + uniform token info:id = "UsdPreviewSurface" + color3f inputs:diffuseColor = (0.04, 0.65, 0.18) + float inputs:roughness = 0.58 + token outputs:surface + } + } + } + + def Mesh "ClassBand" ( + prepend apiSchemas = ["MaterialBindingAPI"] + ) + { + uniform bool doubleSided = 1 + int[] faceVertexCounts = [4, 4, 4, 4] + int[] faceVertexIndices = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15] + rel material:binding = + point3f[] points = [(-0.009, -0.02, 0.02), (0.009, -0.02, 0.02), (0.009, 0.02, 0.02), (-0.009, 0.02, 0.02), (-0.009, -0.02, -0.02), (-0.009, 0.02, -0.02), (0.009, 0.02, -0.02), (0.009, -0.02, -0.02), (-0.009, -0.02, -0.02), (0.009, -0.02, -0.02), (0.009, -0.02, 0.02), (-0.009, -0.02, 0.02), (-0.009, 0.02, -0.02), (-0.009, 0.02, 0.02), (0.009, 0.02, 0.02), (0.009, 0.02, -0.02)] + uniform token subdivisionScheme = "none" + } +} diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/parcel_orange.usda b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/parcel_orange.usda new file mode 100644 index 000000000000..8bb032846167 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/parcel_orange.usda @@ -0,0 +1,60 @@ +#usda 1.0 +# 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 +( + defaultPrim = "Parcel" + metersPerUnit = 1 + upAxis = "Z" +) + +def Xform "Parcel" ( + prepend references = @parcel.usda@ +) +{ + over "Cardboard" + { + over "cardbox_a1_inst" + { + over "Looks" + { + over "opaque__cardboard__cardbox" + { + over "Shader" + { + custom color3f inputs:diffuse_tint = (1, 0.55, 0.22) + } + } + } + } + } + + def "Looks" + { + def Material "ClassBand" + { + token outputs:surface.connect = + + def Shader "Shader" + { + uniform token info:id = "UsdPreviewSurface" + color3f inputs:diffuseColor = (0.95, 0.24, 0.025) + float inputs:roughness = 0.58 + token outputs:surface + } + } + } + + def Mesh "ClassBand" ( + prepend apiSchemas = ["MaterialBindingAPI"] + ) + { + uniform bool doubleSided = 1 + int[] faceVertexCounts = [4, 4, 4, 4] + int[] faceVertexIndices = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15] + rel material:binding = + point3f[] points = [(-0.009, -0.02, 0.02), (0.009, -0.02, 0.02), (0.009, 0.02, 0.02), (-0.009, 0.02, 0.02), (-0.009, -0.02, -0.02), (-0.009, 0.02, -0.02), (0.009, 0.02, -0.02), (0.009, -0.02, -0.02), (-0.009, -0.02, -0.02), (0.009, -0.02, -0.02), (0.009, -0.02, 0.02), (-0.009, -0.02, 0.02), (-0.009, 0.02, -0.02), (-0.009, 0.02, 0.02), (0.009, 0.02, 0.02), (0.009, 0.02, -0.02)] + uniform token subdivisionScheme = "none" + } +} diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/parcel_purple.usda b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/parcel_purple.usda new file mode 100644 index 000000000000..b57791c3d559 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/parcel_purple.usda @@ -0,0 +1,60 @@ +#usda 1.0 +# 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 +( + defaultPrim = "Parcel" + metersPerUnit = 1 + upAxis = "Z" +) + +def Xform "Parcel" ( + prepend references = @parcel.usda@ +) +{ + over "Cardboard" + { + over "cardbox_a1_inst" + { + over "Looks" + { + over "opaque__cardboard__cardbox" + { + over "Shader" + { + custom color3f inputs:diffuse_tint = (0.8, 0.27, 1) + } + } + } + } + } + + def "Looks" + { + def Material "ClassBand" + { + token outputs:surface.connect = + + def Shader "Shader" + { + uniform token info:id = "UsdPreviewSurface" + color3f inputs:diffuseColor = (0.5, 0.065, 0.8) + float inputs:roughness = 0.58 + token outputs:surface + } + } + } + + def Mesh "ClassBand" ( + prepend apiSchemas = ["MaterialBindingAPI"] + ) + { + uniform bool doubleSided = 1 + int[] faceVertexCounts = [4, 4, 4, 4] + int[] faceVertexIndices = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15] + rel material:binding = + point3f[] points = [(-0.009, -0.02, 0.02), (0.009, -0.02, 0.02), (0.009, 0.02, 0.02), (-0.009, 0.02, 0.02), (-0.009, -0.02, -0.02), (-0.009, 0.02, -0.02), (0.009, 0.02, -0.02), (0.009, -0.02, -0.02), (-0.009, -0.02, -0.02), (0.009, -0.02, -0.02), (0.009, -0.02, 0.02), (-0.009, -0.02, 0.02), (-0.009, 0.02, -0.02), (-0.009, 0.02, 0.02), (0.009, 0.02, 0.02), (0.009, 0.02, -0.02)] + uniform token subdivisionScheme = "none" + } +} diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/warehouse.usda b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/warehouse.usda new file mode 100644 index 000000000000..b1761c6a78e7 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/warehouse.usda @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:83be26fa2e0473b4d0d1da2978e3f8ca6e566fc64cea1fb19ede9b81fc444274 +size 458282 diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_cube_pool.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_cube_pool.py new file mode 100644 index 000000000000..f1477f79a266 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_cube_pool.py @@ -0,0 +1,80 @@ +# 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 + +"""Stable physical parcel identities behind the checkpoint's four observation slots.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import TYPE_CHECKING + +import torch + +from .mdp.reset_events import CUBE_COUNT + +if TYPE_CHECKING: + from isaaclab.assets import RigidObject + + +class ConveyorCubePool: + """Bind four policy slots to a larger pool without moving or copying physical bodies. + + Assignments are independent per environment. Local parcels keep their slots; + the active grasp is additionally pinned even if it moves outside the workcell. + """ + + def __init__(self, assets: tuple[RigidObject, ...], num_envs: int, device: str) -> None: + if len(assets) < CUBE_COUNT: + raise ValueError("The conveyor policy requires at least four physical parcels.") + self.assets = assets + self.slot_ids = torch.arange(CUBE_COUNT, device=device).repeat(num_envs, 1) + self.assignment_counts = torch.zeros((num_envs, len(assets)), device=device, dtype=torch.long) + self.assignment_counts[:, :CUBE_COUNT] = 1 + self.transfer_counts = torch.zeros_like(self.assignment_counts) + + def reset(self, env_ids: Sequence[int] | torch.Tensor) -> None: + """Restore slot identities and fresh batch counters for selected environments.""" + self.slot_ids[env_ids] = torch.arange(CUBE_COUNT, device=self.slot_ids.device) + self.assignment_counts[env_ids] = 0 + self.assignment_counts[env_ids, :CUBE_COUNT] = 1 + self.transfer_counts[env_ids] = 0 + + def refresh( + self, + positions: torch.Tensor, + local: torch.Tensor, + candidates: torch.Tensor, + target_slots: torch.Tensor, + pinned: torch.Tensor, + ) -> None: + """Assign arriving parcels to remote slots, preserving all local and pinned identities. + + Args: + positions: Physical parcel positions in the workspace [m], shape ``(N, P, 3)``. + local: Parcels still in the manipulation region, shape ``(N, P)``. + candidates: Parcels eligible for pickup, shape ``(N, P)``. + target_slots: Current command's policy slot, shape ``(N,)``. + pinned: Whether the active parcel must retain its slot, shape ``(N,)``. + """ + mapped = torch.zeros_like(candidates).scatter_(1, self.slot_ids, True) + available = candidates & ~mapped + reusable = ~local.gather(1, self.slot_ids) + reusable.scatter_(1, target_slots[:, None], ~pinned[:, None] & reusable.gather(1, target_slots[:, None])) + # Prefer parcels assigned less often, then the arrival with more belt travel remaining. + priority = positions[..., 0] - 10.0 * self.assignment_counts + for slot in range(CUBE_COUNT): + eligible = reusable[:, slot] & available.any(dim=1) + rows = eligible.nonzero(as_tuple=False).flatten() + if not rows.numel(): + continue + choices = torch.where(available, priority, -torch.inf).argmax(dim=1)[rows] + self.slot_ids[rows, slot] = choices + self.assignment_counts[rows, choices] += 1 + available[rows, choices] = False + + def record_transfers(self, env_ids: torch.Tensor, target_slots: torch.Tensor) -> None: + """Credit stable placements to physical parcels before the command selects its next slot.""" + physical_ids = self.slot_ids[env_ids, target_slots] + self.transfer_counts[env_ids, physical_ids] += 1 diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_franka_asset_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_franka_asset_env_cfg.py new file mode 100644 index 000000000000..22b4ef771ab2 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_franka_asset_env_cfg.py @@ -0,0 +1,366 @@ +# 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 + +"""Digital Twin warehouse visuals for conveyor-Franka playback.""" + +from __future__ import annotations + +from collections.abc import Callable +from functools import cache +from pathlib import Path +from typing import TYPE_CHECKING + +import isaaclab.sim as sim_utils +from isaaclab.assets import AssetBaseCfg +from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR, retrieve_file_path +from isaaclab.utils.configclass import configclass + +from isaaclab_contrib.conveyors.surface_velocity import SurfaceVelocitySpec + +if TYPE_CHECKING: + from pxr import Sdf, Usd + +from . import mdp +from .conveyor_franka_env_cfg import ( + ConveyorFrankaEnvCfg, + ConveyorFrankaSceneCfg, + _conveyor_asset, +) +from .conveyor_geometry import ( + BELT_CENTER_X, + BELT_CENTER_Y, + BELT_HALF_STRAIGHT, + BELT_TOP_Z, + BELT_TURN_RADIUS, +) +from .conveyor_warehouse_geometry import warehouse_belt_sections, warehouse_guard_meshes + +_THOR_TABLE_ASSET_PATH = f"{ISAAC_NUCLEUS_DIR}/Props/Mounts/thor_table.usd" + +# The A12 endpoints are 2.9922 m apart, its belt crown is 1.78053 m above the +# asset origin. USD point overrides extend the legs to the floor while keeping +# the upper frames unchanged; presentation roots align the deck with the policy workspace. +_A12_ENDPOINT_SEPARATION = 2.9922 +_ASSET_BELT_TOP_Z = 1.78053 +_ASSET_XY_SCALE = 2.0 * BELT_TURN_RADIUS / _A12_ENDPOINT_SEPARATION +# Preserve the upper frames' proportions. The scaled belt crown determines +# the presentation translation into the policy workspace. +_ASSET_Z_SCALE = _ASSET_XY_SCALE +_ASSET_ROOT_Z = 0.55 +_ASSET_BELT_WORLD_Z = _ASSET_ROOT_Z + _ASSET_BELT_TOP_Z * _ASSET_Z_SCALE +_PRESENTATION_OFFSET_Z = BELT_TOP_Z - _ASSET_BELT_WORLD_Z + +# A09 is a 4 m straight. Its travel-axis scale is independent of the common +# lateral scale so one asset spans the task's complete 0.88 m straight run. +_A09_LENGTH = 4.0 +_A09_X_SCALE = 2.0 * BELT_HALF_STRAIGHT / _A09_LENGTH + +# The Thor table authors its mounting surface at local z=0 and its lowest foot +# at z=-0.795 m. Uniformly scaling that distance to the robot's height above +# the warehouse floor puts every foot on the ground without distorting the table. +_THOR_TABLE_LOWEST_Z = -0.795 +_THOR_TABLE_SCALE = _PRESENTATION_OFFSET_Z / _THOR_TABLE_LOWEST_Z + + +_PHYSICS_SCHEMA_PREFIXES = ("Physics", "Physx", "Newton", "Mujoco") +_PRESENTATION_ASSETS = Path(__file__).parent / "assets" + + +@cache +def _presentation_layer(usd_path: str) -> Sdf.Layer: + """Resolve a local composition's remote assets without modifying its source layer.""" + from pxr import Sdf, UsdUtils + + source = Sdf.Layer.FindOrOpen(usd_path) + layer = Sdf.Layer.CreateAnonymous(Path(usd_path).name) + layer.TransferContent(source) + resolved = {} + for path in source.GetExternalReferences(): + if path.startswith("https://"): + resolved[path] = retrieve_file_path(path) + else: + local_path = Sdf.ComputeAssetPathRelativeToLayer(source, path) + resolved[path] = _presentation_layer(local_path).identifier + UsdUtils.ModifyAssetPaths( + layer, + lambda path: resolved[path] if path in resolved else Sdf.ComputeAssetPathRelativeToLayer(source, path), + ) + return layer + + +@sim_utils.clone +def _spawn_authored_visual( + prim_path: str, + cfg: sim_utils.UsdFileCfg, + translation: tuple[float, float, float] | None = None, + orientation: tuple[float, float, float, float] | None = None, + **kwargs, +) -> Usd.Prim: + """Compose USD scenery with cached dependencies and no simulation ownership.""" + layer = _presentation_layer(cfg.usd_path) + prim = sim_utils.create_prim( + prim_path, + translation=translation, + orientation=orientation, + scale=cfg.scale, + ) + prim.GetReferences().AddReference(layer.identifier) + sim_utils.make_uninstanceable(prim_path) + _make_usd_subtree_visual_only(prim) + _extend_conveyor_supports(prim) + return prim + + +def _extend_conveyor_supports(root_prim: Usd.Prim) -> None: + """Apply authored support parameters to referenced meshes without bundling geometry.""" + import numpy as np + + from pxr import Usd, UsdGeom, Vt + + for prim in Usd.PrimRange(root_prim): + anchor_attr = prim.GetAttribute("conveyor:supportAnchor") + if not anchor_attr or not anchor_attr.HasAuthoredValueOpinion(): + continue + anchor = anchor_attr.Get() + scale = prim.GetAttribute("conveyor:supportScale").Get() + mesh = UsdGeom.Mesh(prim) + points = np.asarray(mesh.GetPointsAttr().Get()).copy() + lower = points[:, 2] < anchor + points[lower, 2] = anchor + scale * (points[lower, 2] - anchor) + mesh.GetPointsAttr().Set(Vt.Vec3fArray.FromNumpy(points)) + mesh.GetExtentAttr().Set(Vt.Vec3fArray.FromNumpy(np.stack((points.min(axis=0), points.max(axis=0))))) + # Face-varying normals follow the deformed surface, including faces crossing the anchor. + counts = np.asarray(mesh.GetFaceVertexCountsAttr().Get()) + indices = np.asarray(mesh.GetFaceVertexIndicesAttr().Get()) + starts = np.concatenate(([0], np.cumsum(counts)[:-1])) + normal = np.cross( + points[indices[starts + 1]] - points[indices[starts]], + points[indices[starts + 2]] - points[indices[starts]], + ) + normal /= np.maximum(np.linalg.norm(normal, axis=1, keepdims=True), 1e-12) + mesh.GetNormalsAttr().Set(Vt.Vec3fArray.FromNumpy(np.repeat(normal, counts, axis=0))) + mesh.SetNormalsInterpolation(UsdGeom.Tokens.faceVarying) + + +@sim_utils.clone +def _spawn_carton_cube( + prim_path: str, + cfg: _ParcelCuboidCfg, + translation: tuple[float, float, float] | None = None, + orientation: tuple[float, float, float, float] | None = None, + **kwargs, +) -> Usd.Prim: + """Dress the original 40 mm collider with a centered, equally sized SimReady carton.""" + from pxr import UsdGeom + + prim = sim_utils.spawn_cuboid(prim_path, cfg, translation, orientation, **kwargs) + UsdGeom.Imageable(prim.GetStage().GetPrimAtPath(f"{prim_path}/geometry/mesh")).MakeInvisible() + visual_cfg = sim_utils.UsdFileCfg(usd_path=cfg.parcel_usd_path) + _spawn_authored_visual(f"{prim_path}/CartonVisual", visual_cfg) + return prim + + +@configclass +class _ParcelCuboidCfg(sim_utils.CuboidCfg): + """Original task collider with a separately authored carton appearance.""" + + parcel_usd_path: str = str(_PRESENTATION_ASSETS / "parcel.usda") + """USD visual, normalized to the task's 40 mm cube.""" + + +def _make_usd_subtree_visual_only(root_prim: Usd.Prim) -> None: + """Remove physics ownership with stage overrides, preserving the source asset.""" + from pxr import Sdf, Usd, UsdPhysics + + children = tuple(Usd.PrimRange(root_prim)) + with Sdf.ChangeBlock(): + for child in children: + if not child.IsValid() or not child.IsActive(): + continue + if ( + child.GetTypeName().startswith("OmniGraph") + or child.IsA(UsdPhysics.Scene) + or child.IsA(UsdPhysics.Joint) + ): + child.SetActive(False) + continue + for schema_name in tuple(child.GetAppliedSchemas()): + is_physics = schema_name == "IsaacConveyorAPI" or schema_name.startswith(_PHYSICS_SCHEMA_PREFIXES) + if is_physics and not child.RemoveAppliedSchema(schema_name): + raise RuntimeError(f"Failed to remove physics schema {schema_name!r} from {child.GetPath()}.") + + +@sim_utils.clone +def _spawn_visual_only_usd( + prim_path: str, + cfg: sim_utils.UsdFileCfg, + translation: tuple[float, float, float] | None = None, + orientation: tuple[float, float, float, float] | None = None, + **kwargs, +) -> Usd.Prim: + """Spawn a USD asset and strip its authored physics metadata.""" + prim = sim_utils.spawn_from_usd(prim_path, cfg, translation, orientation, **kwargs) + + # Nested rigid bodies and joints must also stay outside the physics model. + _make_usd_subtree_visual_only(prim) + return prim + + +def _visual_usd_asset( + prim_path: str, + usd_path: str, + position: tuple[float, float, float], + scale: tuple[float, float, float], + spawn_func: Callable = _spawn_authored_visual, +) -> AssetBaseCfg: + """Build one visual asset whose authored colliders are explicitly disabled.""" + return AssetBaseCfg( + prim_path=prim_path, + init_state=AssetBaseCfg.InitialStateCfg(pos=position), + spawn=sim_utils.UsdFileCfg(func=spawn_func, usd_path=usd_path, scale=scale, make_uninstanceable=True), + ) + + +@configclass +class ConveyorFrankaA09A12SceneCfg(ConveyorFrankaSceneCfg): + """Checkpoint-compatible Digital Twin scene with visual warehouse dressing.""" + + def __post_init__(self) -> None: + """Dress the workcell and extend its returns around the fixed manipulation sections.""" + super().__post_init__() + + # Translate the warehouse dressing around the unchanged policy workspace. + self.ground = AssetBaseCfg( + prim_path="/World/GroundPlane", + collision_group=-1, + init_state=AssetBaseCfg.InitialStateCfg(pos=(0.0, 0.0, _PRESENTATION_OFFSET_Z)), + spawn=sim_utils.GroundPlaneCfg(visible=False, color=(0.055, 0.065, 0.082)), + ) + self.dome_light.spawn.color = (0.70, 0.78, 0.92) + self.dome_light.spawn.intensity = 700.0 + + right_x = BELT_CENTER_X + BELT_HALF_STRAIGHT + straight_scale = (_A09_X_SCALE, _ASSET_XY_SCALE, _ASSET_Z_SCALE) + + # The visual table supports the fixed robot base and reaches the ground. + self.tabletop = _visual_usd_asset( + prim_path="{ENV_REGEX_NS}/RobotThorTableVisual", + usd_path=_THOR_TABLE_ASSET_PATH, + position=(0.0, 0.0, 0.0), + scale=(_THOR_TABLE_SCALE,) * 3, + spawn_func=_spawn_visual_only_usd, + ) + self.table_pedestal = None + + for side in ("Left", "Right"): + side_key = side.lower() + center_y = BELT_CENTER_Y if side == "Left" else -BELT_CENTER_Y + + # Authored modules replace the procedural visuals; task colliders remain authoritative. + setattr(self, f"conveyor_{side_key}_belt_visual", None) + for boundary in ("inner", "outer"): + setattr(self, f"guard_{side_key}_{boundary}", None) + + run = "bottom" if side == "Left" else "top" + y_position = center_y + (-BELT_TURN_RADIUS if side == "Left" else BELT_TURN_RADIUS) + setattr( + self, + f"conveyor_{side_key}_{run}_a09_visual", + _visual_usd_asset( + prim_path=f"{{ENV_REGEX_NS}}/Conveyor{side}{run.title()}A09Visual", + usd_path=str(_PRESENTATION_ASSETS / "conveyor_straight_supported.usda"), + position=(right_x, y_position, _ASSET_ROOT_Z + _PRESENTATION_OFFSET_Z), + scale=straight_scale, + ), + ) + + self.warehouse_visual = _visual_usd_asset( + prim_path="{ENV_REGEX_NS}/WarehouseVisual", + usd_path=str(_PRESENTATION_ASSETS / "warehouse.usda"), + position=(0.0, 0.0, _PRESENTATION_OFFSET_Z), + scale=(1.0, 1.0, 1.0), + ) + for cube_id in range(4): + cube = getattr(self, f"cube_{cube_id}") + cube.spawn = _ParcelCuboidCfg(**vars(cube.spawn)) + cube.spawn.func = _spawn_carton_cube + + def _configure_route_assets(self, parcel_colors: tuple[str, ...]) -> None: + """Read USD route geometry after the application selects its USD runtime.""" + from .conveyor_warehouse_geometry import warehouse_parcel_positions + + positions = warehouse_parcel_positions() + if len(parcel_colors) != len(positions) or not set(parcel_colors) <= {"blue", "orange", "green", "purple"}: + raise ValueError("Each authored parcel requires a color: blue, orange, green, or purple.") + for cube_id, (position, color) in enumerate(zip(positions, parcel_colors, strict=True)): + cube = self.cube_0.copy() + cube.spawn.parcel_usd_path = str(_PRESENTATION_ASSETS / f"parcel_{color}.usda") + cube.prim_path = f"{{ENV_REGEX_NS}}/Cube{cube_id}" + cube.init_state.pos = position + setattr(self, f"cube_{cube_id}", cube) + for side in ("Left", "Right"): + for key in ("top_straight", "bottom_straight", "right_turn", "left_turn"): + setattr(self, f"conveyor_{side.lower()}_{key}_collision", None) + for index, section in enumerate(warehouse_belt_sections(side)): + setattr( + self, + f"warehouse_{side.lower()}_section_{index}", + _conveyor_asset(section.belt.prim_path, section.geometry), + ) + for guard in warehouse_guard_meshes(side): + setattr( + self, + f"warehouse_{guard.name}_collision", + _conveyor_asset(f"{{ENV_REGEX_NS}}/{guard.name}Collision", guard), + ) + + def build_conveyor_belt_specs(self, **kwargs: float) -> tuple[SurfaceVelocitySpec, ...]: + """Describe the linked warehouse surfaces to the existing Newton conveyor driver.""" + return tuple(section.belt for side in ("Left", "Right") for section in warehouse_belt_sections(side, **kwargs)) + + +@configclass +class ConveyorFrankaA09A12EnvCfg(ConveyorFrankaEnvCfg): + """Newton presentation variant with Digital Twin visuals and extended return routes.""" + + class_type: type | str = "{DIR}.conveyor_franka_warehouse_env:ConveyorFrankaWarehouseEnv" + park_when_idle: bool = False + """Replace idle actions with home-position control during pretrained-policy playback only.""" + scene: ConveyorFrankaA09A12SceneCfg = ConveyorFrankaA09A12SceneCfg( + num_envs=1, + env_spacing=24.0, + ) + + def __post_init__(self) -> None: + """Frame the presentation scene and allow packages to travel along its extended returns.""" + super().__post_init__() + from .mdp.sorting import ConveyorSortCommandCfg + + self.commands.transfer = ConveyorSortCommandCfg() + # Warehouse arrivals replace the phase reset bank; fine-tuning starts from home. + self.curriculum = None + self.events.reset_from_state_table.params.update( + fixed_recipe=int(mdp.ConveyorResetRecipe.BELT), fixed_variant_id=mdp.BELT_DEPLOYMENT_VARIANT + ) + # Kit renders the authored USD directly; Newton needs only the physical geometry. + self.sim.physics.load_visual_shapes = False + # Adjacent static belt sections must not consume the parcel contact budget. + self.sim.physics.collision_cfg.include_static_kinematic_pairs = False + self.conveyor_force.transported_body_count_per_env = len(self.commands.transfer.parcel_destinations) + self.sim.physics.solver_cfg.nconmax = 400 + self.sim.physics.solver_cfg.njmax = 600 + self.conveyor_force.transported_body_pattern = r"(?:^|/)Cube_?[0-9]+(?:/|$)" + self.terminations.cube_out_of_workspace.params = { + "minimum": (-0.4, -1.2, -0.05), + "maximum": (2.85, 1.2, 0.8), + } + self.sim.visualizer_cfgs.eye = (4.8, -5.2, 3.0 + _PRESENTATION_OFFSET_Z) + self.sim.visualizer_cfgs.lookat = (0.9, 0.6, 0.95 + _PRESENTATION_OFFSET_Z) + self.sim.visualizer_cfgs.focal_length = 24.0 + + def play_mode(self) -> None: + """Enable the idle controller only for pretrained-policy playback.""" + super().play_mode() + self.park_when_idle = True 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_franka_warehouse_env.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_franka_warehouse_env.py new file mode 100644 index 000000000000..ee115d7fe312 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_franka_warehouse_env.py @@ -0,0 +1,156 @@ +# 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-slot policy control over a physical parcel pool and USD-authored warehouse.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import TYPE_CHECKING + +import torch + +from isaaclab.envs.common import VecEnvStepReturn + +from .conveyor_cube_pool import ConveyorCubePool +from .conveyor_franka_env import ConveyorFrankaEnv +from .conveyor_warehouse_geometry import warehouse_parcel_positions + +if TYPE_CHECKING: + from pxr import Usd + + from .conveyor_franka_asset_env_cfg import ConveyorFrankaA09A12EnvCfg + + +class ConveyorFrankaWarehouseEnv(ConveyorFrankaEnv): + """Use the four-cube policy contract over an individually tracked workcell parcel pool. + + Kit renders the authored USD materials, lights, and animation. Lightweight + Newton viewers show a static approximation of the warehouse dressing. + """ + + cfg: ConveyorFrankaA09A12EnvCfg + + def __init__(self, cfg: ConveyorFrankaA09A12EnvCfg, render_mode: str | None = None, **kwargs) -> None: + self._warehouse_animation: list[tuple[Usd.Attribute, Usd.Attribute]] = [] + cfg.scene._configure_route_assets(cfg.commands.transfer.parcel_colors) + super().__init__(cfg, render_mode=render_mode, **kwargs) + viewer = next((viz.cfg for viz in self.sim.visualizers if viz.cfg.backend.visualizer_type == "kit"), None) + if viewer is None: + return + self.sim.set_setting("/app/viewport/grid/enabled", False) + + from pxr import Usd + + from .conveyor_franka_asset_env_cfg import _presentation_layer + + camera = self.sim.stage.GetPrimAtPath("/OmniverseKit_Persp") + if camera: + # Kit owns this camera in its session layer; weaker root-layer edits are ignored. + with Usd.EditContext(self.sim.stage, self.sim.stage.GetSessionLayer()): + self.sim.set_camera_view(viewer.eye, viewer.lookat) + camera.GetAttribute("focalLength").Set(viewer.focal_length) + self._warehouse_source = Usd.Stage.Open(_presentation_layer(cfg.scene.warehouse_visual.spawn.usd_path)) + self._warehouse_period = self._warehouse_source.GetEndTimeCode() + self._warehouse_fps = self._warehouse_source.GetTimeCodesPerSecond() + for env_path in self.scene.env_prim_paths: + group = self._warehouse_source.GetPrimAtPath("/Warehouse/Parcels") + for source in group.GetChildren(): + target = self.sim.stage.GetPrimAtPath(f"{env_path}/WarehouseVisual/Parcels/{source.GetName()}") + for name in ("xformOp:translate", "xformOp:rotateXYZ"): + self._warehouse_animation.append((source.GetAttribute(name), target.GetAttribute(name))) + + def load_managers(self) -> None: + """Create slot identities before command, reward, and observation terms inspect cubes.""" + assets = tuple(self.scene[f"cube_{i}"] for i in range(self.cfg.conveyor_force.transported_body_count_per_env)) + self.conveyor_cube_pool = ConveyorCubePool(assets, self.num_envs, self.device) + self._parcel_start_positions = torch.tensor(warehouse_parcel_positions(), device=self.device) + self._warehouse_arm_joint_ids = self.scene["robot"].find_joints( + self.cfg.actions.arm_action.joint_names, preserve_order=True + )[0] + super().load_managers() + + def step(self, action: torch.Tensor) -> VecEnvStepReturn: + """Advance presentation motion; pass training actions through or park during idle playback.""" + if self._warehouse_animation: + self._animate_warehouse() + if not self.cfg.park_when_idle: + return super().step(action) + command = self.command_manager.get_term("transfer") + robot = self.scene["robot"] + joint_ids = self._warehouse_arm_joint_ids + parked = torch.zeros_like(action) + parked[:, :7] = ( + (robot.data.default_joint_pos.torch[:, joint_ids] - robot.data.joint_pos.torch[:, joint_ids]) + / self.cfg.actions.arm_action.scale + ).clamp(-0.25, 0.25) + # Keep invalid actions visible to the existing sanitization and termination terms. + use_policy = command.has_target | ~torch.isfinite(action).all(dim=1) + return super().step(torch.where(use_policy[:, None], action, parked)) + + def close(self) -> None: + """Release presentation attributes before closing the shared task resources.""" + self._warehouse_animation.clear() + super().close() + + def _adapt_policy_cube_state( + self, positions: torch.Tensor, quaternions: torch.Tensor, velocities: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Represent remote inventory as waiting slots while keeping local manipulation states exact.""" + origins = self.scene.env_origins[:, None, :] + local = positions - origins + remote = ~self._in_workcell(local) + waiting = local.clone() + waiting[..., 0] = 0.14 + 0.88 * torch.arange(1, 5, device=self.device)[None, :] / 5 + waiting[..., 1] = torch.where(local[..., 1] >= 0, 0.75, -0.75) + waiting[..., 2] = 0.06 + upright = torch.zeros_like(quaternions) + upright[..., 3] = 1.0 + transport_velocity = torch.zeros_like(velocities) + transport_velocity[..., 0] = self.cfg.conveyor_force.speed + return ( + torch.where(remote[..., None], waiting + origins, positions), + torch.where(remote[..., None], upright, quaternions), + torch.where(remote[..., None], transport_velocity, velocities), + ) + + def _animate_warehouse(self) -> None: + """Sample USD motion using policy time, independently of render frame rate.""" + from pxr import Sdf + + time_code = (self.common_step_counter * self.step_dt * self._warehouse_fps) % self._warehouse_period + with Sdf.ChangeBlock(): + for source, target in self._warehouse_animation: + target.Set(source.Get(time_code)) + + def _reset_idx(self, env_ids: Sequence[int] | torch.Tensor | slice) -> None: + """Shuffle a mixed batch across the authored feeds in selected environments.""" + if isinstance(env_ids, slice): + env_ids = self.scene._ALL_INDICES[env_ids] + pool = self.conveyor_cube_pool + pool.reset(env_ids) + super()._reset_idx(env_ids) + if self.cfg.commands.transfer.randomize_arrivals: + assignments = torch.rand((len(env_ids), len(pool.assets)), device=self.device).argsort(dim=1) + else: + assignments = torch.arange(len(pool.assets), device=self.device).expand(len(env_ids), -1) + for cube_id, cube in enumerate(pool.assets): + pose = cube.data.root_pose_w.torch[env_ids].clone() + pose[:, :3] = self._parcel_start_positions[assignments[:, cube_id]] + self.scene.env_origins[env_ids] + pose[:, 3:] = pose.new_tensor((0.0, 0.0, 0.0, 1.0)) + cube.write_root_pose_to_sim_index(root_pose=pose, env_ids=env_ids) + cube.write_root_velocity_to_sim_index(root_velocity=pose.new_zeros((len(env_ids), 6)), env_ids=env_ids) + + @staticmethod + def _in_workcell(positions: torch.Tensor) -> torch.Tensor: + """Identify parcels within the trained controller's local manipulation region [m].""" + return ( + (positions[..., 0] > -0.25) + & (positions[..., 0] < 1.5) + & (positions[..., 1].abs() < 0.6) + & (positions[..., 2] < 0.4) + # The low merge passes beside the placement bend but is still remote transport. + & ((positions[..., 2] < 0.12) | (positions[..., 0] < 1.05)) + ) 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/conveyor_warehouse_geometry.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_warehouse_geometry.py new file mode 100644 index 000000000000..85e6b6d8d401 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_warehouse_geometry.py @@ -0,0 +1,230 @@ +# 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 + +"""Physical surfaces for the USD-authored elevated conveyor network.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass, replace +from functools import lru_cache +from pathlib import Path +from typing import TYPE_CHECKING + +import numpy as np + +from isaaclab_contrib.conveyors.surface_velocity import SurfaceVelocitySpec + +from .conveyor_geometry import ( + BELT_CENTER_X, + BELT_CENTER_Y, + BELT_COLLISION_OVERHANG, + BELT_COLLISION_SEAM_OVERLAP, + BELT_HALF_STRAIGHT, + BELT_THICKNESS, + GUARD_BASE_OVERLAP, + GUARD_HEIGHT, + GUARD_THICKNESS, + ConveyorSectionSpec, + MeshSpec, + _turn_collision_mesh, + belt_collision_section_specs, + belt_direction, +) + +if TYPE_CHECKING: + from pxr import Sdf + +_ROUTE_ASSET = Path(__file__).parent / "assets" / "conveyor_routes.usda" + + +@lru_cache(maxsize=1) +def _route_layer() -> Sdf.Layer: + # Import Usd to register file formats in kitless mode, without composing asset dependencies. + from pxr import Sdf, Usd # noqa: F401 + + return Sdf.Layer.FindOrOpen(str(_ROUTE_ASSET)) + + +@dataclass(frozen=True) +class _RouteSegment: + name: str + points: tuple[tuple[float, float, float], ...] + direction: tuple[float, float, float] + width: float + pivot: tuple[float, float, float] | None + radius: float | None + friction: float | None + velocity: float | None + path: str + + +@lru_cache(maxsize=2) +def _route_segments(side: str) -> tuple[_RouteSegment, ...]: + """Read the same centerlines that author the visible conveyor modules.""" + from pxr import Sdf + + belt_direction(side) + layer = _route_layer() + root = layer.GetPrimAtPath(Sdf.Path(f"/ConveyorRoutes/{side}")) + segments = [] + for name in root.attributes["conveyor:order"].default: + prim = layer.GetPrimAtPath(root.path.AppendChild(name)) + points = layer.GetPrimAtPath(prim.path.AppendChild("Centerline")).attributes["points"].default + curved = "conveyor:pivot" in prim.attributes + segments.append( + _RouteSegment( + name=name, + velocity=prim.attributes["conveyor:velocity"].default + if "conveyor:velocity" in prim.attributes + else None, + path=prim.attributes["conveyor:path"].default if "conveyor:path" in prim.attributes else "main", + points=tuple(tuple(point) for point in points), + direction=tuple(prim.attributes["conveyor:direction"].default), + width=prim.attributes["conveyor:width"].default, + pivot=tuple(prim.attributes["conveyor:pivot"].default) if curved else None, + radius=prim.attributes["conveyor:radius"].default if curved else None, + friction=prim.attributes["conveyor:friction"].default + if "conveyor:friction" in prim.attributes + else None, + ) + ) + return tuple(segments) + + +def warehouse_parcel_positions() -> tuple[tuple[float, float, float], ...]: + """Return physical parcel infeed spawn positions in workspace coordinates [m].""" + root = _route_layer().GetPrimAtPath("/ConveyorRoutes") + return tuple(tuple(position) for position in root.attributes["conveyor:parcelSpawnPositions"].default) + + +def _normals(segment: _RouteSegment) -> np.ndarray: + if segment.pivot is None: + tangent = np.asarray(segment.direction) + normal = np.array((-tangent[1], tangent[0], 0.0)) + return np.tile(normal / np.linalg.norm(normal), (len(segment.points), 1)) + radial = np.asarray(segment.points) - segment.pivot + radial[:, 2] = 0 + return -math.copysign(1, segment.direction[2]) * radial / np.linalg.norm(radial, axis=1, keepdims=True) + + +def _prism( + name: str, + points: np.ndarray, + normals: np.ndarray, + offset: float | np.ndarray, + width: float, + bottom: float, + top: float, + *, + closed: bool = False, +) -> MeshSpec: + """Sweep a closed solid along a three-dimensional centerline [m].""" + offset = np.broadcast_to(offset, (len(points),))[:, None] + left = points + normals * (offset - width / 2) + right = points + normals * (offset + width / 2) + vertices = np.concatenate((left + (0, 0, top), right + (0, 0, top), left + (0, 0, bottom), right + (0, 0, bottom))) + count = len(points) + quads = [] + for i in range(count if closed else count - 1): + j = (i + 1) % count + quads.extend( + ( + (i, j, count + j, count + i), + (2 * count + i, 3 * count + i, 3 * count + j, 2 * count + j), + (i, 2 * count + i, 2 * count + j, j), + (count + i, count + j, 3 * count + j, 3 * count + i), + ) + ) + if not closed: + quads.extend(((0, count, 3 * count, 2 * count), (count - 1, 3 * count - 1, 4 * count - 1, 2 * count - 1))) + faces = tuple(triangle for a, b, c, d in quads for triangle in ((a, b, c), (a, c, d))) + return MeshSpec(name, tuple(tuple(vertex) for vertex in vertices), faces) + + +def warehouse_belt_sections(side: str, **kwargs: float) -> tuple[ConveyorSectionSpec, ...]: + """Build linked belt surfaces from USD while preserving the original manipulation geometry.""" + sign = belt_direction(side) + sections = [] + for segment in _route_segments(side): + name = f"Conveyor{side}{segment.name}Collision" + pivot = segment.pivot + surface_kwargs = dict(kwargs) + if segment.velocity is not None: + surface_kwargs["velocity"] = segment.velocity + if segment.friction is not None: + surface_kwargs["friction_coefficient"] = segment.friction + if segment.name == "Working": + original = belt_collision_section_specs(side, **kwargs)[1 if sign > 0 else 0] + sections.append(original) + continue + if segment.name in {"PickupBend", "PlacementBend"}: + pickup = segment.name == "PickupBend" + name = f"Conveyor{side}{'Left' if pickup else 'Right'}InnerTurnCollision" + pivot = (BELT_CENTER_X + (-BELT_HALF_STRAIGHT if pickup else BELT_HALF_STRAIGHT), sign * BELT_CENTER_Y, 0.0) + geometry = _turn_collision_mesh( + name, pivot[0], BELT_CENTER_Y, math.pi if pickup else -math.pi / 2, math.pi / 2 + ) + if sign < 0: + geometry = replace( + geometry, + vertices=tuple((x, -y, z) for x, y, z in geometry.vertices), + faces=tuple((a, c, b) for a, b, c in geometry.faces), + ) + else: + points = np.array(segment.points) + # Overlap adjacent panels without exposing a vertical collision seam. + for index, neighbor, direction in ((0, 1, -1), (-1, -2, 1)): + tangent = points[neighbor] - points[index] if index == 0 else points[index] - points[neighbor] + points[index] += direction * BELT_COLLISION_SEAM_OVERLAP * tangent / np.linalg.norm(tangent) + geometry = _prism( + name, points, _normals(segment), 0, segment.width + 2 * BELT_COLLISION_OVERHANG, -BELT_THICKNESS, 0 + ) + sections.append( + ConveyorSectionSpec( + geometry, + SurfaceVelocitySpec( + prim_path=f"{{ENV_REGEX_NS}}/{geometry.name}", + direction=segment.direction, + surface_normal=(0.0, 0.0, 1.0) + if segment.pivot is not None + else tuple(np.cross(segment.direction, _normals(segment)[0])), + curved=segment.pivot is not None, + pivot_point=pivot or (0, 0, 0), + radius=segment.radius, + **surface_kwargs, + ), + ) + ) + return tuple(sections) + + +def warehouse_guard_meshes(side: str) -> tuple[MeshSpec, ...]: + """Build continuous guides for each closed circulation route and open supply belt.""" + segments = _route_segments(side) + guards = [] + for path in dict.fromkeys(segment.path for segment in segments): + points, normals, offsets = [], [], [] + route = [segment for segment in segments if segment.path == path] + for index, segment in enumerate(route): + stop = None if path != "main" and index == len(route) - 1 else -1 + for point, normal in zip(segment.points[:stop], _normals(segment)[:stop]): + points.append(point) + normals.append(normal) + offsets.append((segment.width + GUARD_THICKNESS) / 2) + for boundary, sign in (("Inner", -1), ("Outer", 1)): + guards.append( + _prism( + f"Guard{side}{path.title()}{boundary}", + np.asarray(points), + np.asarray(normals), + sign * np.asarray(offsets), + GUARD_THICKNESS, + -GUARD_BASE_OVERLAP, + GUARD_HEIGHT, + closed=path == "main", + ) + ) + return tuple(guards) 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..b5bc6e573632 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/commands.py @@ -0,0 +1,289 @@ +# 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 + pool = getattr(self._env, "conveyor_cube_pool", None) + if pool is not None: + pool.record_transfers(success_ids, self.target_cube_ids[success_ids]) + + 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..c583f9c955f1 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/cube_state.py @@ -0,0 +1,32 @@ +# 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 policy-slot or physical inventory order.""" + +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 a rigid-object data attribute in policy-slot or complete physical-inventory order. + + Values retain the attribute's units and world frame. The result has shape + ``(num_envs, num_slots_or_parcels, attribute_size)``. Tasks without a parcel pool + retain the original four fixed cube identities. + """ + pool = getattr(env, "conveyor_cube_pool", None) + assets = pool.assets if pool is not None else tuple(env.scene[f"cube_{i}"] for i in range(CUBE_COUNT)) + values = torch.stack([getattr(asset.data, attribute).torch for asset in assets], dim=1) + if pool is not None and not all_cubes: + values = values.gather(1, pool.slot_ids[..., None].expand(-1, -1, values.shape[-1])) + return values 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..ebc7e8679074 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/observations.py @@ -0,0 +1,135 @@ +# 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"), + ) + adapter = getattr(env, "_adapt_policy_cube_state", None) + return state if adapter is None else adapter(*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. The base task keeps fixed identities; + warehouse playback can refill remote slots from its physical parcel pool. + :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/sorting.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/sorting.py new file mode 100644 index 000000000000..597ec6445034 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/sorting.py @@ -0,0 +1,141 @@ +# 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 + +"""Parcel-class dispatch for the warehouse's unchanged four-slot transfer policy.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import TYPE_CHECKING + +import torch + +from isaaclab.utils.configclass import configclass + +from .commands import ConveyorTransferCommand, ConveyorTransferCommandCfg +from .cube_state import cube_values +from .kinematics import end_effector_pose +from .reset_events import CUBE_REST_Z + +if TYPE_CHECKING: + from ..conveyor_franka_warehouse_env import ConveyorFrankaWarehouseEnv + + +class ConveyorSortCommand(ConveyorTransferCommand): + """Transfer misplaced cartons and leave correctly sorted inventory circulating. + + Physical parcel IDs carry immutable destination classes. The dispatcher presents + one wrong-lane arrival through the pretrained policy's existing cube/side command; + color is a visual class label, not an additional policy observation. + """ + + cfg: ConveyorSortCommandCfg + + def __init__(self, cfg: ConveyorSortCommandCfg, env: ConveyorFrankaWarehouseEnv) -> None: + super().__init__(cfg, env) + if not 0.14 <= cfg.pickup_x_range[0] < cfg.pickup_x_range[1] <= 1.02: + raise ValueError("The sorting pickup window must lie within the original working straight [0.14, 1.02] m.") + if len(cfg.parcel_destinations) != len(env.conveyor_cube_pool.assets) or set(cfg.parcel_destinations) != {0, 1}: + raise ValueError("Parcel destinations must match the physical pool and include both conveyor IDs, 0 and 1.") + if len(cfg.parcel_colors) != len(cfg.parcel_destinations): + raise ValueError("Each physical parcel requires a color and a destination.") + for color in set(cfg.parcel_colors): + destinations = { + side for shade, side in zip(cfg.parcel_colors, cfg.parcel_destinations, strict=True) if shade == color + } + if len(destinations) != 1: + raise ValueError(f"All {color} parcels must share one destination conveyor.") + self.parcel_destinations = torch.tensor(cfg.parcel_destinations, device=self.device) + self.has_target = torch.zeros(self.num_envs, device=self.device, dtype=torch.bool) + self.metrics["sorted_parcels"] = torch.zeros(self.num_envs, device=self.device) + self.metrics["batch_complete"] = torch.zeros(self.num_envs, device=self.device) + + def evaluate(self) -> None: + """Credit stable placements only while the dispatcher owns an active transfer.""" + self.new_success[~self.has_target] = False + self.is_success[~self.has_target] = False + self._last_evaluation_steps[~self.has_target] = self._env.episode_length_buf[~self.has_target] + super().evaluate() + + def _resample_command(self, env_ids: Sequence[int]) -> None: + # Reset metadata remains compatible with the shared reward and observation terms. + if self._resampling_from_reset: + super()._resample_command(env_ids) + self.has_target[env_ids] = False + self.held_cube_ids[env_ids] = -1 + self.pending_success[env_ids] = False + + def _physical_cube_acquisition_mask(self, active_position_w: torch.Tensor) -> torch.Tensor: + """Pin a lifted, closed-gripper parcel near the tool while the arm transfers it.""" + tool_position, _ = end_effector_pose(self._env) + finger_positions = self._robot.data.joint_pos.torch[:, self._finger_joint_ids] + local_cube_z = active_position_w[:, 2] - self._env.scene.env_origins[:, 2] + return ( + (local_cube_z >= CUBE_REST_Z + 0.025) + & (torch.linalg.vector_norm(active_position_w - tool_position, dim=1) <= 0.075) + & (torch.amax(finger_positions, dim=1) <= 0.030) + ) + + def _update_command(self) -> None: + """Publish slot assignments before the manager computes the next policy observation.""" + env = self._env + pool = env.conveyor_cube_pool + positions_w = cube_values(env, "root_pos_w", all_cubes=True) + positions = positions_w - env.scene.env_origins[:, None] + local = env._in_workcell(positions) + rows = env.scene._ALL_INDICES + physical_target = pool.slot_ids[rows, self.target_cube_ids] + held = self._physical_cube_acquisition_mask(positions_w[rows, physical_target]) & self.has_target + self.has_target &= ~self.pending_success & (local[rows, physical_target] | held) + self.pending_success.zero_() + wrong_lane = (positions[..., 1] < 0).long() != self.parcel_destinations + candidates = ( + wrong_lane + & (positions[..., 0] > self.cfg.pickup_x_range[0]) + & (positions[..., 0] < self.cfg.pickup_x_range[1]) + & (positions[..., 1].abs() > 0.20) + & (positions[..., 1].abs() < 0.36) + & (positions[..., 2] > 0.04) + & (positions[..., 2] < 0.10) + ) + pool.refresh(positions, local, candidates, self.target_cube_ids, self.has_target) + available = candidates.gather(1, pool.slot_ids) + eligible = available.any(dim=1) & ~self.has_target + choices = torch.where(available, positions[..., 0].gather(1, pool.slot_ids), -torch.inf).argmax(dim=1) + env_ids = eligible.nonzero(as_tuple=False).flatten() + slots = choices[env_ids] + physical_ids = pool.slot_ids[env_ids, slots] + self._assign_goal(env_ids, slots, (positions[env_ids, physical_ids, 1] < 0.0).long()) + self.has_target[env_ids] = True + self.command_counter[env_ids] += 1 + # A class is counted only after leaving the elevated supply and landing on a loop. + velocities = cube_values(env, "root_lin_vel_w", all_cubes=True) + on_loop = (positions[..., 2] > 0.04) & (positions[..., 2] < 0.20) & (velocities[..., 2].abs() < 0.15) + active = torch.zeros_like(on_loop).scatter_( + 1, pool.slot_ids[rows, self.target_cube_ids, None], self.has_target[:, None] + ) + on_loop &= ~active + sorted_parcels = (~wrong_lane & on_loop).sum(dim=1) + self.metrics["sorted_parcels"].copy_(sorted_parcels) + self.metrics["batch_complete"].copy_(sorted_parcels == len(pool.assets)) + + +@configclass +class ConveyorSortCommandCfg(ConveyorTransferCommandCfg): + """Batch sorting with fixed physical classes and checkpoint-compatible commands.""" + + class_type: type[ConveyorSortCommand] | str = "{DIR}.sorting:ConveyorSortCommand" + + parcel_destinations: tuple[int, ...] = (0, 1) * 12 + """Destination per physical parcel: 0 is the positive-Y loop, 1 the negative-Y loop.""" + + parcel_colors: tuple[str, ...] = ("blue", "orange", "green", "purple") * 6 + """Authored color per physical parcel; blue, orange, green, and purple are available.""" + + randomize_arrivals: bool = True + """Shuffle physical parcel identities across authored start positions at each reset.""" + + pickup_x_range: tuple[float, float] = (0.35, 1.02) + """Longitudinal arrival window in policy workspace coordinates [m].""" 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..bb4b2e7eab8d 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" @@ -39,3 +43,4 @@ include = ["isaaclab_tasks", "isaaclab_tasks.*"] [tool.setuptools.package-data] "*" = ["*.pyi"] +"isaaclab_tasks.contrib.conveyor_franka" = ["assets/*.usda"] diff --git a/source/isaaclab_tasks/test/contrib/test_conveyor_franka_asset_cfg.py b/source/isaaclab_tasks/test/contrib/test_conveyor_franka_asset_cfg.py new file mode 100644 index 000000000000..bf99c17ca99d --- /dev/null +++ b/source/isaaclab_tasks/test/contrib/test_conveyor_franka_asset_cfg.py @@ -0,0 +1,236 @@ +# 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 checks for the Digital Twin conveyor playback scene.""" + +import math +import subprocess +import sys +from types import SimpleNamespace + +import pytest + +import isaaclab.sim as sim_utils + +from isaaclab_tasks.contrib.conveyor_franka.conveyor_franka_asset_env_cfg import ( + _PRESENTATION_ASSETS, + ConveyorFrankaA09A12EnvCfg, + _presentation_layer, +) +from isaaclab_tasks.contrib.conveyor_franka.conveyor_franka_env_cfg import ConveyorFrankaEnvCfg +from isaaclab_tasks.contrib.conveyor_franka.conveyor_geometry import ( + BELT_CENTER_X, + BELT_CENTER_Y, + BELT_HALF_STRAIGHT, + BELT_TOP_Z, + BELT_TURN_RADIUS, +) + + +def test_a09_a12_config_import_does_not_preload_usd() -> None: + """Task discovery must not import USD before a requested Kit application starts.""" + code = ( + "import sys; " + "import isaaclab_tasks.contrib.conveyor_franka.conveyor_franka_asset_env_cfg; " + "raise SystemExit('pxr' in sys.modules)" + ) + result = subprocess.run([sys.executable, "-c", code], check=False) + assert result.returncode == 0 + + +def test_carton_visual_is_centered_on_the_original_40_mm_collider(tmp_path, monkeypatch) -> None: + """Asset normalization changes appearance without moving or resizing the grasp surface.""" + from pxr import Sdf, Usd, UsdGeom, UsdPhysics + + from isaaclab_tasks.contrib.conveyor_franka import conveyor_franka_asset_env_cfg as asset_cfg + + # Measured unscaled Cardbox_A1 bounds, used as an offline stand-in for the remote mesh. + lower = (-0.34971755743026733, -0.260576993227005, 0.0) + upper = (0.34971755743026733, 0.2605747878551483, 0.5099270939826965) + source = Usd.Stage.CreateNew(str(tmp_path / "carton.usda")) + root = UsdGeom.Xform.Define(source, "/Carton").GetPrim() + source.SetDefaultPrim(root) + UsdPhysics.RigidBodyAPI.Apply(root) + UsdPhysics.ArticulationRootAPI.Apply(root) + UsdPhysics.MassAPI.Apply(root) + UsdPhysics.FixedJoint.Define(source, "/Carton/Joint") + source.DefinePrim("/Carton/Graph", "OmniGraph") + UsdPhysics.Scene.Define(source, "/Carton/PhysicsScene") + mesh = UsdGeom.Cube.Define(source, "/Carton/Mesh") + mesh.CreateSizeAttr(1.0) + mesh.AddTranslateOp().Set(tuple((a + b) / 2 for a, b in zip(lower, upper))) + mesh.AddScaleOp().Set(tuple(b - a for a, b in zip(lower, upper))) + UsdPhysics.CollisionAPI.Apply(mesh.GetPrim()) + UsdPhysics.MeshCollisionAPI.Apply(mesh.GetPrim()) + UsdPhysics.FilteredPairsAPI.Apply(mesh.GetPrim()) + mesh.GetPrim().AddAppliedSchema("PhysxCollisionAPI") + support = UsdGeom.Mesh.Define(source, "/Carton/Support") + support.CreatePointsAttr([(0.1, 0, 0.2), (0.1, 0.1, 0.255), (0, 0.1, 0.3)]) + support.CreateFaceVertexCountsAttr([3]) + support.CreateFaceVertexIndicesAttr([0, 1, 2]) + support.GetPrim().CreateAttribute("conveyor:supportAnchor", Sdf.ValueTypeNames.Double).Set(0.255) + support.GetPrim().CreateAttribute("conveyor:supportScale", Sdf.ValueTypeNames.Double).Set(2.0) + source.GetRootLayer().Save() + monkeypatch.setattr(asset_cfg, "retrieve_file_path", lambda path: str(tmp_path / "carton.usda")) + _presentation_layer.cache_clear() + try: + stage = sim_utils.create_new_stage() + cfg = ConveyorFrankaA09A12EnvCfg().scene.cube_0.spawn + cfg.parcel_usd_path = str(_PRESENTATION_ASSETS / "parcel_blue.usda") + prim = cfg.func("/Cube", cfg) + visual = stage.GetPrimAtPath("/Cube/CartonVisual") + bounds = UsdGeom.BBoxCache(Usd.TimeCode.Default(), ["default", "render"]).ComputeWorldBound(visual) + assert tuple(bounds.ComputeAlignedRange().GetMin()) == pytest.approx((-0.02,) * 3, abs=1e-8) + assert tuple(bounds.ComputeAlignedRange().GetMax()) == pytest.approx((0.02,) * 3, abs=1e-8) + extended = UsdGeom.Mesh(next(p for p in Usd.PrimRange(visual) if p.GetName() == "Support")) + assert [p[2] for p in extended.GetPointsAttr().Get()] == pytest.approx((0.145, 0.255, 0.3)) + assert [p[2] for p in support.GetPointsAttr().Get()] == pytest.approx((0.2, 0.255, 0.3)) + assert prim.HasAPI(UsdPhysics.RigidBodyAPI) + collider = stage.GetPrimAtPath("/Cube/geometry/mesh") + assert collider.HasAPI(UsdPhysics.CollisionAPI) + assert UsdGeom.Imageable(collider).ComputeVisibility() == "invisible" + assert sum(p.HasAPI(UsdPhysics.RigidBodyAPI) for p in stage.Traverse()) == 1 + assert sum(p.HasAPI(UsdPhysics.CollisionAPI) for p in stage.Traverse()) == 1 + imported = tuple(Usd.PrimRange(visual, Usd.TraverseInstanceProxies())) + for api in ( + UsdPhysics.ArticulationRootAPI, + UsdPhysics.RigidBodyAPI, + UsdPhysics.MassAPI, + UsdPhysics.CollisionAPI, + UsdPhysics.MeshCollisionAPI, + UsdPhysics.FilteredPairsAPI, + ): + assert not any(p.HasAPI(api) for p in imported) + assert all("PhysxCollisionAPI" not in p.GetAppliedSchemas() for p in imported) + # Inactive imported graphs, joints and scenes must not survive composition as executable prims. + assert not any(p.GetTypeName() in {"PhysicsFixedJoint", "OmniGraph", "PhysicsScene"} for p in imported) + # Resolving cached assets must leave the authored layer portable. + assert all( + path.startswith("https://") + for path in Sdf.Layer.FindOrOpen(str(_PRESENTATION_ASSETS / "parcel.usda")).GetExternalReferences() + ) + finally: + _presentation_layer.cache_clear() + + +def test_asset_transforms_preserve_the_original_workcell() -> None: + """The working straights and adjoining quarter-turns keep their original locations and radii.""" + from pxr import Sdf + + from isaaclab_tasks.contrib.conveyor_franka.conveyor_geometry import belt_collision_section_specs + from isaaclab_tasks.contrib.conveyor_franka.conveyor_warehouse_geometry import ( + warehouse_belt_sections, + ) + + cfg = ConveyorFrankaA09A12EnvCfg() + scene = cfg.scene + layer = Sdf.Layer.FindOrOpen(scene.warehouse_visual.spawn.usd_path) + scene._configure_route_assets(cfg.commands.transfer.parcel_colors) + assert scene.robot.init_state.pos == ConveyorFrankaEnvCfg().scene.robot.init_state.pos + assert math.isclose(scene.cube_0.init_state.pos[2], 0.28) + base_z = ConveyorFrankaEnvCfg().scene.conveyor_left_top_straight_collision.init_state.pos[2] + assert math.isclose(scene.warehouse_left_section_0.init_state.pos[2], base_z) + for side, sign, original_index in (("Left", 1, 1), ("Right", -1, 0)): + sections = warehouse_belt_sections(side, velocity=0.35) + original = belt_collision_section_specs(side)[original_index].geometry + inner = sections[0].geometry + assert inner.position == original.position + assert inner.size == original.size + visual = getattr(scene, f"conveyor_{side.lower()}_{'bottom' if sign > 0 else 'top'}_a09_visual") + assert visual.init_state.pos[:2] == ( + BELT_CENTER_X + BELT_HALF_STRAIGHT, + sign * (BELT_CENTER_Y - BELT_TURN_RADIUS), + ) + assert 4 * visual.spawn.scale[0] == pytest.approx(2 * BELT_HALF_STRAIGHT) + assert visual.init_state.pos[2] + 1.78053 * visual.spawn.scale[2] == pytest.approx(BELT_TOP_Z) + for name, x in ( + ("RightInner", BELT_CENTER_X + BELT_HALF_STRAIGHT), + ("LeftInner", BELT_CENTER_X - BELT_HALF_STRAIGHT), + ): + turn = next(section.belt for section in sections if f"{name}Turn" in section.geometry.name) + assert turn.pivot_point == (x, sign * BELT_CENTER_Y, 0) + assert turn.radius == BELT_TURN_RADIUS + authored = layer.GetPrimAtPath(f"/Warehouse/WorkcellConveyors/{side}{name}") + position = authored.attributes["xformOp:translate"].default + scale = authored.attributes["xformOp:scale"].default + angle = math.radians(authored.attributes["xformOp:rotateZ"].default) + # A03's original circular pivot is at (0, -1.4961); the referenced arc lands on the physical bend. + radius = 1.4961 * scale[0] + assert radius == pytest.approx(BELT_TURN_RADIUS) + assert position[0] + radius * math.sin(angle) == pytest.approx(x) + assert position[1] - radius * math.cos(angle) == pytest.approx(sign * BELT_CENTER_Y) + + +def test_warehouse_animation_uses_active_kit_viewer_and_policy_time(tmp_path, monkeypatch) -> None: + """A CLI-selected Kit viewer animates and loops authored parcels without changing task state.""" + import torch + + from pxr import Usd, UsdGeom + + from isaaclab_tasks.contrib.conveyor_franka.conveyor_franka_env import ConveyorFrankaEnv + from isaaclab_tasks.contrib.conveyor_franka.conveyor_franka_warehouse_env import ConveyorFrankaWarehouseEnv + + path = str(tmp_path / "traffic.usda") + source = Usd.Stage.CreateNew(path) + root = UsdGeom.Xform.Define(source, "/Warehouse").GetPrim() + source.SetDefaultPrim(root) + source.SetTimeCodesPerSecond(60) + source.SetEndTimeCode(120) + parcel = UsdGeom.Xform.Define(source, "/Warehouse/Parcels/Parcel00") + translate = parcel.AddTranslateOp() + translate.Set((0, 0, 1), 0) + translate.Set((2, 0, 1), 120) + parcel.AddRotateXYZOp().Set((0, 0, 0)) + source.GetRootLayer().Save() + stage = Usd.Stage.CreateInMemory() + stage.DefinePrim("/World/envs/env_0/WarehouseVisual").GetReferences().AddReference(path) + camera = UsdGeom.Camera.Define(stage, "/OmniverseKit_Persp") + + def initialize(env, cfg, **kwargs): + env.cfg = cfg + env.common_step_counter = 60 + env.scene = SimpleNamespace(env_prim_paths=["/World/envs/env_0"]) + env.sim = SimpleNamespace( + stage=stage, + visualizers=[ + SimpleNamespace( + cfg=SimpleNamespace( + backend=SimpleNamespace(visualizer_type="kit"), + eye=(4.0, -5.0, 3.0), + lookat=(0.9, 0.6, 0.95), + focal_length=37.0, + ) + ) + ], + set_camera_view=lambda eye, lookat: None, + set_setting=lambda name, value: None, + add_render_callback=lambda name, callback: None, + remove_render_callback=lambda name: None, + ) + + monkeypatch.setattr(ConveyorFrankaEnv, "__init__", initialize) + monkeypatch.setattr(ConveyorFrankaEnv, "close", lambda env: None) + monkeypatch.setattr(ConveyorFrankaEnv, "step", lambda env, action: None) + cfg = ConveyorFrankaA09A12EnvCfg() + # CLI viewer selection is resolved by SimulationContext, not written back into this list. + cfg.sim.visualizer_cfgs = [] + cfg.scene.warehouse_visual.spawn.usd_path = path + env = ConveyorFrankaWarehouseEnv(cfg) + assert camera.GetFocalLengthAttr().Get() == 37.0 + action = torch.zeros(1, 8) + try: + # Headless Kit recording captures on demand, without invoking render callbacks. + env.step(action) + position = stage.GetPrimAtPath("/World/envs/env_0/WarehouseVisual/Parcels/Parcel00").GetAttribute( + "xformOp:translate" + ) + assert tuple(position.Get()) == pytest.approx((1, 0, 1)) + env.common_step_counter = 180 + env.step(action) + assert tuple(position.Get()) == pytest.approx((1, 0, 1)) + finally: + env.close() + _presentation_layer.cache_clear() + assert not env._warehouse_animation 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..c016ec1562d7 --- /dev/null +++ b/source/isaaclab_tasks/test/contrib/test_conveyor_franka_geometry.py @@ -0,0 +1,91 @@ +# 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 + +import pytest + +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 + + +@pytest.mark.parametrize("warehouse", [False, True]) +def test_belt_top_faces_point_upward(warehouse): + """One-sided triangle-mesh surfaces support parcels from above.""" + for side in ("Left", "Right"): + if warehouse: + from isaaclab_tasks.contrib.conveyor_franka.conveyor_warehouse_geometry import ( + warehouse_belt_sections, + warehouse_guard_meshes, + ) + + specs = [ + section.geometry for section in warehouse_belt_sections(side) if isinstance(section.geometry, MeshSpec) + ] + specs.extend(warehouse_guard_meshes(side)) + else: + 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 + + +def test_elevated_belt_normals_accept_contacts_on_every_ramp_panel(): + """Inclined parcels receive traction instead of sliding into the bottom transition.""" + import numpy as np + + from isaaclab_tasks.contrib.conveyor_franka.conveyor_warehouse_geometry import warehouse_belt_sections + + for side in ("Left", "Right"): + for section in warehouse_belt_sections(side): + if section.belt.curved or abs(section.belt.direction[2]) < 1e-6: + continue + vertices = np.asarray(section.geometry.vertices) + triangles = vertices[np.asarray(section.geometry.faces)] + normals = np.cross(triangles[:, 1] - triangles[:, 0], triangles[:, 2] - triangles[:, 0]) + top = normals[normals[:, 2] > 1e-8] + top /= np.linalg.norm(top, axis=1, keepdims=True) + assert len(top) > 0 + assert np.all(top @ section.belt.surface_normal >= section.belt.contact_threshold) + assert abs(np.dot(section.belt.direction, section.belt.surface_normal)) < 1e-6 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..e9f3702a5f85 --- /dev/null +++ b/source/isaaclab_tasks/test/contrib/test_conveyor_franka_mdp.py @@ -0,0 +1,507 @@ +# 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 + +"""Conveyor policy and sorting 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.kinematics import end_effector_pose +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" +_SORTER = "IsaacContrib-Conveyor-Warehouse-Sorting-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) + + +@pytest.mark.parametrize("task", [_BASE, _SORTER]) +def test_checkpoint_interface_and_invalid_action_recovery(task): + """Both real tasks preserve learned feature ordering, finite rewards and independent reset flags.""" + with closing(_environment(task)) as env: + obs, _ = env.reset() + torch.testing.assert_close(env.scene["robot"].data.root_pos_w.torch[:, 2], env.scene.env_origins[:, 2]) + 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() + pool = getattr(env, "conveyor_cube_pool", None) + if pool is not None: + assert not env.curriculum_manager.active_terms, "Warehouse training must not use phase reset rows." + env.command_manager.get_term("transfer").has_target[:] = torch.tensor([False, True]) + env.step(torch.full((2, 8), 0.5)) + torch.testing.assert_close(env.action_manager.action, torch.full((2, 8), 0.5)) + env.command_manager.get_term("transfer").has_target[:] = torch.tensor([False, True]) + env.cfg.park_when_idle = True + env.step(torch.full((2, 8), 0.5)) + torch.testing.assert_close(env.action_manager.action[1], torch.full((8,), 0.5)) + assert env.action_manager.action[0, :7].abs().max() <= 0.25 + assert env.action_manager.action[0, 7] == 0 + 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]) + if pool is not None: + pool.slot_ids[:] = torch.tensor([[4, 1, 2, 3], [0, 5, 2, 3]]) + positions[0, 4] = torch.tensor([0.52, 0.27, 0.06]) + positions[1, 5] = torch.tensor([0.72, -0.27, 0.06]) + _place(env, positions) + if pool is not None: + env.sim.step(render=False) + model = env.sim.backend_manager.get_model() + contacts = env.sim.backend_manager.get_contacts() + count = int(contacts.rigid_contact_count.numpy()[0]) + assert count > 0 + bodies = model.shape_body.numpy() + first = bodies[contacts.rigid_contact_shape0.numpy()[:count]] + second = bodies[contacts.rigid_contact_shape1.numpy()[:count]] + assert ((first >= 0) | (second >= 0)).all() + _place(env, positions) + command = env.command_manager.get_term("transfer") + env.episode_length_buf[:] = torch.tensor([11, 19]) + command.set_goal(2) + if pool is not None: + command.has_target.fill_(True) + if pool is not None: + pose = pool.assets[4].data.root_pose_w.torch.clone() + pose[0, 3:] = torch.tensor([2**-0.5, 0, 0, 2**-0.5]) + velocity = torch.zeros(2, 6) + velocity[0] = torch.tensor([0.1, 0.2, 0.3, 0.4, 0.5, 0.6]) + pool.assets[4].write_root_pose_to_sim_index(root_pose=pose, skip_forward=True) + pool.assets[4].write_root_velocity_to_sim_index(root_velocity=velocity, skip_forward=True) + env.sim.forward() + env.scene.update(env.step_dt) + policy = env.observation_manager.compute()["policy"] + expected = positions[:, :4] if pool is None else positions[torch.arange(2)[:, None], pool.slot_ids] + 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(), + ) + if pool is not None: + torch.testing.assert_close(policy[0, 40:43], torch.tensor([0.0, -1.0, 0.0]), atol=1e-6, rtol=0) + torch.testing.assert_close(policy[0, 52:58], velocity[0]) + remote = positions.clone() + remote[:, 2:4] = torch.tensor([[2.8, 0.1, 0.56], [1.2, -0.59, 0.16]]) + _place(env, remote) + adapted = env.observation_manager.compute()["policy"][:, 16:28].reshape(2, 4, 3) + torch.testing.assert_close(adapted[:, 2:, 1:], torch.tensor([[[0.75, 0.06], [-0.75, 0.06]]] * 2)) + physical = ( + torch.stack([p.data.root_pos_w.torch for p in pool.assets], dim=1) - env.scene.env_origins[:, None] + ) + torch.testing.assert_close(physical, remote) + _place(env, positions) + assert command.held_cube_ids.tolist() == [-1, -1] + assert command.subgoal_start_steps.tolist() == [11, 19] + if pool is None: + 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() + # Interior joint positions make the learned 0.12-rad residual independently observable. + 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() + # Isolate an arm-only failure from a gripper-only failure. + 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() + if pool is not None: + command.has_target.zero_() + 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_sorter_dispatch_inventory_and_selected_resets(): + """Actual parcels are shuffled, individually dispatched, pinned while grasped, and left sorted.""" + from isaaclab_tasks.contrib.conveyor_franka.conveyor_warehouse_geometry import warehouse_parcel_positions + + with closing(_environment(_SORTER)) as env: + env.reset() + pool = env.conveyor_cube_pool + before_assignments = pool.assignment_counts.clone() + before_transfers = pool.transfer_counts.clone() + before_slots = pool.slot_ids[0].clone() + before_velocity = torch.stack([p.data.root_vel_w.torch.clone() for p in pool.assets], dim=1) + before = torch.stack([p.data.root_pose_w.torch.clone() for p in pool.assets], dim=1) + repeated = None + for ids in (torch.tensor([1]), slice(1, 2)): + pool.assignment_counts.fill_(7) + pool.transfer_counts.fill_(3) + velocity = pool.assets[-1].data.root_vel_w.torch.clone() + velocity[1].fill_(0.2) + pool.assets[-1].write_root_velocity_to_sim_index(root_velocity=velocity) + torch.manual_seed(42) + env._reset_idx(ids) + assert not pool.assets[-1].data.root_vel_w.torch[1].any() + torch.testing.assert_close( + pool.assignment_counts[1], torch.tensor([1] * CUBE_COUNT + [0] * (len(pool.assets) - CUBE_COUNT)) + ) + assert not pool.transfer_counts[1].any() + assert (pool.assignment_counts[0] == 7).all() and (pool.transfer_counts[0] == 3).all() + actual = torch.stack([p.data.root_pose_w.torch.clone() for p in pool.assets], dim=1) + torch.testing.assert_close(actual[0], before[0]) + torch.testing.assert_close(pool.slot_ids[0], before_slots) + torch.testing.assert_close(pool.slot_ids[1], torch.arange(4)) + torch.testing.assert_close( + torch.stack([p.data.root_vel_w.torch[0] for p in pool.assets]), before_velocity[0] + ) + local = actual[1, :, :3] - env.scene.env_origins[1] + destinations = torch.tensor(warehouse_parcel_positions()) + distance = torch.linalg.vector_norm(local[:, None] - destinations[None], dim=-1) + assert distance.min(1).values.max() < 1e-5 + assert distance.argmin(1).unique().numel() == len(pool.assets) + assert not torch.allclose(local, destinations) + if repeated is not None: + torch.testing.assert_close(actual, repeated) + repeated = actual.clone() + pool.assignment_counts.copy_(before_assignments) + pool.transfer_counts.copy_(before_transfers) + positions = _waiting_positions(env) + positions[:, 0] = torch.tensor([0.7, 0.27, 0.06]) # Correct class, so leave it alone. + positions[0, 2:4] = torch.tensor([[0.6, 0.27, 0.06], [0.6, -0.27, 0.06]]) + positions[0, 7] = torch.tensor([0.9, 0.27, 0.06]) + pool.assignment_counts[0, 7] = 100 # A repeatedly assigned, nearer arrival must not starve a new one. + positions[0, 5] = torch.tensor([0.8, 0.27, 0.06]) # Arrival outside the initial four policy slots. + _place(env, positions) + command = env.command_manager.get_term("transfer") + command.has_target.zero_() + command._update_command() + assert command.has_target.tolist() == [True, False] + slot = int(command.target_cube_ids[0]) + assert pool.slot_ids[0, slot] == 5 + assert command.command[0, -2:].tolist() == [0, 1] + assert pool.slot_ids[0, 0] == 0 + assert command.metrics["sorted_parcels"].tolist() == [3, 1] + torch.testing.assert_close(pool.assets[5].data.root_pos_w.torch[0] - env.scene.env_origins[0], positions[0, 5]) + # Simultaneous arrivals can select different slots and opposite transfer directions. + waiting = positions[1, 2].clone() + positions[1, 2] = torch.tensor([0.75, -0.27, 0.06]) + _place(env, positions) + command.has_target.zero_() + command._update_command() + assert command.target_cube_ids.tolist() == [slot, 2] + assert command.source_side_ids.tolist() == [0, 1] + positions[1, 2] = waiting + _place(env, positions) + command._update_command() + # Carry the assigned parcel above the workcell; a remote held target must retain its slot. + robot = env.scene["robot"] + fingers, _ = robot.find_joints("panda_finger_joint[1-2]", preserve_order=True) + joints = robot.data.joint_pos.torch.clone() + joints[:, fingers] = 0.019 + robot.write_joint_position_to_sim_index(position=joints) + arms, _ = robot.find_joints("panda_joint[1-7]", preserve_order=True) + joints[:, arms] = robot.data.default_joint_pos.torch[:, arms] + joints[:, arms[1]] = -0.8 + robot.write_joint_position_to_sim_index(position=joints) + env.sim.forward() + env.scene.update(env.step_dt) + tool, _ = end_effector_pose(env) + positions[0, 5] = tool[0] - env.scene.env_origins[0] + _place(env, positions) + assert positions[0, 5, 2] > 0.4 + command._update_command() + assert pool.slot_ids[0, slot] == 5 and command.has_target[0] + positions[0, 5] = torch.tensor([0.8, 0.27, 0.06]) + joints[:, fingers] = 0.04 + robot.write_joint_position_to_sim_index(position=joints) + _place(env, positions) + env.episode_length_buf += command.cfg.minimum_subgoal_steps + command.evaluate() + assert not command.pending_success[0] + positions[0, 5] = torch.tensor([0.8, -0.27, 0.06]) + joints[:, fingers] = 0.04 + robot.write_joint_position_to_sim_index(position=joints) + _place(env, positions) + env.episode_length_buf += command.cfg.minimum_subgoal_steps + for _ in range(command.cfg.hold_steps): + env.episode_length_buf += 1 + command.evaluate() + assert pool.transfer_counts[0, 5] == 1 and pool.transfer_counts[1].sum() == 0 + # Completed classes stay circulating and cannot keep paying a previous success reward. + positions[:, :, 0] = 0.6 + positions[:, :, 1] = torch.tensor([0.27, -0.27] * (len(pool.assets) // 2)) + positions[:, :, 2] = 0.06 + _place(env, positions) + command._update_command() + assert not command.has_target.any() + assert command.metrics["batch_complete"].tolist() == [1, 1] + command.new_success.fill_(True) + command.is_success.fill_(True) + command.evaluate() + assert not command.new_success.any() and not command.is_success.any() + assert not mdp.cube_out_of_workspace(env, **env.cfg.terminations.cube_out_of_workspace.params).any() + # Present every identity as a separate arrival and let the real dispatcher refill its slots. + for parcel_id in range(len(pool.assets)): + positions[0] = _waiting_positions(env)[0] + source_y = -0.27 if command.parcel_destinations[parcel_id] == 0 else 0.27 + positions[0, parcel_id] = torch.tensor([0.8, source_y, 0.06]) + _place(env, positions) + command.pending_success[0] = True + command._update_command() + assert pool.slot_ids[0, command.target_cube_ids[0]] == parcel_id + assert (pool.assignment_counts[0] > 0).all() + unassigned = next(i for i in range(len(pool.assets)) if i not in pool.slot_ids[0]) + positions[0, unassigned, 2] = -1 + _place(env, positions) + assert mdp.cube_out_of_workspace(env, **env.cfg.terminations.cube_out_of_workspace.params).tolist() == [ + True, + False, + ] + + +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]]