Add NotNextTo loss strategy and solver#732
Conversation
There was a problem hiding this comment.
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
-
Clean separation of concerns: Relation classes in
relations.pyare pure data containers; all loss computation logic lives in the strategy classes. -
Consistent API: The new strategies follow the exact same interface pattern as
OnLossStrategyandNextToLossStrategy, accepting the same parameters and returning tensors in the same shape. -
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. -
Configurable margin: The
margin_mparameter 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 # inside4. 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) |
There was a problem hiding this comment.
Suggestion to move the all loss computation into loss_primitives.py into respective functions next to the others.
| 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) |
There was a problem hiding this comment.
Same here, suggestion to move the loss computation into loss_primitives.py to line this up stylistically with the existing strategies.
| # 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] |
There was a problem hiding this comment.
I think we could reuse some existing losses like the linear_band_loss() from loss_primitives.py in this function
There was a problem hiding this comment.
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). |
There was a problem hiding this comment.
This looks like the half-plane logic that single_boundary_linear_loss() implements
There was a problem hiding this comment.
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
- Clean architecture: The three-file structure (relation class, loss strategy, solver params) mirrors the existing
NextTopattern perfectly - 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 - Excellent documentation: Docstrings clearly explain the margin_m rule of thumb and the relationship between parameters
- 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:
- Boundary precision test: Position child exactly at
margin_mdistance from escape to verify loss transitions correctly - 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). |
There was a problem hiding this comment.
This doesn't exist Not(NextTo) no?
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
NotNextTostrategy, which is basically the negation ofNexTO: breaking any one condition by at leastmargin_mmeters 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
Regarding margin_m & distance_m value, I include a rule-of-thumb guidance in docstring.
