Skip to content

Add NotNextTo loss strategy and solver#732

Draft
xyao-nv wants to merge 2 commits into
mainfrom
xyao/dev/not_relation
Draft

Add NotNextTo loss strategy and solver#732
xyao-nv wants to merge 2 commits into
mainfrom
xyao/dev/not_relation

Conversation

@xyao-nv
Copy link
Copy Markdown
Collaborator

@xyao-nv xyao-nv commented May 28, 2026

Summary

Add NotNextTo loss strategy and solver

Detailed description

  • Why: when a task goal is NextTo(parent, side, distance) like Place A on +x side of B. The init state must NOT already satisfy it.

  • What: Added NotNextTo strategy, which is basically the negation of NexTO: breaking any one condition by at least margin_m meters drives the loss to zero. It's a local keep-out bubble around the one NextTo spot, than staying away from the entire whole side of the parent.

  • See below for cases where distance_m is the same between NextTo and NotNextTo

image

Regarding margin_m & distance_m value, I include a rule-of-thumb guidance in docstring.
image

Copy link
Copy Markdown
Contributor

@isaaclab-review-bot isaaclab-review-bot Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review Summary

This PR introduces NotOn and NotNextTo relation constraints—the geometric inverses of the existing On and NextTo relations. The implementation is well-designed and follows the established patterns in the codebase.

Verdict: APPROVE


Architecture & Design

The implementation correctly identifies the core insight: negated spatial relations can be satisfied by escaping along any single axis rather than satisfying all constraints simultaneously. The use of torch.minimum() to take the minimum across escape dimensions is mathematically sound and produces correct gradients for the optimizer.

Strengths

  1. Clean separation of concerns: Relation classes in relations.py are pure data containers; all loss computation logic lives in the strategy classes.

  2. Consistent API: The new strategies follow the exact same interface pattern as OnLossStrategy and NextToLossStrategy, accepting the same parameters and returning tensors in the same shape.

  3. Correct gradient behavior: The test test_not_on_gradient_pushes_outward() verifies that gradients push toward the nearer edge—this is the expected behavior for efficient optimization.

  4. Configurable margin: The margin_m parameter in both strategies allows fine-tuning the "escape distance" required before loss reaches zero. This is important for avoiding jitter at the boundary.


Potential Issues & Suggestions

1. Asymmetric margin_m defaults

NotOnLossStrategy defaults to margin_m=0.0, while NotNextToLossStrategy defaults to margin_m=0.05 and requires it to be positive (assert margin_m > 0.0).

This asymmetry may be intentional (different geometric semantics), but it could surprise users who expect consistent behavior. Consider adding a brief note in the docstring explaining why the default differs.

2. Missing Z-axis consideration in NotOn

The docstring states: "Z is ignored — with no XY overlap there is nothing to stack on."

This is correct for the stated use case, but there may be edge cases where an object is directly below the parent (same XY footprint, different Z). If the goal is to prevent any spatial overlap, Z should also be checked. If the goal is purely to prevent "resting on top," the current behavior is correct—but this assumption should be documented more explicitly for future maintainers.

3. Test coverage for batch inputs

The tests use single-position inputs (child_pos.dim() == 1). The strategies handle batched inputs (child_pos.dim() == 2), but this code path is not tested. Consider adding a parametrized test that verifies correct behavior with batched positions:

def test_not_on_batched_input():
    """Verify batched positions produce correct output shape."""
    child_pos = torch.tensor([[2.0, 0.4, 0.5], [0.4, 0.4, 0.5]])  # outside, inside
    loss = NotOnLossStrategy(slope=100.0).compute_loss(
        NotOn(_table()), child_pos, _box().bounding_box, _table().bounding_box
    )
    assert loss.shape == (2,)
    assert torch.isclose(loss[0], torch.tensor(0.0), atol=1e-6)  # outside
    assert loss[1] > 0.0  # inside

4. Solver integration test starts mug in wrong location

In test_solver_drives_mug_off_forbidden_table():

  • The mug starts at (1.9, 0.4, 0.13) which is already outside the right table's footprint ([1.5, 2.5] × [0, 1]).
  • The test verifies the final position, but the starting position already satisfies NotOn(right).

This test proves the solver converges, but doesn't verify the NotOn loss actively pushes the mug off. Consider starting the mug inside the forbidden table's footprint (e.g., (1.7, 0.4, 0.13)) to validate the repulsive behavior.


Minor Suggestions

  • Docstring typo (non-blocking): In NotNextToLossStrategy.__init__, the docstring says "margin_m: Meters of clearance required along whichever escape axis the optimizer takes." The word "clearance" could be misleading since this is the penalty zone width, not a physical clearance. Consider rephrasing to "safety margin around the forbidden zone."

Test Coverage Assessment

Test Area Coverage
Zero loss when constraint satisfied ✅ Both
Positive loss when violated ✅ Both
Gradient direction ✅ NotOn
Margin behavior ✅ NotOn
All Side variants ✅ NotNextTo
Solver integration
Batched input ❌ Not tested
Relation weight scaling ❌ Not tested

The test suite is comprehensive for the core functionality. The two gaps noted above are minor and can be addressed in a follow-up.


Final Notes

This is a clean, well-documented implementation that solves a real problem (preventing reset states from satisfying goal conditions). The mathematical approach is sound, the code follows established patterns, and the test coverage is good. Approving with the minor suggestions above for consideration.

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)
Copy link
Copy Markdown
Collaborator

@cvolkcvolk cvolkcvolk May 28, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggestion to move the all loss computation into loss_primitives.py into respective functions next to the others.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

done

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)
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same here, suggestion to move the loss computation into loss_primitives.py to line this up stylistically with the existing strategies.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

done

# 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]
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we could reuse some existing losses like the linear_band_loss() from loss_primitives.py in this function

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

done

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the refactor! 👍

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).
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This looks like the half-plane logic that single_boundary_linear_loss() implements

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

done

@xyao-nv xyao-nv changed the title Add NotOn, NotNextTo loss strategy and solver Add NotNextTo loss strategy and solver May 29, 2026
Copy link
Copy Markdown
Contributor

@isaaclab-review-bot isaaclab-review-bot Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review Summary

This PR adds NotNextTo — the logical inverse of NextTo — enabling exclusion zones where objects must not be placed beside a parent. The implementation is well-designed and follows the existing pattern established by NextTo and NextToLossStrategy.

✅ Strengths

  1. Clean architecture: The three-file structure (relation class, loss strategy, solver params) mirrors the existing NextTo pattern perfectly
  2. Solid mathematical foundation: The escape-based loss (min(gap_side, gap_cross, gap_dist)) correctly implements the logical negation — breaking ANY single condition satisfies NotNextTo
  3. Excellent documentation: Docstrings clearly explain the margin_m rule of thumb and the relationship between parameters
  4. Comprehensive test coverage: 8 new tests covering core functionality, all four sides, batched input, and solver integration

📝 Suggestions for Improvement

1. Consider torch.amin for cleaner multi-way minimum (Low Priority)

File: relation_loss_strategies.py, line 113

loss = self.slope * torch.minimum(torch.minimum(gap_side, gap_cross), gap_dist)

Consider using torch.amin(torch.stack([gap_side, gap_cross, gap_dist]), dim=0) for readability and extensibility if more escape routes are added in the future. The nested minimum is functionally correct but slightly harder to read.

2. Missing Z-axis Side variants (Clarification Needed)

File: test_relation_loss_strategies.py, line 319

The parametrized test covers POSITIVE_X, NEGATIVE_X, POSITIVE_Y, NEGATIVE_Y but not Z-axis variants. If Side includes POSITIVE_Z/NEGATIVE_Z, consider adding coverage. If these are intentionally excluded (2D planar constraints), this is fine as-is.

3. Potential saddle point at exact target position (Info)

File: relation_loss_strategies.py, line 107

The distance escape escape_dist = single_point_linear_loss(primary, target_pos, slope=1.0) is zero exactly at the target position, which is mathematically correct. However, at this point all three escapes are zero, making it a saddle point where the gradient direction is solely determined by numerical noise. The test at line 397 (test_solver_drives_box_out_of_forbidden_next_to_zone) wisely starts the box at 1.08 rather than exactly at the target to avoid this.

This is well-handled — just noting it for awareness.

4. Consider validating relation.distance_m against self.margin_m (Optional)

File: relation_loss_strategies.py, around line 65

The docstring mentions that margin_m >= distance_m is a rule of thumb. Consider adding a warning (not an assertion) when margin_m < relation.distance_m to help users avoid the "gap" issue described in the docs:

if self.debug and self.margin_m < relation.distance_m:
    print(f"Warning: margin_m ({self.margin_m}) < distance_m ({relation.distance_m}) may leave a gap")

🔍 Test Coverage Analysis

The test suite is solid. For completeness, consider adding:

  1. Boundary precision test: Position child exactly at margin_m distance from escape to verify loss transitions correctly
  2. Gradient verification test: Assert loss.backward() produces non-zero gradients in expected directions (ensures the optimizer can actually move the object)

✅ Overall Assessment

LGTM with minor suggestions. The implementation is correct, well-tested, and follows established patterns. The suggestions above are quality-of-life improvements, not blockers.


Reviewed at commit b6999b6

escape_dist = single_point_linear_loss(primary, target_pos, slope=1.0)
gap_dist = single_boundary_linear_loss(escape_dist, self.margin_m, slope=1.0, penalty_side="less")

# min(): a single escape past the margin is enough to satisfy Not(NextTo).
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This doesn't exist Not(NextTo) no?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants