diff --git a/isaaclab_arena/relations/relation_loss_strategies.py b/isaaclab_arena/relations/relation_loss_strategies.py index af42a2e98..1d0fdddd9 100644 --- a/isaaclab_arena/relations/relation_loss_strategies.py +++ b/isaaclab_arena/relations/relation_loss_strategies.py @@ -18,7 +18,7 @@ from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox if TYPE_CHECKING: - from isaaclab_arena.relations.relations import AtPosition, NextTo, On, PositionLimits, Relation + from isaaclab_arena.relations.relations import AtPosition, NextTo, NotNextTo, NotOn, On, PositionLimits, Relation from isaaclab_arena.relations.relations import Side @@ -319,6 +319,184 @@ def compute_loss( return result.squeeze(0) if single_input else result +class NotOnLossStrategy(RelationLossStrategy): + """Loss strategy for ``NotOn`` — push child out of parent's XY footprint. + + The inverse of ``OnLossStrategy``. ``NotOn`` is satisfied as soon + as either the X or Y footprint has *escaped* the parent — Z is + incidental, since with no XY overlap there is nothing to stack on. + + Loss = ``slope * min(inside_x, inside_y)``, where ``inside_axis`` + is the child center's "depth inside" the placement band along that + axis (0 when outside, > 0 when inside). The ``min`` means a + single-axis escape is sufficient. The gradient on the smaller + inside-axis points toward the nearer edge — the optimizer pops + the child out along the path of least travel. + + If ``margin_m > 0``, the parent's effective footprint is widened + by ``margin_m`` along both XY axes, so the loss does not drop to + zero until the child has cleared the parent by that margin. + """ + + def __init__(self, slope: float = 100.0, margin_m: float = 0.0, debug: bool = False): + """ + Args: + slope: Loss magnitude per meter of inside-penetration. + Matches the default of ``OnLossStrategy``. + margin_m: Optional safety margin (meters) added to each XY + extent of the parent's footprint. The loss falls to + zero only when the child has cleared the widened + footprint. Default 0.0. + debug: If True, print the per-axis inside-penetration. + """ + assert slope >= 0.0, f"slope must be non-negative, got {slope}" + assert margin_m >= 0.0, f"margin_m must be non-negative, got {margin_m}" + self.slope = slope + self.margin_m = margin_m + self.debug = debug + + def compute_loss( + self, + relation: "NotOn", + child_pos: torch.Tensor, + child_bbox: AxisAlignedBoundingBox, + parent_world_bbox: AxisAlignedBoundingBox, + ) -> torch.Tensor: + """Compute loss for ``NotOn``.""" + single_input = child_pos.dim() == 1 + if single_input: + child_pos = child_pos.unsqueeze(0) + + # The On-valid band: child center positions for which the + # child's footprint is entirely inside the parent's footprint. + # Inflated by margin_m so Not(On) keeps pushing past the rim. + m = self.margin_m + valid_x_min = parent_world_bbox.min_point[:, 0] - child_bbox.min_point[:, 0] - m + valid_x_max = parent_world_bbox.max_point[:, 0] - child_bbox.max_point[:, 0] + m + valid_y_min = parent_world_bbox.min_point[:, 1] - child_bbox.min_point[:, 1] - m + valid_y_max = parent_world_bbox.max_point[:, 1] - child_bbox.max_point[:, 1] + m + + # Inside-band penetration: distance from child center to the + # nearer edge, clamped to >= 0. Zero when the child has + # escaped the band along that axis. + zero = torch.zeros((), dtype=child_pos.dtype, device=child_pos.device) + inside_x = torch.maximum(zero, torch.minimum(child_pos[:, 0] - valid_x_min, valid_x_max - child_pos[:, 0])) + inside_y = torch.maximum(zero, torch.minimum(child_pos[:, 1] - valid_y_min, valid_y_max - child_pos[:, 1])) + + # min(): a single-axis escape is enough to satisfy Not(On). + loss = self.slope * torch.minimum(inside_x, inside_y) + + if self.debug and child_pos.shape[0] == 1: + print( + f" [NotOn] inside_x={inside_x[0].item():.6f} inside_y={inside_y[0].item():.6f} " + f"-> loss={loss[0].item():.6f}" + ) + + result = relation.relation_loss_weight * loss + return result.squeeze(0) if single_input else result + + +class NotNextToLossStrategy(RelationLossStrategy): + """Loss strategy for ``NotNextTo`` — push child out of the NextTo zone on the given side. + + The inverse of ``NextToLossStrategy``. The NextTo "satisfied + region" is the conjunction of three conditions: + 1. **Half-plane:** child on the correct side of the parent edge. + 2. **Cross band:** child's perpendicular position inside the + parent's perpendicular extent. + 3. **Target distance:** primary-axis position equals + ``parent_edge + direction * distance_m``. + + ``NotNextTo`` is satisfied as soon as *any one* condition is + violated by at least ``margin_m`` meters. + + Loss = ``slope * min(relu(margin_m - escape_side), + relu(margin_m - escape_cross), relu(margin_m - escape_dist))``. + The gradient points along whichever escape is cheapest, so the + optimizer naturally escapes by the easiest of the three directions. + """ + + def __init__(self, slope: float = 10.0, margin_m: float = 0.05, debug: bool = False): + """ + Args: + slope: Loss magnitude per meter inside the safety margin. + Matches the default of ``NextToLossStrategy``. + margin_m: Meters of clearance required along whichever + escape axis the optimizer takes. Default 5 cm. + debug: If True, print the per-condition escape distances. + """ + assert slope >= 0.0, f"slope must be non-negative, got {slope}" + assert margin_m > 0.0, f"margin_m must be positive, got {margin_m}" + self.slope = slope + self.margin_m = margin_m + self.debug = debug + + def compute_loss( + self, + relation: "NotNextTo", + child_pos: torch.Tensor, + child_bbox: AxisAlignedBoundingBox, + parent_world_bbox: AxisAlignedBoundingBox, + ) -> torch.Tensor: + """Compute loss for ``NotNextTo``.""" + single_input = child_pos.dim() == 1 + if single_input: + child_pos = child_pos.unsqueeze(0) + + cfg = SIDE_CONFIGS[relation.side] + distance = relation.distance_m + + # Mirror NextToLossStrategy's target-position derivation. + if cfg.direction == Direction.POSITIVE: + parent_edge = parent_world_bbox.max_point[:, cfg.primary_axis] + child_offset = child_bbox.min_point[:, cfg.primary_axis] + else: + parent_edge = parent_world_bbox.min_point[:, cfg.primary_axis] + child_offset = child_bbox.max_point[:, cfg.primary_axis] + target_pos = parent_edge + cfg.direction * distance - child_offset + + # Cross band: child placed at target position within parent's perpendicular extent. + parent_band_min = parent_world_bbox.min_point[:, cfg.band_axis] + parent_band_max = parent_world_bbox.max_point[:, cfg.band_axis] + valid_band_min = parent_band_min - child_bbox.min_point[:, cfg.band_axis] + valid_band_max = parent_band_max - child_bbox.max_point[:, cfg.band_axis] + + primary = child_pos[:, cfg.primary_axis] + cross = child_pos[:, cfg.band_axis] + zero = torch.zeros((), dtype=child_pos.dtype, device=child_pos.device) + + # escape_side: how far on the WRONG side of parent's edge (0 if on correct side). + # direction = +1 means child should be > parent_edge; wrong-side amount = parent_edge - child. + escape_side = torch.maximum(zero, (parent_edge - primary) * cfg.direction) + + # escape_cross: how far OUTSIDE the perpendicular band (0 inside). + escape_cross = torch.maximum(zero, valid_band_min - cross) + torch.maximum(zero, cross - valid_band_max) + + # escape_dist: how far the primary-axis position is from the target distance (always >= 0). + escape_dist = torch.abs(primary - target_pos) + + # Per-condition "how much of the margin is unfilled". Zero once that escape passes margin. + margin = self.margin_m + gap_side = torch.maximum(zero, margin - escape_side) + gap_cross = torch.maximum(zero, margin - escape_cross) + gap_dist = torch.maximum(zero, margin - escape_dist) + + # min(): a single escape past the margin is enough to satisfy Not(NextTo). + loss = self.slope * torch.minimum(torch.minimum(gap_side, gap_cross), gap_dist) + + if self.debug and child_pos.shape[0] == 1: + print( + f" [NotNextTo] {relation.side.value}: " + f"escape_side={escape_side[0].item():.4f} " + f"escape_cross={escape_cross[0].item():.4f} " + f"escape_dist={escape_dist[0].item():.4f} " + f"-> loss={loss[0].item():.6f}" + ) + + result = relation.relation_loss_weight * loss + return result.squeeze(0) if single_input else result + + class NoCollisionLossStrategy: """Loss strategy for no-overlap constraints between objects. diff --git a/isaaclab_arena/relations/relation_solver_params.py b/isaaclab_arena/relations/relation_solver_params.py index 24a53d079..1e730b9ab 100644 --- a/isaaclab_arena/relations/relation_solver_params.py +++ b/isaaclab_arena/relations/relation_solver_params.py @@ -8,12 +8,14 @@ from isaaclab_arena.relations.relation_loss_strategies import ( AtPositionLossStrategy, NextToLossStrategy, + NotNextToLossStrategy, + NotOnLossStrategy, OnLossStrategy, PositionLimitsLossStrategy, RelationLossStrategy, UnaryRelationLossStrategy, ) -from isaaclab_arena.relations.relations import AtPosition, NextTo, On, PositionLimits, RelationBase +from isaaclab_arena.relations.relations import AtPosition, NextTo, NotNextTo, NotOn, On, PositionLimits, RelationBase def _default_strategies() -> dict[type[RelationBase], RelationLossStrategy | UnaryRelationLossStrategy]: @@ -21,6 +23,8 @@ def _default_strategies() -> dict[type[RelationBase], RelationLossStrategy | Una return { NextTo: NextToLossStrategy(slope=10.0), On: OnLossStrategy(slope=100.0), + NotOn: NotOnLossStrategy(slope=100.0), + NotNextTo: NotNextToLossStrategy(slope=10.0, margin_m=0.05), AtPosition: AtPositionLossStrategy(slope=100.0), PositionLimits: PositionLimitsLossStrategy(slope=100.0), } diff --git a/isaaclab_arena/relations/relations.py b/isaaclab_arena/relations/relations.py index c204c34ad..226d16322 100644 --- a/isaaclab_arena/relations/relations.py +++ b/isaaclab_arena/relations/relations.py @@ -126,6 +126,58 @@ def __init__( self.clearance_m = clearance_m +class NotOn(Relation): + """Forbids the child from sitting on the parent. + + The inverse of ``On``: penalizes any child position inside the + parent's XY footprint and contributes zero loss once the child has + cleared the footprint along X or Y. Z is ignored — with no XY + overlap there is nothing to stack on. + + Note: Loss computation is handled by NotOnLossStrategy in relation_loss_strategies.py. + """ + + def __init__(self, parent: ObjectBase, relation_loss_weight: float = 1.0): + """ + Args: + parent: The asset the child must NOT be on. + relation_loss_weight: Weight for the relationship loss function. + """ + super().__init__(parent, relation_loss_weight) + + +class NotNextTo(Relation): + """Forbids the child from sitting in the NextTo zone on the given side of the parent. + + The inverse of ``NextTo``: the child must be far from the target + position on the primary axis, *or* outside the perpendicular band, + *or* on the wrong side of the parent's edge. Any one escape is + sufficient for zero loss. + + Note: Loss computation is handled by NotNextToLossStrategy in relation_loss_strategies.py. + """ + + def __init__( + self, + parent: ObjectBase, + relation_loss_weight: float = 1.0, + distance_m: float = 0.05, + side: Side = Side.POSITIVE_X, + ): + """ + Args: + parent: The parent asset whose NextTo zone is forbidden. + relation_loss_weight: Weight for the relationship loss function. + distance_m: Target distance from parent's boundary in meters (default: 5cm). + Defines where the forbidden zone is centered along the primary axis. + side: Which axis direction (default: Side.POSITIVE_X). + """ + super().__init__(parent, relation_loss_weight) + assert distance_m > 0.0, f"Distance must be positive, got {distance_m}" + self.distance_m = distance_m + self.side = side + + class IsAnchor(RelationBase): """Marker indicating this object is an anchor for relation solving. diff --git a/isaaclab_arena/tests/test_not_relation.py b/isaaclab_arena/tests/test_not_relation.py new file mode 100644 index 000000000..2e51b0d77 --- /dev/null +++ b/isaaclab_arena/tests/test_not_relation.py @@ -0,0 +1,156 @@ +# Copyright (c) 2025-2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for NotOn / NotNextTo relations and their loss strategies.""" + +import torch + +import pytest + +from isaaclab_arena.assets.dummy_object import DummyObject +from isaaclab_arena.relations.relation_loss_strategies import NotNextToLossStrategy, NotOnLossStrategy +from isaaclab_arena.relations.relation_solver import RelationSolver +from isaaclab_arena.relations.relation_solver_params import RelationSolverParams +from isaaclab_arena.relations.relations import IsAnchor, NotNextTo, NotOn, On, Side +from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox +from isaaclab_arena.utils.pose import Pose + + +def _table(): + """1m x 1m x 0.1m table at the local origin (used as anchor at world (0, 0, 0)).""" + return DummyObject( + name="table", bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(1.0, 1.0, 0.1)) + ) + + +def _box(): + """0.2m cube.""" + return DummyObject( + name="box", bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.2)) + ) + + +# ============================================================================= +# NotOnLossStrategy +# ============================================================================= + + +@pytest.mark.parametrize("child_pos", [(2.0, 0.4, 0.5), (0.4, 2.0, 0.5)], ids=["outside_x", "outside_y"]) +def test_not_on_zero_when_child_outside_parent(child_pos): + """A single-axis escape (X or Y) is enough to drive NotOn loss to zero.""" + loss = NotOnLossStrategy(slope=100.0).compute_loss( + NotOn(_table()), torch.tensor(child_pos), _box().bounding_box, _table().bounding_box + ) + assert torch.isclose(loss, torch.tensor(0.0), atol=1e-6) + + +def test_not_on_positive_when_child_centered_on_parent(): + """At the center of the parent's footprint the loss is slope * min(inside_x, inside_y).""" + loss = NotOnLossStrategy(slope=100.0).compute_loss( + NotOn(_table()), torch.tensor([0.4, 0.4, 0.5]), _box().bounding_box, _table().bounding_box + ) + # valid_x in [0, 0.8], child_x=0.4 -> inside_x = 0.4; symmetric on Y. loss = 100 * 0.4 = 40. + assert torch.isclose(loss, torch.tensor(40.0), atol=1e-4) + + +def test_not_on_gradient_pushes_outward(): + """The gradient at an interior point pushes the child toward the nearer edge of the binding axis.""" + child_pos = torch.tensor([0.3, 0.4, 0.5], requires_grad=True) + # X is the binding axis (inside_x=0.3 < inside_y=0.4); descent direction is -X. + NotOnLossStrategy(slope=100.0).compute_loss( + NotOn(_table()), child_pos, _box().bounding_box, _table().bounding_box + ).backward() + assert child_pos.grad[0].item() > 0 + assert torch.isclose(child_pos.grad[1], torch.tensor(0.0), atol=1e-6) + + +def test_not_on_margin_extends_penalty_past_parent_edge(): + """margin_m widens the penalized region beyond the parent's footprint.""" + child_pos = torch.tensor([0.85, 0.4, 0.5]) # 5 cm past parent's X edge + loss_no_margin = NotOnLossStrategy(slope=100.0, margin_m=0.0).compute_loss( + NotOn(_table()), child_pos, _box().bounding_box, _table().bounding_box + ) + loss_with_margin = NotOnLossStrategy(slope=100.0, margin_m=0.1).compute_loss( + NotOn(_table()), child_pos, _box().bounding_box, _table().bounding_box + ) + assert torch.isclose(loss_no_margin, torch.tensor(0.0), atol=1e-6) + assert loss_with_margin > 0.0 + + +# ============================================================================= +# NotNextToLossStrategy +# ============================================================================= + + +def test_not_next_to_positive_at_target_position(): + """Loss is positive when the child sits at the NextTo zone's center (all escapes = 0).""" + relation = NotNextTo(_table(), side=Side.POSITIVE_X, distance_m=0.05) + loss = NotNextToLossStrategy(slope=10.0, margin_m=0.05).compute_loss( + relation, torch.tensor([1.05, 0.4, 0.5]), _box().bounding_box, _table().bounding_box + ) + # All three escapes = 0, all gaps = margin = 0.05. loss = slope * min(0.05, 0.05, 0.05) = 0.5. + assert torch.isclose(loss, torch.tensor(0.5), atol=1e-4) + + +@pytest.mark.parametrize( + "child_pos,reason", + [ + ((2.0, 0.4, 0.5), "far past target"), + ((0.3, 0.4, 0.5), "wrong side"), + ((1.05, 5.0, 0.5), "outside cross band"), + ], +) +def test_not_next_to_zero_when_any_escape_past_margin(child_pos, reason): + """Any single escape (distance / half-plane / cross-band) past margin -> loss = 0.""" + relation = NotNextTo(_table(), side=Side.POSITIVE_X, distance_m=0.05) + loss = NotNextToLossStrategy(slope=10.0, margin_m=0.05).compute_loss( + relation, torch.tensor(child_pos), _box().bounding_box, _table().bounding_box + ) + assert torch.isclose(loss, torch.tensor(0.0), atol=1e-6), f"Expected zero when {reason}" + + +@pytest.mark.parametrize( + "side,target", + [ + (Side.POSITIVE_X, (1.05, 0.4, 0.5)), + (Side.NEGATIVE_X, (-0.25, 0.4, 0.5)), + (Side.POSITIVE_Y, (0.4, 1.05, 0.5)), + (Side.NEGATIVE_Y, (0.4, -0.25, 0.5)), + ], +) +def test_not_next_to_each_side_variant_is_positive_at_its_target(side, target): + """All four Side variants produce a positive loss at their respective NextTo target.""" + relation = NotNextTo(_table(), side=side, distance_m=0.05) + loss = NotNextToLossStrategy(slope=10.0, margin_m=0.05).compute_loss( + relation, torch.tensor(target), _box().bounding_box, _table().bounding_box + ) + assert loss > 0.0 + + +# ============================================================================= +# Solver integration — proves the strategy is wired through end-to-end. +# ============================================================================= + + +def test_solver_drives_mug_off_forbidden_table(): + """Mug must sit on left_table but NOT on right_table. Starts on right, ends on left.""" + left, right = _table(), _table() + mug = DummyObject( + name="mug", bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.1, 0.1, 0.12)) + ) + left.set_initial_pose(Pose(position_xyz=(0.0, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) + right.set_initial_pose(Pose(position_xyz=(1.5, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) + left.add_relation(IsAnchor()) + right.add_relation(IsAnchor()) + mug.add_relation(On(left)) + mug.add_relation(NotOn(right)) + + initial = {left: (0.0, 0.0, 0.0), right: (1.5, 0.0, 0.0), mug: (1.9, 0.4, 0.13)} + solver = RelationSolver(params=RelationSolverParams(verbose=False, save_position_history=False, max_iters=400)) + final = solver.solve(objects=[left, right, mug], initial_positions=[initial])[0] + + mug_x, mug_y, _ = final[mug] + assert 0.0 <= mug_x <= 1.0 and 0.0 <= mug_y <= 1.0, f"Mug should be on the left table; got {final[mug]}" + assert not (1.5 <= mug_x <= 2.5), f"Mug should NOT be on the right table; got {final[mug]}"