code clean-up

Author: zhx06

isaaclab_arena/relations/layout_pool_serialization.py

@@ -114,7 +114,7 @@ def write_pool_document(path: Path, document: PoolDocument) -> None:
     path.parent.mkdir(parents=True, exist_ok=True)
     tmp_path = path.with_name(f"{path.name}.tmp")
     try:
-        tmp_path.write_text(payload)
+        tmp_path.write_text(payload, encoding="utf-8")
         os.replace(tmp_path, path)
     except OSError:
         tmp_path.unlink(missing_ok=True)
@@ -129,7 +129,7 @@ def read_pool_document(path: Path) -> PoolDocument:
     """
     assert path.is_file(), f"Layout pool file not found: {path}"
     try:
-        data = json.loads(path.read_text())
+        data = json.loads(path.read_text(encoding="utf-8"))
     except json.JSONDecodeError as exc:
         raise ValueError(f"Layout pool file is not valid JSON: {path}") from exc
     return PoolDocument.from_dict(data, path)

isaaclab_arena/relations/pooled_object_placer.py

@@ -483,6 +483,7 @@ def sample_with_replacement(self, count: int) -> list[PlacementResult]:
         # Reusable layouts are interchangeable: draw each slot from the flattened pool, with the
         # per-env RNG only selecting which stream the slot draws from.
         all_layouts = [layout for pool in self._env_pools for layout in pool.layouts]
+        assert all_layouts, "No accepted layouts to sample from."
         return [self._draw(self._env_rngs[i % self._num_envs], all_layouts) for i in range(count)]
 
     @staticmethod
@@ -525,6 +526,22 @@ def total_remaining(self) -> int:
         """Total unread layouts across all env pools."""
         return self._total_available()
 
+    @property
+    def stored_layouts(self) -> tuple[tuple[PlacementResult, ...], ...]:
+        """Read-only view of every stored layout, grouped by absolute env id.
+
+        Returns one inner tuple per env (outer length num_envs), each holding that env's stored
+        PlacementResults in pool order, including already-consumed ones. The grouping is meaningful
+        for env-specific layouts (see requires_env_indexed_layouts) and an arbitrary partition for
+        reusable ones, which can be flattened.
+
+        The tuples are a snapshot (later refills do not appear), but the PlacementResult objects are
+        live, so a post-pool check (e.g. a simulation collision test) records its outcome on a layout
+        in place via result.validation.with_check(...) and accepts() then reflects it. Resampling is
+        left to the caller via the existing sample_* methods.
+        """
+        return tuple(tuple(pooled.result for pooled in pool.layouts) for pool in self._env_pools)
+
     # ------------------------------------------------------------------
     # Persistence: save/load solved layouts to reuse poses without re-solving
     # ------------------------------------------------------------------
@@ -565,8 +582,9 @@ def load(
         env-count/heterogeneity/object-name mismatches fail loudly (see layout_pool_serialization
         for the structural checks). The saved placement_seed is restored (placer_params.placement_seed is
         ignored for seeding) so sampling matches the saved run; refill offset is not persisted, so
-        a refill restarts from the first solve batch. pool_size becomes the loaded layout count,
-        governing refill size.
+        a refill restarts from the first solve batch. pool_size becomes the total loaded layout count
+        across all envs (not per-env), so a multi-env refill batches larger than the original
+        per-batch pool_size; harmless because refills on a loaded pool are rare.
         """
         path = Path(path)
         document = read_pool_document(path)

isaaclab_arena/tests/test_layout_pool_serialization.py

@@ -5,6 +5,8 @@
 
 """Unit tests for the pure layout-pool JSON serialization helpers (no simulation)."""
 
+from pathlib import Path
+
 import pytest
 
 from isaaclab_arena.relations.layout_pool_serialization import (
@@ -165,6 +167,21 @@ def test_write_pool_document_rejects_non_finite_leaving_no_files(tmp_path):
     assert not path.with_name(f"{path.name}.tmp").exists()
 
 
+def test_write_pool_document_cleans_up_tmp_on_os_error(tmp_path, monkeypatch):
+    path = tmp_path / "pool.json"
+    document = PoolDocument.from_dict(_pool_dict(num_envs=1), path=path)
+
+    def fail_mid_write(self, *args, **kwargs):
+        self.write_bytes(b"partial")  # leave an orphan .tmp, then fail like a full disk
+        raise OSError("No space left on device")
+
+    monkeypatch.setattr(Path, "write_text", fail_mid_write)
+    with pytest.raises(OSError, match="No space left on device"):
+        write_pool_document(path, document)
+    assert not path.exists()
+    assert not path.with_name(f"{path.name}.tmp").exists()
+
+
 def test_read_pool_document_missing_file(tmp_path):
     with pytest.raises(AssertionError, match="not found"):
         read_pool_document(tmp_path / "missing.json")

isaaclab_arena/tests/test_object_placer_reproducibility.py

@@ -477,6 +477,39 @@ def test_pooled_placer_homogeneous_stored_layouts_have_distinct_positions_dicts(
             ), f"Layouts {i} and {j} share the same positions dict reference"
 
 
+def test_pooled_placer_stored_layouts_groups_live_results_by_env():
+    """stored_layouts must return a per-env snapshot of the live stored PlacementResults."""
+    num_envs = 3
+    solver_params = RelationSolverParams(max_iters=50)
+    placer_params = ObjectPlacerParams(placement_seed=42, solver_params=solver_params, apply_positions_to_objects=False)
+
+    pool = PooledObjectPlacer(
+        objects=list(_create_test_objects()), placer_params=placer_params, pool_size=12, num_envs=num_envs
+    )
+
+    grouped = pool.stored_layouts
+    assert len(grouped) == num_envs
+    assert all(isinstance(env_layouts, tuple) for env_layouts in grouped)
+    # Each returned result is the same live object the pool stores (no copy), grouped per env.
+    for cur_env, env_layouts in enumerate(grouped):
+        assert list(env_layouts) == [pooled.result for pooled in pool._env_pools[cur_env].layouts]
+
+
+def test_pooled_placer_stored_layouts_post_validation_flows_into_accepts():
+    """Enriching a stored layout's report in place (the post-validation use case) must change accepts()."""
+    solver_params = RelationSolverParams(max_iters=50)
+    placer_params = ObjectPlacerParams(placement_seed=42, solver_params=solver_params, apply_positions_to_objects=False)
+
+    pool = PooledObjectPlacer(objects=list(_create_test_objects()), placer_params=placer_params, pool_size=4)
+    result = pool.stored_layouts[0][0]
+    assert pool.accepts(result)
+
+    # Record a failing post-pool check (e.g. a simulation collision test) on the live layout.
+    result.validation = result.validation.with_check("sim_collision_free", False)
+    assert "sim_collision_free" in result.validation.failed_checks
+    assert not pool.accepts(result)
+
+
 def test_pooled_placer_homogeneous_sample_without_replacement_count_exceeds_pool_size():
     """sample_without_replacement(count) where count > pool_size must solve a larger batch in one shot."""
     solver_params = RelationSolverParams(max_iters=50)
@@ -714,10 +747,7 @@ def test_pooled_placer_save_load_round_trip_preserves_layouts(tmp_path):
 
 
 def test_pooled_placer_save_load_round_trip_multi_env_homogeneous(tmp_path):
-    """A multi-env homogeneous pool's per-env layouts survive the round trip intact.
-
-    Env-specific (variant) correctness is covered by the heterogeneous round-trip test.
-    """
+    """A multi-env homogeneous pool's per-env layouts survive the round trip (variants covered separately)."""
     num_envs = 3
     placer_params = _make_seeded_params()
     pool = PooledObjectPlacer(