Skip to content

Resolve initial state specs to full state specs in env graph#753

Merged
xyao-nv merged 22 commits into
mainfrom
xyao/dev/state_spec_resolver
Jun 9, 2026
Merged

Resolve initial state specs to full state specs in env graph#753
xyao-nv merged 22 commits into
mainfrom
xyao/dev/state_spec_resolver

Conversation

@xyao-nv

@xyao-nv xyao-nv commented Jun 3, 2026

Copy link
Copy Markdown
Collaborator

Summary

Resolve initial state specs to full state specs in env graph

Detailed description

  • Why:
  1. upstream produces only nodes, the task list, and the initial state estimation(state_spec_0)
  2. Populating it to intermediate & final states to reflect how Env State changes along the task execution
  3. Note that those Env State changes may pose constraints on the initial Env placement.
  • What:
  1. Add a declarative TaskTransition model such that it describes what effect a task causes to the Env if success condition is met. E.g. Relocate to another spatial relation, or SetState to another state for embodiment/object itself. Each task class has an optional success_state_transition and can be access thru the TaskRegistry.
  2. Add resolve_constraints method in ArenaEnvGraphSpec: 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.
  3. Add a constraint update mechanism, such that unrelated ones are kept, related ones are updated from dropped based on the task state transition. Also include reachability of both a state's postcondition and precondition target.

APIs

unresolved_graph_spec  = ArenaEnvGraphSpec.from_yaml('<unresolved_yaml_path>', is_task_wiring_enabled=False)
full_graph_spec = unresolved_graph_spec.resolve_constraints()

Known limitation/ TODOs

  • Only work for sequential tasks so far
  • Only implemented Relocate realization for PnP task using "hardcode" object ON destination

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

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.

🤖 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_posdroid). 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 → UnresolvedArenaEnvGraphSpecStateSpecResolver.resolve_path()ArenaEnvGraphSpec
  • TaskTransition / Relocate / SetState effect algebra is well-designed and extensible
  • RelationBase.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_yaml enable 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

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

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.

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

✅ 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 generic NotImplementedError.

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.
  • TaskTransition as 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=False to parameterize existing validation rather than duplicating it.

📝 Update (94c1dce)

Reviewed incremental changes from c0b45d894c1dce. 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_state is 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 94c1dcee3b471e37. 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 TaskRegistry lookups / double _transition_for call): ✅ Fixed. Transitions are now pre-computed once via transitions = [_get_task_state_transition(task) for task in spec.tasks].
  • 🟡 #2 (NotImplementedError for SetState): 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 3b471e37c5ec371c.

Previous concerns status:

  • 🔴 Keyword argument mismatch (prev_state vs prev_state_spec): ✅ Fixed. Line 72 now correctly uses prev_state_spec=states[-1].
  • 🟡 #2 (NotImplementedError for SetState): 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.py and arena_env_graph_spec.py from dataclasses to Pydantic BaseModel with validators — cleaner validation and serialization.
  • Removed ArenaEnvGraphSpatialConstraintType enum in favor of string-based constraint types validated against ObjectRelationLibraryRegistry — more extensible, single source of truth.
  • Simplified _attach_spatial_constraints_to_assets using relation_class.is_unary() dispatch — replaces the old AT_POSE special-casing with uniform relation handling.
  • Added pydantic>=2.0 to setup.py dependencies.

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 c5ec371cffbbb137. 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 (NotImplementedError for SetState): 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 ffbbb1378da01adf. This is a documentation and test consolidation pass — no logic changes.

Changes:

  • Docstrings trimmed across arena_env_graph_spec.py, graph_spec_utils.py, and task_transition.py — shorter and more focused, no semantic loss.
  • test_state_spec_resolver.py deleted; its tests consolidated into test_arena_env_graph_spec.py with minor renames (populatedresolved) and one new assertion (len(state_specs) == len(tasks) + 1).
  • from_dict replaced explicit ValueError with assert (consistent with existing assert usage 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 (NotImplementedError for SetState): 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 8da01adfff95d269. 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: SamplerBaseContinuousSamplerUniformSampler, and separately ChoiceSampler for discrete pools.

Concrete variations: CameraExtrinsicsVariation (run-time), HDRImageVariation (build-time).

Key Changes

  1. Asset.variations — Assets now hold a dict[str, VariationBase] with add_variation/get_variation/get_variations.
  2. EmbodimentBase now inherits from Asset via super().__init__(name=self.name, tags=self.tags) and auto-tags "embodiment".
  3. ArenaEnvBuilder gains get_all_variations(), _compose_variations_event_cfg() (run-time), and _apply_build_time_variations() (build-time, called before scene materialization).
  4. Graph spec: ArenaEnvGraphSpec.resolve_constraints() chains an unresolved graph into full state specs. from_yaml/from_dict accept is_task_wiring_enabled to support partial graphs. Validation factored into assert_task_wiring (separate from assert_references_exist).
  5. Task transitions: TaskTransition, Relocate, SetState dataclasses. TaskBase.success_state_transition() as abstract classmethod, implemented by PickAndPlaceTask and RotateRevoluteJointTask.
  6. RelationBase.is_spawn_pose_constraint() — overridden True by AtPosition, PositionLimits, RandomAroundSolution, RotateAroundSolution.
  7. Camera offset fix: Franka _DEFAULT_CAMERA_OFFSET rotation 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 = None

These 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()SamplerBase pattern is extensible and cfg-driven.
  • resolve_constraints(): The worked example in the docstring makes the algorithm easy to follow.
  • Defensive assertions: add_variation prevents 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 (NotImplementedError for SetState): Still present in _get_task_relocations — non-blocking.
  • 🟡 #3 (env_name discrepancy): Init fixture uses pick_and_place_maple_table_init; test passes env_name explicitly — 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 ff95d2694de779aa. 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

  1. Delta-based state specsArenaEnvGraphStateSpec gains an is_delta: bool field (default True). 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.

  2. Keyword-argument–based success_state_transition — Task classes now declare the assets they act on as named Asset-typed parameters (e.g. pick_up_object: Asset, destination_location: Asset) instead of receiving a raw Mapping[str, Any]. The resolver constructs Asset(name=value) for each string task arg and passes them as **kwargs. This gives type safety and self-documenting signatures.

  3. Removal of task constraints (reach constraints) — The task_constraints field, ArenaEnvGraphTaskConstraintSpec, reach-target logic, and embodiment-id extraction are all removed. This significantly simplifies both the resolver and the YAML fixtures.

  4. Removal of is_spawn_pose_constraint() — No longer needed since delta states don't carry forward or filter constraints by type. Removed from RelationBase, AtPosition, PositionLimits, RandomAroundSolution, RotateAroundSolution, and the spatial_constraint_is_spawn_pose free function.

Previous concerns status

  • 🟡 #2 (NotImplementedError for SetState): ✅ Resolved. _get_task_effects now gracefully filters to only Relocate instances instead of raising — SetState effects 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_transition is called once per task in the loop.
  • 🟢 #4 (spatial_constraint_is_spawn_pose indirection): ✅ 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 Asset parameters 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_delta values.
  • Effect type alias exported: task_transition.py now exports the Effect = Relocate | SetState union, making the protocol clearer.
  • @dataclass()@dataclass: Minor but correct style cleanup.
  • The ObjectRelationLibraryRegistry lookup in PickAndPlaceTask.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 4de779aab9d42f92. 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

  1. TaskBase.success_transition_asset_args — New class attribute (tuple[str, ...]) on TaskBase that declares exactly which task args the success transition needs. Default is (). Subclasses override alongside success_state_transition.

  2. Concrete declarationsPickAndPlaceTask declares ("pick_up_object", "destination_location"); RotateRevoluteJointTask declares ("openable_object",).

  3. Resolver uses declaration_get_task_state_transition now iterates over task_cls.success_transition_asset_args with an assert that each name exists in task.task_args, then builds only those into Asset kwargs. This replaces the old pattern of passing all string-valued args (which could accidentally include non-asset args like scene or episode config).

  4. TaskBase.success_state_transition signature — 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_effects name): Not addressed (non-blocking nit).
  • 🟢 Previous #2 (is_delta defaults): Not addressed (non-blocking nit).
  • 🟢 Previous #3 (redundant assert): Not addressed (non-blocking nit).

Findings

🟢 Low / Nits:

  1. assert for missing task args will be stripped in optimized mode — The resolver uses assert name in task.task_args (line ~235). If Python is run with -O (assertions stripped), missing args will produce a less informative KeyError on the next line rather than the custom message. Consider if name not in task.task_args: raise ValueError(...) for production robustness. Non-blocking since tests run without -O.

  2. 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 raise TypeError. 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 alongside success_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: str being incorrectly wrapped in Asset).
  • 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 self from _run_validation() — it's a void method, this was dead code
  • state_spec_0 is now explicitly set to is_delta: False during 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].id instead of hardcoded f"state_spec_{i}" — more robust if IDs ever diverge from positional naming
  • cli_override_specs is now forwarded to the resolved graph — previously these were silently dropped during resolution
  • Test updated to use resolved.tasks directly (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 (NotImplementedError for SetState): ✅ Resolved (earlier push) — _get_task_effects gracefully filters to Relocate only.
  • 🟡 #3 (env_name discrepancy): Partially mitigated — the test uses resolve_constraints() without explicit env_name, and the init/full fixtures use distinct names by design.
  • 🟡 EmbodimentBase.__init__ accessing class-level self.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.py guards it).
  • 🟡 ArenaEnvGraphTaskSpec Optional state-spec fields: Not addressed (non-blocking; assert_task_wiring guard is adequate).
  • 🟢 _get_task_effects naming nit: ✅ Addressed — function is now clearly documented as filtering to relocations only.
  • 🟢 is_delta default nit: Not addressed (non-blocking; works correctly for resolved output).
  • 🟢 Redundant assert nit: Not addressed (harmless defensive code).

Changes since last reviewed state

  1. success_transition_asset_args removed in favor of **_ kwargs — The explicit class-attribute approach from the b9d42f92 push has been simplified: success_state_transition now 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_transition passes **task.task_args directly without Asset wrapping.

  2. Docker multi-clone supportrun_docker.sh derives container suffix from repo directory name; agent skills updated to use $ARENA_CONTAINER variable instead of hardcoded names. Clean infrastructure improvement.

  3. entrypoint.sh ordering fixuserdel calls moved before groupadd to prevent stale user conflicts. Defensive improvement.

  4. Galileo G1 environment — Default task_description simplified; dynamic template generation removed. Unrelated cleanup.

  5. RelationBase body removed — Empty pass dropped (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 c869f6847db64b36. 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

  1. UnresolvedArenaEnvGraphSpec.initial_state_spec replaces spatial_constraints: list[SpatialRelationSpec] — the LLM pipeline now produces a complete ArenaEnvGraphStateSpec (with id, is_delta: false, and individually-identified constraints) rather than a flat list that required wrapping.

  2. _build_initial_state_spec removed — 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.

  3. assert_references_existassert_constraint_references — renamed and simplified (drops tasks parameter). Task wiring is now validated exclusively by assert_task_wiring, giving clearer separation of concerns.

  4. UnresolvedArenaEnvGraphSpec.validate() — new Pydantic model validator that runs assert_unique_ids, assert_constraint_references, and assert_spatial_constraint_shapes on the unresolved graph at parse time. Good — catches malformed LLM output early.

  5. Test importsArenaEnvGraphNodeType / ArenaEnvGraphObjectReferenceNodeSpec moved to import from arena_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 qianl-nv left a comment

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.

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.

Comment thread isaaclab_arena/environments/arena_env_graph_spec.py
Comment thread isaaclab_arena/environments/arena_env_graph_spec.py Outdated
Comment thread isaaclab_arena/agentic_environment_generation/state_spec_resolver.py Outdated
@xyao-nv xyao-nv changed the base branch from main to qianl/dev/pydantic_spec June 3, 2026 22:24
@xyao-nv xyao-nv force-pushed the xyao/dev/state_spec_resolver branch from 6048ef6 to 4fbd3dd Compare June 3, 2026 22:25

@qianl-nv qianl-nv left a comment

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.

Thanks for rebasing on the pydantic refacotor MR. This change is now much simpler. LGTM!

Base automatically changed from qianl/dev/pydantic_spec to main June 4, 2026 06:06
@xyao-nv xyao-nv force-pushed the xyao/dev/state_spec_resolver branch 2 times, most recently from 3b471e3 to 9f19eea Compare June 4, 2026 06:22
Comment thread isaaclab_arena/agentic_environment_generation/state_spec_resolver.py Outdated
@xyao-nv xyao-nv changed the title Resolve initial state specs to full env graph Resolve initial state specs to full state specs in env graph Jun 4, 2026
Comment thread isaaclab_arena/environments/arena_env_graph_spec.py Outdated

@alexmillane alexmillane left a comment

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.

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!

Comment thread isaaclab_arena/environments/arena_env_graph_spec.py Outdated
Comment thread isaaclab_arena/environments/arena_env_graph_spec.py Outdated
Comment thread isaaclab_arena/tasks/task_transition.py Outdated
Comment thread isaaclab_arena/environments/arena_env_graph_spec.py Outdated
Comment thread isaaclab_arena/tasks/pick_and_place_task.py Outdated
Comment thread isaaclab_arena/environments/arena_env_graph_spec.py Outdated
Comment thread isaaclab_arena/environments/arena_env_graph_spec.py Outdated
Comment thread isaaclab_arena/tests/test_data/pick_and_place_maple_table_init_env_graph.yaml Outdated
Comment thread isaaclab_arena/tasks/pick_and_place_task.py Outdated
Comment thread isaaclab_arena/tasks/pick_and_place_task.py Outdated
@xyao-nv xyao-nv marked this pull request as ready for review June 4, 2026 20:30
@greptile-apps

greptile-apps Bot commented Jun 4, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR wires the LLM-produced initial-state-only graph into a fully-resolved ArenaEnvGraphSpec by chaining each task's declared success_state_transition into sequential delta states. The approach cleanly separates the unresolved schema (UnresolvedArenaEnvGraphSpec) from the validated full schema via a shared ArenaEnvGraphSpecBase, and introduces a declarative TaskTransition model (Relocate / SetState) that task classes implement to describe their side-effects.

  • UnresolvedArenaEnvGraphSpec.resolve(): iterates over the sequential task list, looks up each task class via TaskRegistry, calls success_state_transition(**params), and wires the resulting Relocate effects into delta ArenaEnvGraphStateSpec objects indexed as state_spec_{i+1}; the output is validated end-to-end by the existing ArenaEnvGraphSpec model validator.
  • graph_spec_utils refactor: assert_references_exist is split into assert_constraint_references (node cross-references) and assert_task_wiring (state-spec ID references), and task constraint ID uniqueness tracking is intentionally dropped in line with the removal of task_constraints from the YAML format.
  • Known limitations (documented in the PR): only sequential topologies are supported, only Relocate effects are materialized (joint-state SetState is a tracked TODO), and cli_override_specs is not yet threaded through UnresolvedArenaEnvGraphSpec.

Confidence Score: 5/5

Safe 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

Filename Overview
isaaclab_arena/environments/arena_env_graph_spec.py Introduces ArenaEnvGraphSpecBase + UnresolvedArenaEnvGraphSpec + resolve(). Logic is sound; shared node-object references between the unresolved and resolved specs mean apply_cli_override_args mutations on the resolved spec will silently bleed into the original unresolved spec.
isaaclab_arena/environments/graph_spec_utils.py Splits assert_references_exist into assert_constraint_references + assert_task_wiring; correctly preserves the child-None guard while safely removing the redundant parent-None guard (parent is str, not Optional[str]).
isaaclab_arena/tasks/task_transition.py New declarative model (Relocate, SetState, TaskTransition) with reach_target_on_success property; pure Python, well-documented, no issues.
isaaclab_arena/tasks/task_base.py Adds success_state_transition classmethod with NotImplementedError fallback; intentionally not @AbstractMethod (classmethods can't be enforced abstractly at call time), but the failure surface is deferred to resolve() rather than class definition.
isaaclab_arena/tasks/pick_and_place_task.py Implements success_state_transition by looking up 'on' from the registry and emitting a Relocate effect; correctly delegates relation lookup to registry rather than hardcoding the string.
isaaclab_arena/tasks/rotate_revolute_joint_task.py Stub success_state_transition that emits no effects; clearly documented as a TODO pending joint-openness support in the graph schema.
isaaclab_arena/environments/arena_env_graph_types.py Adds is_delta: bool = True field to ArenaEnvGraphStateSpec and introduces TaskSpec base class; straightforward schema change.
isaaclab_arena/tests/test_arena_env_graph_spec.py Adds resolve() integration tests including chain wiring, spatial content comparison, and not-implemented task rejection; uses model-level mutation to inject NoTask which bypasses field validators but works because NoTask IS registered.
isaaclab_arena/tests/test_task_transitions.py New unit tests for TaskTransition logic; correctly uses lazy imports to avoid loading pxr before SimulationApp startup.
isaaclab_arena/tests/test_data/pick_and_place_maple_table_env_graph.yaml Refactored to use differential delta states (is_delta: true) with one constraint per state instead of full snapshots; task_constraints removed from all state specs.
isaaclab_arena/tests/test_data/pick_and_place_maple_table_init_env_graph.yaml New unresolved graph fixture; initial_state_spec correctly sets is_delta: false; tasks carry only kind/params with no state wiring.

Sequence Diagram

sequenceDiagram
    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
Loading

Reviews (11): Last reviewed commit: "Apply suggestion from @greptile-apps[bot..." | Re-trigger Greptile

Comment thread isaaclab_arena/environments/arena_env_graph_spec.py Outdated
Comment thread isaaclab_arena/environments/arena_env_graph_spec.py Outdated
Comment thread isaaclab_arena/environments/arena_env_graph_spec.py Outdated
Comment thread isaaclab_arena/environments/graph_spec_utils.py
Comment thread isaaclab_arena/tests/test_arena_env_graph_spec.py
@xyao-nv xyao-nv force-pushed the xyao/dev/state_spec_resolver branch from 8da01ad to ff95d26 Compare June 4, 2026 22:34
Comment thread isaaclab_arena/environments/arena_env_graph_spec.py Outdated
Comment thread isaaclab_arena/environments/arena_env_graph_spec.py Outdated

@alexmillane alexmillane left a comment

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.

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?

Comment thread isaaclab_arena/environments/arena_env_graph_spec.py Outdated
Comment thread isaaclab_arena/environments/arena_env_graph_spec.py Outdated
Comment thread isaaclab_arena/tasks/rotate_revolute_joint_task.py Outdated
Comment thread isaaclab_arena/tasks/task_base.py Outdated

@alexmillane alexmillane left a comment

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.

Feel free to merge when you address the above

@xyao-nv xyao-nv force-pushed the xyao/dev/state_spec_resolver branch from 2bb7002 to 6eef4da Compare June 5, 2026 19:36
Comment thread isaaclab_arena/environments/arena_env_graph_spec.py Outdated
@xyao-nv xyao-nv enabled auto-merge (squash) June 5, 2026 21:38
qianl-nv added 3 commits June 9, 2026 02:29
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>
Comment thread isaaclab_arena/environments/graph_spec_utils.py Outdated
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
@xyao-nv xyao-nv merged commit 306626f into main Jun 9, 2026
6 checks passed
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.

3 participants