Resolve initial state specs to full state specs in env graph#753
Conversation
There was a problem hiding this comment.
🤖 Isaac Lab Review Bot — PR #753
Summary: This PR introduces a StateSpecResolver that chains per-task success transitions into intermediate/final states from a partially-populated env graph. The architecture is clean and extensible.
🔴 Issues (3)
1. Missing blank line before @configclass (lint failure likely)
File: isaaclab_arena/tasks/pick_and_place_task.py
The new success_state_transition method ends immediately before @configclass class SceneCfg: with no separating blank line. This will likely trip PEP 8 / pre-commit checks (E302: expected 2 blank lines).
2. _relocations helper will break when SetState effects are used
File: isaaclab_arena/agentic_environment_generation/state_spec_resolver.py (lines ~115-120)
def _relocations(task: UnresolvedArenaEnvGraphTaskSpec) -> list[Relocate]:
transition = _transition_for(task)
for effect in transition.effects:
if not isinstance(effect, Relocate):
raise NotImplementedError(f"Effect {effect} is not yet supported.")
return list(transition.effects)SetState is defined and exported but any task that uses it will crash here. Consider filtering to Relocate effects only and handling (or ignoring) SetState separately, or at least document this as intentional until SetState is supported.
3. assert used for runtime validation in production code
Files: state_spec_resolver.py, arena_env_graph_spec.py
Multiple assert statements guard runtime invariants (e.g., assert states_in, assert ids, assert old_id.startswith(old_prefix)). These are stripped with python -O. For production validation, prefer explicit if not ...: raise ValueError(...).
🟡 Suggestions (3)
4. embodiment_id.split('_')[0] is fragile
File: state_spec_resolver.py (line ~164)
The reach-constraint id uses embodiment_id.split('_')[0] as a prefix. This relies on a naming convention (e.g. droid_abs_joint_pos → droid). If an embodiment id doesn't follow this pattern, the generated id will be wrong. Consider using the full embodiment id or a dedicated short-name field.
5. Multiple embodiments silently use the first
File: state_spec_resolver.py (_embodiment_id)
If multiple embodiment nodes exist, the resolver silently picks the first. Either assert exactly one or accept the embodiment id as a parameter.
6. Test mutation of resolver internals
File: tests/test_state_spec_resolver.py (test_task_without_a_transition_is_rejected)
Directly mutating resolver.graph.tasks[0].type works but is brittle. A dedicated YAML fixture with an unsupported task type would be more robust.
✅ Strengths
- Clear data-flow pipeline: YAML →
UnresolvedArenaEnvGraphSpec→StateSpecResolver.resolve_path()→ArenaEnvGraphSpec TaskTransition/Relocate/SetStateeffect algebra is well-designed and extensibleRelationBase.is_spawn_pose_constraint()as single source of truth is clean- Good test coverage: golden-fixture comparison, chain-wiring assertion, rejection tests
to_dict/write_yamlenable round-trip serialization
Verdict: Approve-with-nits — the architecture is sound. Fix the blank-line lint issue (blocking CI), and consider addressing the assert → explicit-raise and SetState handling before merge.
🏷️ Findings: 6 | Blocking: 1 (lint) | Non-blocking: 5
There was a problem hiding this comment.
Review Summary
This PR introduces a StateSpecResolver that chains task success conditions into intermediate/final state specs from a single initial state. The architecture is clean — TaskTransition as the single source of truth per task type, the resolver as a deterministic chain builder, and golden-fixture testing for correctness. Overall a solid design that will scale well.
Below are findings ordered by severity.
🔴 Critical
1. Missing blank line before top-level @configclass in pick_and_place_task.py
@configclass in pick_and_place_task.py✅ Fixed in 94c1dce.
🟡 Medium
2. _relocations() raises NotImplementedError for SetState effects — fragile for future use
File: isaaclab_arena/agentic_environment_generation/state_spec_resolver.py (lines 117-121)
for effect in transition.effects:
if not isinstance(effect, Relocate):
raise NotImplementedError(f"Effect {effect} is not yet supported.")SetState is defined and exported from task_transition.py, but any task returning it will crash the resolver. Since SetState is part of the public API of this PR, consider either:
- Handling it (skip gracefully since it doesn't change spatial constraints), or
- Adding an explicit
isinstance(effect, SetState)check with a descriptive skip/warning instead of a genericNotImplementedError.
This prevents a confusing error message when the next developer adds a SetState effect to a task.
3. env_name discrepancy between init and full fixture may confuse future maintainers
File: isaaclab_arena/tests/test_data/pick_and_place_maple_table_init_env_graph.yaml
The init graph uses env_name: pick_and_place_maple_table_init while the golden full graph uses env_name: pick_and_place_maple_table_default. The test's resolve_path() call (without an env_name argument) will produce a spec with env_name = "pick_and_place_maple_table_init", but the test only compares state_specs and tasks — so it passes. However, if someone later adds an env_name assertion or uses resolve_path() output for round-trip YAML comparison, this will silently differ.
Suggestion: Either make the init fixture use env_name: pick_and_place_maple_table_default, or have the test explicitly pass env_name="pick_and_place_maple_table_default" to resolve_path() to signal intent.
4. _transition_for instantiates TaskRegistry() on every call — redundant function call overhead
File: isaaclab_arena/agentic_environment_generation/state_spec_resolver.py (line 112)
def _transition_for(task: UnresolvedArenaEnvGraphTaskSpec) -> TaskTransition:
task_cls = TaskRegistry().get_task_by_name(task.type)
return task_cls.success_state_transition(task.task_args)While TaskRegistry uses SingletonMeta (so no actual re-registration), calling _transition_for → _relocations → _transition_for again for the same task (lines 87 and 113) means the registry lookup + success_state_transition is called twice per task. Consider refactoring to call _transition_for once per task in the loop and pass the result to both the reach-target logic and _relocations.
🟢 Low / Nits
5. _relocations calls _transition_for redundantly
File: isaaclab_arena/agentic_environment_generation/state_spec_resolver.py (lines 87 vs 113)
This double-resolution is harmless but wasteful. A small refactor:
transition = _transition_for(task)
relocations = [e for e in transition.effects if isinstance(e, Relocate)]…would eliminate the redundancy and make the code easier to follow.
6. reach_target_on_success property could return None for effect-less tasks in terminal position
File: isaaclab_arena/tasks/task_transition.py (lines 57-59)
For tasks with subject=None and no effects, this returns None. The resolver now handles this correctly via the if reach_target is not None guard in the deduplication loop. Consider adding a brief docstring note about this edge case.
7. Test test_task_without_a_transition_is_rejected mutates shared state
File: isaaclab_arena/tests/test_state_spec_resolver.py (line 71)
Consider using copy.deepcopy or constructing a minimal inline fixture.
8. Consider adding __all__ to task_transition.py
✅ Positives
- The golden-fixture testing strategy (
test_populate_reproduces_golden_full_graph) is excellent — deterministic output pinned against a hand-authored YAML. TaskTransitionas a declarative, classmethod-based contract is a clean design that separates concerns well.RelationBase.is_spawn_pose_constraint()as a single source of truth for which constraints carry forward is the right abstraction.- Good use of
check_task_wiring=Falseto parameterize existing validation rather than duplicating it.
📝 Update (94c1dce)
Reviewed incremental changes from c0b45d8 → 94c1dce. This is a high-quality refactor that migrates the resolver from raw dict manipulation to typed Pydantic model operations (model_copy, direct constructors). Key improvements:
Previous concerns status:
- 🔴 #1 (missing blank line before
@configclass): ✅ Fixed. - 🟡 #2–#4, 🟢 #5–#8: Not addressed in this push (non-blocking).
Semantic change — multiple reach targets per state: Interior success states now correctly carry two reach constraints (postcondition of the completed task + precondition of the next), rather than choosing one. This is semantically more correct — a success state must assert both that the embodiment reached the target of the finished task and that it can reach the next task's subject. The golden fixture is updated accordingly. The logic is cleanly structured with the seen_targets deduplication set.
New observations:
- The extensive docstring with worked example in
_successor_stateis a significant readability improvement. 👍 - Using
model_copy(update=...)instead of dict construction is safer (schema-validated) and more maintainable. - The
reach_targets: list[str | None]parameter + dedup pattern handles edge cases well. - No new issues found in this push.
Verdict: Approve-ready (remaining items are non-blocking nits).
Update (3b471e3): Reviewed incremental changes from 94c1dcee → 3b471e37. This is a significant refactor: the resolver now uses typed imports and cleaner function names (_get_next_state_spec, _get_task_relocations, etc.), and the PR adds CLI override specs, chunked eval dispatch, and validation context for unresolved graphs.
Previous concerns status:
- 🟡 #4/#5 (redundant
TaskRegistrylookups / double_transition_forcall): ✅ Fixed. Transitions are now pre-computed once viatransitions = [_get_task_state_transition(task) for task in spec.tasks]. - 🟡 #2 (
NotImplementedErrorforSetState): Not addressed (non-blocking). - 🟡 #3 (env_name discrepancy): Not addressed (non-blocking).
- 🟢 #6–#8: Not addressed (non-blocking nits).
New issue:
🔴 Keyword argument mismatch in _get_next_state_spec call (runtime TypeError)
state_spec_resolver.py line ~72 calls _get_next_state_spec(prev_state=states[-1], ...) but the function signature (line ~115) declares the parameter as prev_state_spec. This will raise TypeError: _get_next_state_spec() got an unexpected keyword argument 'prev_state' at runtime when resolve_constraints() is invoked. Fix: rename either the call-site kwarg to prev_state_spec= or the parameter to prev_state.
Verdict: One new critical bug (naming mismatch) introduced; otherwise a clean, well-structured refactor.
Update (c5ec371): Reviewed incremental changes from 3b471e37 → c5ec371c.
Previous concerns status:
- 🔴 Keyword argument mismatch (
prev_statevsprev_state_spec): ✅ Fixed. Line 72 now correctly usesprev_state_spec=states[-1]. - 🟡 #2 (
NotImplementedErrorforSetState): Not addressed (non-blocking). - 🟡 #3 (env_name discrepancy): Not addressed (non-blocking).
- 🟢 #6–#8: Not addressed (non-blocking nits).
New changes in this push:
- Migrated
arena_env_graph_types.pyandarena_env_graph_spec.pyfrom dataclasses to PydanticBaseModelwith validators — cleaner validation and serialization. - Removed
ArenaEnvGraphSpatialConstraintTypeenum in favor of string-based constraint types validated againstObjectRelationLibraryRegistry— more extensible, single source of truth. - Simplified
_attach_spatial_constraints_to_assetsusingrelation_class.is_unary()dispatch — replaces the oldAT_POSEspecial-casing with uniform relation handling. - Added
pydantic>=2.0tosetup.pydependencies.
No new issues found. The Pydantic migration is well-structured, validators are correctly placed, and the removal of the hardcoded spatial constraint enum is a clean design improvement.
Update (ffbbb13): Reviewed incremental changes from c5ec371c → ffbbb137. This commit removes the standalone isaaclab_arena/agentic_environment_generation/ package and inlines the resolver logic directly into arena_env_graph_spec.py as the resolve_constraints() method + module-level helper functions. The logic is unchanged — pure code reorganization for better colocation.
Previous concerns status:
- 🟡 #2 (
NotImplementedErrorforSetState): Not addressed (non-blocking). - 🟡 #3 (env_name discrepancy): Not addressed (non-blocking).
- 🟢 #6–#8: Not addressed (non-blocking nits).
No new issues found. Clean refactor — module elimination reduces indirection, helpers are well-organized at module level.
Update (8da01ad): Reviewed incremental changes from ffbbb137 → 8da01adf. This is a documentation and test consolidation pass — no logic changes.
Changes:
- Docstrings trimmed across
arena_env_graph_spec.py,graph_spec_utils.py, andtask_transition.py— shorter and more focused, no semantic loss. test_state_spec_resolver.pydeleted; its tests consolidated intotest_arena_env_graph_spec.pywith minor renames (populated→resolved) and one new assertion (len(state_specs) == len(tasks) + 1).from_dictreplaced explicitValueErrorwithassert(consistent with existingassertusage throughout, but note this compounds concern #3 above).- Removed the multi-line comment block explaining dual reach constraints in
TaskTransition— that reasoning is now implicit.
Previous concerns status:
- 🟡 #2 (
NotImplementedErrorforSetState): Not addressed (non-blocking). - 🟡 #3 (env_name discrepancy): Not addressed (non-blocking).
- 🟢 #6–#8: Not addressed (non-blocking nits).
No new issues found. Clean consolidation pass.
Update (ff95d26): Reviewed incremental changes from 8da01adf → ff95d269. This is a substantial expansion introducing the Variations system — a full framework for build-time and run-time environment variations, plus task-transition–driven constraint resolution, all wired end-to-end with comprehensive tests.
High-Level Architecture
The isaaclab_arena/variations/ package introduces a well-layered stack:
VariationBaseCfg ← configclass (enabled, sampler_cfg)
└── VariationBase ← abstract (enable/disable, apply_cfg, sampler)
├── RunTimeVariationBase ← build_event_cfg() → EventTermCfg
└── BuildTimeVariationBase ← apply() → mutates asset in place
Samplers: SamplerBase → ContinuousSampler → UniformSampler, and separately ChoiceSampler for discrete pools.
Concrete variations: CameraExtrinsicsVariation (run-time), HDRImageVariation (build-time).
Key Changes
Asset.variations— Assets now hold adict[str, VariationBase]withadd_variation/get_variation/get_variations.EmbodimentBasenow inherits fromAssetviasuper().__init__(name=self.name, tags=self.tags)and auto-tags"embodiment".ArenaEnvBuildergainsget_all_variations(),_compose_variations_event_cfg()(run-time), and_apply_build_time_variations()(build-time, called before scene materialization).- Graph spec:
ArenaEnvGraphSpec.resolve_constraints()chains an unresolved graph into full state specs.from_yaml/from_dictacceptis_task_wiring_enabledto support partial graphs. Validation factored intoassert_task_wiring(separate fromassert_references_exist). - Task transitions:
TaskTransition,Relocate,SetStatedataclasses.TaskBase.success_state_transition()as abstract classmethod, implemented byPickAndPlaceTaskandRotateRevoluteJointTask. RelationBase.is_spawn_pose_constraint()— overriddenTruebyAtPosition,PositionLimits,RandomAroundSolution,RotateAroundSolution.- Camera offset fix: Franka
_DEFAULT_CAMERA_OFFSETrotation updated from(0.0, 0.0, -0.66262, -0.74896)to(0.0, 0.0, 0.70711, 0.70711).
Findings
🟡 Medium:
1. EmbodimentBase.__init__ accesses self.name / self.tags before assignment
File: isaaclab_arena/embodiments/embodiment_base.py (lines 32-35)
assert self.name is not None, "Embodiment name is required"
super().__init__(name=self.name, tags=self.tags)This assumes self.name and self.tags are set as class-level attributes on the subclass — a convention that must hold for every embodiment. If any future subclass sets name/tags in its own __init__ after calling super().__init__(), this will fail. Consider documenting this requirement or using a classmethod/class variable protocol.
2. ArenaEnvGraphTaskSpec state-spec fields become Optional with no migration path
File: isaaclab_arena/environments/arena_env_graph_types.py (lines 159-160)
initial_state_spec_id: str | None = None
success_state_spec_id: str | None = NoneThese were previously str = Field(min_length=1). Existing YAML files with NULL string values will now pass validation silently. The init fixture uses YAML NULL which maps to Python None — correct. But the assert_task_wiring guard only fires when is_task_wiring_enabled=True, so partially-wired graphs could slip through if misused.
3. CameraExtrinsicsVariation relies on undocumented Isaac Lab quaternion bug
File: isaaclab_arena/variations/camera_extrinsics_variation.py (line 141)
The code notes that get_local_poses() claims (w,x,y,z) but actually returns (x,y,z,w). This is tested, but if Isaac Lab fixes this upstream, the variation will break silently. Consider adding a version guard or a startup assertion that validates the quaternion layout.
🟢 Low / Nits:
4. spatial_constraint_is_spawn_pose as a free function duplicates what could be a method
File: isaaclab_arena/environments/graph_spec_utils.py (lines 170-177)
This function does relation_class_for_spatial_constraint_type(constraint_type).is_spawn_pose_constraint(). Consider making it a method on the constraint spec to reduce indirection.
5. HDRImageVariation.apply() calls self.sampler.sample(num_samples=1, choices=hdr_names) — type mismatch
The type annotation on VariationBase.sampler is SamplerBase | None, so calling .sample(choices=...) requires knowing the concrete type. This works but loses type safety. A minor typing improvement would be to parametrize the base class.
✅ Positives
- Excellent test coverage: 7 new test files covering variations, transitions, camera extrinsics, HDR, graph resolution, and the Isaac Lab quaternion bug.
- Clean separation: Build-time vs run-time variations are properly distinguished architecturally.
- Sampler abstraction: The
SamplerBaseCfg.build()→SamplerBasepattern is extensible and cfg-driven. resolve_constraints(): The worked example in the docstring makes the algorithm easy to follow.- Defensive assertions:
add_variationprevents duplicate names; samplers validate shapes; transitions validate effects. - The camera extrinsics test that validates the actual ROS→OpenGL→parent frame math end-to-end is particularly thorough.
Previous concerns status
- 🟡 #2 (
NotImplementedErrorforSetState): Still present in_get_task_relocations— non-blocking. - 🟡 #3 (env_name discrepancy): Init fixture uses
pick_and_place_maple_table_init; test passesenv_nameexplicitly — partially mitigated.
Verdict
Approve. The variations system is well-designed, properly tested, and cleanly integrated. The remaining items are non-blocking.
Update (4de779a): Reviewed incremental changes from ff95d269 → 4de779aa. This is a major simplification of constraint resolution, shifting from full-state propagation to a delta-based model where derived states record only what their task changed.
Key Architectural Changes
-
Delta-based state specs —
ArenaEnvGraphStateSpecgains anis_delta: boolfield (defaultTrue). The initial state (state_spec_0) is a full snapshot (is_delta=False); every resolved successor state records only the spatial constraints introduced by its task's effects. This eliminates the complex logic for carrying forward, filtering, and reprefixing constraints from previous states. -
Keyword-argument–based
success_state_transition— Task classes now declare the assets they act on as namedAsset-typed parameters (e.g.pick_up_object: Asset, destination_location: Asset) instead of receiving a rawMapping[str, Any]. The resolver constructsAsset(name=value)for each string task arg and passes them as**kwargs. This gives type safety and self-documenting signatures. -
Removal of task constraints (reach constraints) — The
task_constraintsfield,ArenaEnvGraphTaskConstraintSpec, reach-target logic, and embodiment-id extraction are all removed. This significantly simplifies both the resolver and the YAML fixtures. -
Removal of
is_spawn_pose_constraint()— No longer needed since delta states don't carry forward or filter constraints by type. Removed fromRelationBase,AtPosition,PositionLimits,RandomAroundSolution,RotateAroundSolution, and thespatial_constraint_is_spawn_posefree function.
Previous concerns status
- 🟡 #2 (
NotImplementedErrorforSetState): ✅ Resolved._get_task_effectsnow gracefully filters to onlyRelocateinstances instead of raising —SetStateeffects are silently skipped. This is the correct behavior given the TODO for joint state support. - 🟡 Previous #4/#5 (redundant registry lookups): Still clean —
_get_task_state_transitionis called once per task in the loop. - 🟢 #4 (
spatial_constraint_is_spawn_poseindirection): ✅ Resolved — function removed entirely.
Findings
🟢 Low / Nits:
1. _get_task_effects name is slightly misleading
File: isaaclab_arena/environments/arena_env_graph_spec.py
The function is named _get_task_effects but it only returns Relocate effects (silently filtering out SetState). Consider renaming to _get_task_relocations (which was the previous name) or adding a brief inline comment noting the filter.
2. is_delta defaults to True — initial states must explicitly opt out
File: isaaclab_arena/environments/arena_env_graph_types.py
The default is_delta=True means any hand-authored state that forgets to set is_delta: false will be treated as a delta. Since only state_spec_0 in init fixtures should be is_delta=False, this is a reasonable default for the resolution output, but could trip up YAML authors. The init fixture correctly sets it — just worth documenting.
3. assert curr_task_transition is not None is redundant
File: isaaclab_arena/environments/arena_env_graph_spec.py (line ~174)
_get_task_state_transition always returns a TaskTransition (the registry assert + success_state_transition call will raise before returning None). The assert is defensive but can never fire. Harmless — just noise.
4. Test imports use importlib pattern for lazy loading — good
File: isaaclab_arena/tests/test_task_transitions.py
The lazy import pattern to avoid pulling in omni/pxr at collection time is a good practice for keeping test sessions clean.
✅ Positives
- Dramatic simplification: ~130 lines of complex state-propagation logic replaced by ~15 lines of delta construction. The code is now much easier to reason about.
- Type-safe task signatures: Named
Assetparameters make the contract between tasks and the resolver explicit and IDE-friendly. - Clean fixture updates: The YAML golden fixtures are properly trimmed to reflect the delta model, and tests assert
is_deltavalues. Effecttype alias exported:task_transition.pynow exports theEffect = Relocate | SetStateunion, making the protocol clearer.@dataclass()→@dataclass: Minor but correct style cleanup.- The
ObjectRelationLibraryRegistrylookup inPickAndPlaceTask.success_state_transition(for"on") ensures the relation name is validated against the registry — good defensive practice.
Verdict
Approve. This is an excellent simplification that removes significant complexity while preserving correctness. The delta-based model is conceptually cleaner and the kwarg-based transition API is a meaningful type-safety improvement. No blocking issues.
Update (b9d42f9): Reviewed incremental changes from 4de779aa → b9d42f92. This commit refines how the resolver discovers which task args to pass as Asset kwargs to success_state_transition — replacing the previous "pass all string args as Assets" heuristic with an explicit class-level declaration.
Key Changes
-
TaskBase.success_transition_asset_args— New class attribute (tuple[str, ...]) onTaskBasethat declares exactly which task args the success transition needs. Default is(). Subclasses override alongsidesuccess_state_transition. -
Concrete declarations —
PickAndPlaceTaskdeclares("pick_up_object", "destination_location");RotateRevoluteJointTaskdeclares("openable_object",). -
Resolver uses declaration —
_get_task_state_transitionnow iterates overtask_cls.success_transition_asset_argswith anassertthat each name exists intask.task_args, then builds only those intoAssetkwargs. This replaces the old pattern of passing all string-valued args (which could accidentally include non-asset args like scene or episode config). -
TaskBase.success_state_transitionsignature — Base class drops the**_catch-all; now takes no arguments. Subclasses define their own named parameters as before.
Previous concerns status
- 🟢 Previous #1 (
_get_task_effectsname): Not addressed (non-blocking nit). - 🟢 Previous #2 (
is_deltadefaults): Not addressed (non-blocking nit). - 🟢 Previous #3 (redundant assert): Not addressed (non-blocking nit).
Findings
🟢 Low / Nits:
-
assertfor missing task args will be stripped in optimized mode — The resolver usesassert name in task.task_args(line ~235). If Python is run with-O(assertions stripped), missing args will produce a less informativeKeyErroron the next line rather than the custom message. Considerif name not in task.task_args: raise ValueError(...)for production robustness. Non-blocking since tests run without-O. -
Empty
success_transition_asset_args = ()base default means a task that forgets to override will silently pass no assets — If a future task needs assets but forgets the declaration,success_state_transition()will be called with zero arguments and presumably raiseTypeError. The error message won't hint at the missing class attribute. Consider adding a check or a note in the docstring. The current docstring does say "Override this alongsidesuccess_state_transition" which is adequate.
✅ Positives
- Explicit over implicit: Moving from "pass all strings" to a declared manifest eliminates a class of subtle bugs (e.g. a task arg like
scene_name: strbeing incorrectly wrapped inAsset). - Self-documenting: Reading a task class now immediately shows which args drive the transition, without inspecting the resolver.
- Good assertion message: The assert provides task id, type, and missing arg name — helpful for debugging graph specs.
- Minimal diff: Clean, focused change — no unrelated modifications.
Verdict
Approve. This is a clean, targeted improvement that makes the task→resolver contract explicit and type-safe. No blocking issues.
Update (74fce31): Trivial docstring formatting fix in arena_env_graph_spec.py — collapsed a two-line closing """ into a single-line docstring terminator on _get_task_state_transition(). No logic changes, no functional impact. Verdict unchanged: Approve.
Incremental Update (2bb7002) — 2 fix commits since 74fce31:
These commits address correctness and robustness in resolve_constraints():
✅ Good changes:
- Removed stray
return selffrom_run_validation()— it's a void method, this was dead code state_spec_0is now explicitly set tois_delta: Falseduring resolution rather than relying on the YAML fixture to declare it — this is a correctness fix since the initial state is never a delta by definition- Task wiring now uses
states[i].idinstead of hardcodedf"state_spec_{i}"— more robust if IDs ever diverge from positional naming cli_override_specsis now forwarded to the resolved graph — previously these were silently dropped during resolution- Test updated to use
resolved.tasksdirectly (preserved order) instead of sorting by id — cleaner and less fragile
📝 Note: The YAML fixture removal of is_delta: false is the correct companion change — that flag is now set programmatically, making the YAML simpler and the code the single source of truth.
No new concerns. Previous findings still apply (especially #2 regarding SetState handling and #4 on TaskRegistry() instantiation).
Update (48111fc): Reviewed incremental changes — branch rebased/force-pushed (18 commits). The latest commits are labeled "address reviews." This is the final state of the PR.
Previous concerns status
- 🟡 #2 (
NotImplementedErrorforSetState): ✅ Resolved (earlier push) —_get_task_effectsgracefully filters toRelocateonly. - 🟡 #3 (env_name discrepancy): Partially mitigated — the test uses
resolve_constraints()without explicitenv_name, and the init/full fixtures use distinct names by design. - 🟡
EmbodimentBase.__init__accessing class-levelself.name/self.tags: Not addressed (non-blocking; documented by convention). - 🟡 Camera extrinsics quaternion bug reliance: Not addressed (non-blocking; regression test
test_isaaclab_bug_get_local_poses.pyguards it). - 🟡
ArenaEnvGraphTaskSpecOptional state-spec fields: Not addressed (non-blocking;assert_task_wiringguard is adequate). - 🟢
_get_task_effectsnaming nit: ✅ Addressed — function is now clearly documented as filtering to relocations only. - 🟢
is_deltadefault nit: Not addressed (non-blocking; works correctly for resolved output). - 🟢 Redundant assert nit: Not addressed (harmless defensive code).
Changes since last reviewed state
-
success_transition_asset_argsremoved in favor of**_kwargs — The explicit class-attribute approach from theb9d42f92push has been simplified:success_state_transitionnow accepts task arg strings directly as named parameters with**_to absorb extras (e.g.episode_length_s,background_scene). This is a cleaner API — the task signature itself documents which args it needs._get_task_state_transitionpasses**task.task_argsdirectly withoutAssetwrapping. -
Docker multi-clone support —
run_docker.shderives container suffix from repo directory name; agent skills updated to use$ARENA_CONTAINERvariable instead of hardcoded names. Clean infrastructure improvement. -
entrypoint.shordering fix —userdelcalls moved beforegroupaddto prevent stale user conflicts. Defensive improvement. -
Galileo G1 environment — Default
task_descriptionsimplified; dynamic template generation removed. Unrelated cleanup. -
RelationBasebody removed — Emptypassdropped (the class just inherits without adding anything extra). Cosmetic.
No new issues found.
The final state is well-structured. The kwargs-based success_state_transition API is the cleanest iteration yet — self-documenting, type-safe via parameter names, and no auxiliary metadata needed.
Verdict
Approve. All critical and medium issues from prior reviews have been addressed. Remaining items are non-blocking nits/documentation suggestions.
Update (commit c869f68): The new commits completely refactored the state resolution approach. state_spec_resolver.py was removed and replaced with UnresolvedArenaEnvGraphSpec.resolve() in arena_env_graph_spec.py. Previous concerns addressed:
- ✅ Keyword argument mismatch — File removed; logic rewritten cleanly in the new
resolve()method. - ✅ Spawn-pose logic cleanup — Removed entirely. New approach only tracks delta constraints from task transitions.
No new issues found in the incremental changes.
Update (7db64b3): Reviewed incremental changes from c869f684 → 7db64b36. This commit promotes the initial state from a flat constraint list to a proper ArenaEnvGraphStateSpec — the unresolved graph now owns a fully-typed initial state instead of raw SpatialRelationSpec entries.
Key Changes
-
UnresolvedArenaEnvGraphSpec.initial_state_specreplacesspatial_constraints: list[SpatialRelationSpec]— the LLM pipeline now produces a completeArenaEnvGraphStateSpec(withid,is_delta: false, and individually-identified constraints) rather than a flat list that required wrapping. -
_build_initial_state_specremoved — no longer needed since the initial state arrives pre-structured. The synthetic ID generation (state_spec_0_i_kind_subject) now lives in the YAML fixture / upstream producer. -
assert_references_exist→assert_constraint_references— renamed and simplified (dropstasksparameter). Task wiring is now validated exclusively byassert_task_wiring, giving clearer separation of concerns. -
UnresolvedArenaEnvGraphSpec.validate()— new Pydantic model validator that runsassert_unique_ids,assert_constraint_references, andassert_spatial_constraint_shapeson the unresolved graph at parse time. Good — catches malformed LLM output early. -
Test imports —
ArenaEnvGraphNodeType/ArenaEnvGraphObjectReferenceNodeSpecmoved to import fromarena_env_graph_types(reflecting the type module split from earlier commits).
Previous concerns status
- 🟡 #3 (env_name discrepancy): Not addressed (non-blocking).
- 🟢 Remaining nits: Not addressed (non-blocking).
No new issues found.
Clean, focused change that eliminates a layer of indirection. Having the unresolved graph own a validated ArenaEnvGraphStateSpec from the start is the right call — it ensures schema-level validation fires immediately and makes the resolve() method trivially simple (just pass through self.initial_state_spec as the first state).
Verdict: Approve. No blocking issues.
Update (a159634): Applied greptile-bot suggestion in graph_spec_utils.py — adds a None guard on task_constraint.child before asserting membership in node_ids. Correct fix: ArenaEnvGraphTaskConstraintSpec.child is Optional[str], so the unconditional assert would crash on constraints that legitimately have no child (e.g. unary task constraints). No new issues. Verdict unchanged: Approve.
qianl-nv
left a comment
There was a problem hiding this comment.
Thanks for making the change.
I've been playing with unifying IntentSpec with the GraphSpec today and realize the bigger blocker is actually pydantic vs dataclass. Pydantic BaseModel provides much bette default methods for validation/schema generation and is a lot more LLM-interface friendly.
I make a MR to move the whole ArenaEnvGraphSpec to pydantic https://github.com/isaac-sim/IsaacLab-Arena/pull/757/changes
It should also simply a lot of your serialization code in this MR.
6048ef6 to
4fbd3dd
Compare
qianl-nv
left a comment
There was a problem hiding this comment.
Thanks for rebasing on the pydantic refacotor MR. This change is now much simpler. LGTM!
3b471e3 to
9f19eea
Compare
alexmillane
left a comment
There was a problem hiding this comment.
Thanks for this. Looks pretty good.
It's quite complex, what's going on. Let's have a call so I can fully understand the intent?
But it's close I think!
Greptile SummaryThis PR wires the LLM-produced initial-state-only graph into a fully-resolved
Confidence Score: 5/5Safe to merge; the resolution logic is correctly validated end-to-end by the existing ArenaEnvGraphSpec model validator, and both new test files cover the key chain-wiring and error paths. The core resolve() loop is straightforward: it iterates the task list, calls the task class's declared classmethod, and constructs delta states whose constraint references are validated by assert_constraint_references before the result is accepted. No existing callers are broken — the old ArenaEnvGraphSpec API is unchanged. The two minor findings (shared node references and the no-op base validate()) do not affect correctness of the current code paths. arena_env_graph_spec.py — the node-sharing between UnresolvedArenaEnvGraphSpec and the resolved ArenaEnvGraphSpec is worth a follow-up if apply_cli_override_args is ever called on the resolved spec while the unresolved spec is still live. Important Files Changed
Sequence DiagramsequenceDiagram
participant Caller
participant UnresolvedSpec as UnresolvedArenaEnvGraphSpec
participant TaskRegistry
participant TaskCls as TaskBase subclass
participant Helpers as _get_task_state_transition / _get_next_state_spec
participant FullSpec as ArenaEnvGraphSpec
Caller->>UnresolvedSpec: from_yaml(path)
UnresolvedSpec->>UnresolvedSpec: validate() [unique IDs, constraint refs, shapes]
Caller->>UnresolvedSpec: resolve()
Note over UnresolvedSpec: states = [initial_state_spec]
loop for each TaskSpec (i, task)
UnresolvedSpec->>TaskRegistry: get_task_by_name(task.kind)
TaskRegistry-->>UnresolvedSpec: task_cls
UnresolvedSpec->>TaskCls: "success_state_transition(**task.params)"
TaskCls-->>UnresolvedSpec: "TaskTransition(effects=[Relocate(...)])"
UnresolvedSpec->>Helpers: "_get_next_state_spec(state_spec_{i+1}, transition)"
Helpers-->>UnresolvedSpec: "ArenaEnvGraphStateSpec(is_delta=True)"
Note over UnresolvedSpec: Wire task: state_spec_i to state_spec_{i+1}
end
UnresolvedSpec->>FullSpec: ArenaEnvGraphSpec(nodes, tasks, state_specs)
FullSpec->>FullSpec: validate() [unique IDs, constraint refs, wiring, shapes]
FullSpec-->>Caller: ArenaEnvGraphSpec
Reviews (11): Last reviewed commit: "Apply suggestion from @greptile-apps[bot..." | Re-trigger Greptile |
8da01ad to
ff95d26
Compare
alexmillane
left a comment
There was a problem hiding this comment.
Looking much better.
Thanks for working on this. I have a few comments:
- reducing coupling between the run-time tasks and the graphs.
- do we handle prior removing constraints when we get a transition.
One broader question that I have is how these future constraints are actually used. Presumably we will fold thesefuture constraints back to something that operates on the initial layout of the scene?
alexmillane
left a comment
There was a problem hiding this comment.
Feel free to merge when you address the above
2bb7002 to
6eef4da
Compare
Resolve conflicts between PR #753 (state-spec resolver / delta model) and PR #718 (Pydantic migration / agentic env gen): - arena_env_graph_types.py: add is_delta field to ArenaEnvGraphStateSpec; use ArenaEnvGraphSpatialRelationSpec (Pydantic-migrated name); make ArenaEnvGraphTaskSpec.initial_state_spec_id / success_state_spec_id Optional[str] to support unresolved graphs; drop duplicate class def - arena_env_graph_spec.py: update _get_task_state_transition to use task.kind / task.params (TaskSpec field names); replace stale ArenaEnvGraphSpatialConstraintSpec with ArenaEnvGraphSpatialRelationSpec and old type/parent/child fields with kind/reference/subject - pick_and_place_task.py: import both agent_ready and ObjectRelationLibraryRegistry - test_arena_env_graph_spec.py: use delta-model assertions with Pydantic field names (kind/reference/subject) - pick_and_place_maple_table_env_graph.yaml: delta structure from PR #753, Pydantic field names from PR #718, no task_constraints Signed-off-by: Qian Lin <qianl@nvidia.com>
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Summary
Resolve initial state specs to full state specs in env graph
Detailed description
TaskTransitionmodel such that it describes what effect a task causes to the Env if success condition is met. E.g.Relocateto another spatial relation, orSetStateto another state for embodiment/object itself. Each task class has an optionalsuccess_state_transitionand can be access thru theTaskRegistry.resolve_constraintsmethod inArenaEnvGraphSpec: it chains each task's declared success condition onto the previous state to derive the missing states and wire every task. Topology is implicit in the sequential list, so N tasks yield N+1 state specs.APIs
Known limitation/ TODOs
Relocaterealization for PnP task using "hardcode" object ON destination