Batch Support for Relation Solver by zhx06 · Pull Request #512 · isaac-sim/IsaacLab-Arena

zhx06 · 2026-03-30T21:28:34Z

Summary

Add batched multi-env placement to relation solver (1 -> num_envs)

Detailed description

Solver

RelationSolverState positions stored as (num_envs, num_optimizable, 3); get_position() returns (num_envs, 3)
RelationSolver._compute_total_loss() accumulates (num_envs,) loss per env, returns mean of losses; per-env loss exposed via last_loss_per_env
All loss strategies (On, NextTo, NoCollision, AtPosition) updated to accept (num_envs, 3) with (3,) backward compat

ObjectPlacer

place() accepts num_envs and result_per_env; returns MultiEnvPlacementResult (one PlacementResult per env) or a single PlacementResult
Pool-based placement: generates max_placement_attempts * num_results candidates in one batched solver call, sorts by (valid-first, lowest-loss), and selects the best num_results
result_per_env=False solves a single layout applied to all environments
_validate_placement skips On-related pairs in overlap check (fix for false-positive collisions)
mean_loss calculation filters out non-finite values

Unified Placement Path

New PosePerEnv dataclass holds a list of Pose objects, one per environment
ObjectBase._init_event_cfg handles PosePerEnv via set_object_pose_per_env, so each object registers its own reset event for both single-env and multi-env cases
Removed placement_events.py — its functionality is now handled by object-level event registration
ArenaEnvBuilder._solve_relations() simplified: no longer builds a separate placement event cfg or distinguishes single/multi-env code paths

Example Update

Update default objects in gr1_table_multi_object_no_collision_environment, add TODO for potential object collision issues after running the solver

alexmillane

Good work.

I have some suggestions for how we might improve things.

alexmillane · 2026-03-31T16:39:41Z

isaaclab_arena/environments/arena_env_builder.py

+            positions_all_envs_by_name = [
+                {obj.name: result.results[e].positions[obj] for obj in result.results[0].positions}
+                for e in range(len(result.results))
+            ]


Wdyt about moving this into MultiEnvPlacementResult? It's a bit difficult to decipher. Perhaps would be nicers to hide inside a method?

Done, currently it does not build layout dicts

alexmillane · 2026-03-31T16:57:12Z

isaaclab_arena/environments/arena_env_builder.py

+        num_envs = self.args.num_envs
+        result = placer.place(objects_with_relations, num_envs=num_envs)
+        if isinstance(result, MultiEnvPlacementResult) and result.results:
+            positions_all_envs_by_name = [
+                {obj.name: result.results[e].positions[obj] for obj in result.results[0].positions}
+                for e in range(len(result.results))
+            ]
+            object_names = [obj.name for obj in objects_with_relations]
+            anchor_names = [a.name for a in get_anchor_objects(objects_with_relations)]
+            self._placement_event_cfg = make_placement_event_cfg(
+                positions_all_envs_by_name,
+                object_names,
+                anchor_names,
+            )
+        else:
+            self._placement_event_cfg = None


I can see what the motivation is here, but we have two quite different code paths for the single environment and multi-environment case.

I think we should unify these code paths.

Both code paths (single and multi placement) should run the solver to produce a result (of type PlacementResult or MultiEnvPlacementResult).

This result should be passed to ObjectPlacer._apply_positions.

The placement should be passed to the object through it's Object.set_initial_pose, which will have to be expanded to take something like a PosePerEnv or something like that.

Right now, depending on whether your optimizing for a single or multiple environments results in two very different paths. In the single env case, the responsibility for applying the placement is in the ObjectPlacer and Object registers the event. In the multi-env case the placement is effectively applied in the ArenaEnvBuilder, which also registers the event.

I notice you have a note somewhere that perhaps these options could be unified. I suggest we unify now.

We unified these code paths in the current code. Each object's _init_event_cfg generates the appropriate reset event.

alexmillane · 2026-03-31T17:00:16Z

isaaclab_arena/environments/arena_env_builder.py

            task.get_events_cfg(),
-        )
+        ]
+        placement_event = getattr(self, "_placement_event_cfg", None)


Prefer to set this in the constructor (__init__) as None unconditionally, such that the object is garunteed to have the member, rather than checking for the attribute every time

This is removed in the new version of code, we no longer have placement_event

alexmillane · 2026-03-31T17:01:04Z

isaaclab_arena/environments/arena_env_builder.py

+        placement_event = getattr(self, "_placement_event_cfg", None)
+        if placement_event is not None:
+            events_sources.append(placement_event)
+        events_cfg = combine_configclass_instances("EventsCfg", *events_sources)


See comment above. I think that the placement result should be registered through the objects, as it is the single env case.

Done. Both single-env and multi-env paths now register placement in object levels.

alexmillane · 2026-03-31T17:03:51Z

isaaclab_arena/relations/object_placer.py

        self,
        objects: list[Object | ObjectReference],
-    ) -> PlacementResult:
+        num_envs: int = 1,


What do you think about the case that we have num_envs > 1, but we want a placement that is the same across all environments. I feel that that might be a capability we'd like to preserve. Wdyt @cvolkcvolk?

If we do want to preserve that functionality we should introduce another parameter result_per_env: bool = True

Is see you've solved this! Thanks!

alexmillane · 2026-03-31T17:13:21Z

isaaclab_arena/relations/object_placer.py

+        # Per env: use best valid if any, else best-by-loss fallback
+        final_per_env: list[dict] = [
+            (
+                best_valid_positions_per_env[e]
+                if best_valid_positions_per_env[e] is not None
+                else best_any_positions_per_env[e]
+            )
+            for e in range(num_envs)
+        ]


The idea here is that we return some result, even if some of the environments do not result in a valid placement?

Do you think that that's a safe approach.

This one no longer exists when we have the pool-based selection.

alexmillane · 2026-03-31T17:13:48Z

isaaclab_arena/relations/object_placer.py

-            final_loss=best_loss,
-            attempts=attempt + 1,
-        )
+        # TODO(@zhx06): Consider applying via event for consistency with multi_env.


I think we should do this in this MR. See comments above.

It is now implemented in this MR.

alexmillane · 2026-03-31T17:18:15Z

isaaclab_arena/relations/object_placer.py

+                valid = self._validate_placement(positions_per_env[e])
+                if valid and loss_e < best_valid_loss_per_env[e]:
+                    best_valid_loss_per_env[e] = loss_e
+                    best_valid_positions_per_env[e] = positions_per_env[e]
+                if loss_e < best_any_loss_per_env[e]:
+                    best_any_loss_per_env[e] = loss_e
+                    best_any_positions_per_env[e] = positions_per_env[e]


One thing that is inefficient about this loop is that it if a single environment (say "A)) fails N times, by chance, the whole thing fails. Even if environment "B" has 10 successful solutions. Because the environments are homogenous, "A" could just use the solution from "B".

I wonder if we shouldn't just calculate FACTOR * num_env solutions, and take the best num_envs. There's no need to associate a particular solution to a particular environment (because they're homogeneous).

Done. We now generate max_placement_attempts * num_envs candidates in a single batched solver call

alexmillane · 2026-03-31T17:20:32Z

isaaclab_arena/relations/placement_events.py

+        if not hasattr(asset, "write_root_pose_to_sim"):
+            continue


under which conditions would an object in out list not have this method? Is that something we want to handle?

alexmillane · 2026-03-31T17:20:55Z

isaaclab_arena/relations/placement_events.py

+    ordered_names = [n for n in object_names if n in anchor_set]
+    ordered_names += [n for n in object_names if n not in anchor_set]


Why do the names have to be ordered. Consider writing a comment.

This part is outdated and is replaced by the pool-based approach

alexmillane

Thank you for addressing the last comments. It's look a lot better (and simpler) now. Great!

I just have some smaller stylistic comments now.

alexmillane · 2026-04-01T15:45:51Z

isaaclab_arena/environments/arena_env_builder.py

+        events_sources = [
            embodiment.get_events_cfg(),
            self.arena_env.scene.get_events_cfg(),
            task.get_events_cfg(),
-        )
+        ]
+        events_cfg = combine_configclass_instances("EventsCfg", *events_sources)


Suggestion just to revert this change (it has no effect).

Done. This is reverted

alexmillane · 2026-04-01T15:46:23Z

isaaclab_arena/relations/object_placer.py

        self,
        objects: list[Object | ObjectReference],
-    ) -> PlacementResult:
+        num_envs: int = 1,


Is see you've solved this! Thanks!

alexmillane · 2026-04-01T15:49:28Z

isaaclab_arena/relations/object_placer.py

+            rng_state = None
+            if self.params.placement_seed is not None:
+                rng_state = torch.get_rng_state()
+                torch.manual_seed(self.params.placement_seed + candidate_idx)
+            initial_positions.append(self._generate_initial_positions(objects, anchor_objects_set, init_bounds))
+            if rng_state is not None:
+                torch.set_rng_state(rng_state)


Suggestion to abstract this logic into a function. Maybe _generate_initial_positions could be expanded to handle the single and multiple placement cases.

Now the new _generate_initial_positions handles both cases.

alexmillane · 2026-04-01T15:52:00Z

isaaclab_arena/relations/object_placer.py

+        all_losses = (
+            self._solver.last_loss_per_env.cpu().tolist()
+            if self._solver.last_loss_per_env is not None
+            else [float("inf")] * num_candidates
+        )


Why is this needed? Seems like if we've run the solve, the last_loss_per_env should not be None or am I missing something?

This one is removed from the updated code.

alexmillane · 2026-04-01T15:52:40Z

isaaclab_arena/relations/object_placer.py

-            torch.set_rng_state(rng_state)
+        all_candidates: list[tuple[float, dict, bool]] = []
+        for idx in range(num_candidates):
+            loss = all_losses[idx] if idx < len(all_losses) else float("inf")


Do we need this check: idx < len(all_losses). Shouldn't all_losses be defined to have the correct length?

Correct, This extra check is removed.

alexmillane · 2026-04-01T16:00:07Z

isaaclab_arena/relations/object_placer.py

+        n_valid = sum(1 for c in selected if c[2])
+        if self.params.verbose:
+            total_valid = sum(1 for c in all_candidates if c[2])
+            finite_losses = [c[0] for c in all_candidates if math.isfinite(c[0])]
+            mean_loss = sum(finite_losses) / len(finite_losses) if finite_losses else float("inf")
+            print(
+                f"Solved {num_candidates} candidates in one batch: mean loss = {mean_loss:.6f},"
+                f" {total_valid} valid, selected best {num_results} ({n_valid} valid)"
+            )

-        return PlacementResult(
-            success=success,
-            positions=best_positions,
-            final_loss=best_loss,
-            attempts=attempt + 1,
-        )
+        final_per_env: list[dict] = [c[1] for c in selected]
+        results_per_env = [
+            PlacementResult(
+                success=c[2],
+                positions=c[1],
+                final_loss=c[0],
+                attempts=self.params.max_placement_attempts,
+            )
+            for c in selected
+        ]


See comment above. This accesses of c make things a bit difficult to read. Consider using a named composite.

alexmillane · 2026-04-01T16:08:09Z

isaaclab_arena/relations/object_placer.py

+                if (id(a), id(b)) in on_pairs:
+                    continue


Add:

# Pairs related by an OnRelation are excluded from the check.

Done. This comment has been added.

alexmillane · 2026-04-01T16:18:46Z

isaaclab_arena/relations/object_placer.py

        """
-        for obj, pos in positions.items():
+        num_envs = len(positions_per_env)
+        for obj in positions_per_env[0]:


This looks a bit weird because of the [0]. Could be clearer to split into two lines:

# Objects are the same for every environment. Extract them. objects = [obj for obj in positions_per_env[0]] # Apply pose for each object. for obj in objects: ...

The code has been updated following your suggestion.

alexmillane · 2026-04-01T16:46:16Z

isaaclab_arena/relations/relation_solver.py

+        initial_positions: (
+            dict[Object | ObjectReference, tuple[float, float, float]]
+            | list[dict[Object | ObjectReference, tuple[float, float, float]]]
+        ),
+    ) -> (
+        dict[Object | ObjectReference, tuple[float, float, float]]
+        | list[dict[Object | ObjectReference, tuple[float, float, float]]]


These types are getting a bit awkward.

One way we can simply is: Object | ObjectReference -> ObjectBase (their parent class).

Do you see any opportunity to just use a list of length 1 for the single env case?

ll type annotations now use ObjectBase.

solve() and RelationSolverState now always use list[dict].

alexmillane · 2026-04-01T16:46:50Z

isaaclab_arena/relations/relation_solver_state.py

+        if isinstance(initial_positions, dict):
+            initial_positions = [initial_positions]


Seems like maybe we should just use a list no matter what? What do you think?

This is removed and it takes list[dict] for all cases.

zhx06 added 3 commits March 30, 2026 14:10

batch support for solver based on batched bbox

72ed7d1

fix overlap validation, revert hyperparameters

142af71

pre-commit check

7c0079c

zhx06 changed the title ~~Zxiao/solver batch support from batch bbox~~ Zxiao/batch support for relation solver Mar 30, 2026

zhx06 changed the title ~~Zxiao/batch support for relation solver~~ Batch Support for Relation Solver Mar 30, 2026

zhx06 added 3 commits March 30, 2026 15:21

rename env_id related function

debb6bf

add assert and type hints

c433c1c

potential bug fix

ab1bff2

alexmillane reviewed Mar 31, 2026

View reviewed changes

zhx06 added 4 commits March 31, 2026 10:25

pre commit fix

2b1e0a4

address comments

459a745

pre commit check

1680357

add comments for Object Placer

8c7f5a3

alexmillane reviewed Apr 1, 2026

View reviewed changes

zhx06 added 2 commits April 1, 2026 11:08

simplify types to ObjectBase, always use list

2b9f06d

update dataclass name

68784da

		ordered_names = [n for n in object_names if n in anchor_set]
		ordered_names += [n for n in object_names if n not in anchor_set]

		if isinstance(initial_positions, dict):
		initial_positions = [initial_positions]

Conversation

zhx06 commented Mar 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Detailed description

Solver

ObjectPlacer

Unified Placement Path

Example Update

Uh oh!

alexmillane left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zhx06 Apr 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

alexmillane left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zhx06 commented Mar 30, 2026 •

edited

Loading

zhx06 Apr 1, 2026 •

edited

Loading