Add ObjectRelation & Task Registries

[xyao-nv]

Summary

Add ObjectRelation Registry & Task Registry

Detailed description

• Why: When constructing ArenaEnv from EnvGraphSpec, it requires matching spatial relationship & task to its class name. Without this registry, new relation/task need to be added in two places manually.
• What: Added @register_object_relation decorator backed by the added ObjectRelationLibraryRegistry. Added @register_task decorator backed by the added TaskRegistry.
• Leave TODO for AT_POSE and IN registry supports

Note

• Added a Bool to block re-entry when using registration decorator. Without it, it triggers circular deps ImportError. It's caused by object relation registry.

[isaaclab-review-bot]

Code Review: Add ObjectRelation Registry

Summary

This PR introduces a registry pattern for RelationBase subclasses via a new ObjectRelationLibraryRegistry and @register_object_relation decorator. The implementation follows established patterns in the codebase and provides a clean mechanism for dynamically resolving spatial constraint types to their corresponding relation classes.

---

✅ Strengths

1. Consistent Design Pattern: The new ObjectRelationLibraryRegistry follows the same singleton pattern as existing registries (DeviceRegistry, EnvironmentRegistry, etc.), maintaining codebase consistency.

2. Clean Decorator API: The @register_object_relation decorator provides an intuitive way to register classes, mirroring the existing @register_environment pattern.

3. Good Test Coverage: The test additions verify both the registry lookup and the expected None returns for unregistered constraint types (AT_POSE, IN).

4. Proper Lazy Loading: Integration with ensure_assets_registered() ensures relations are loaded when needed, avoiding circular import issues.

---

⚠️ Suggestions & Concerns

1. Silent Warning on Duplicate Registration (Medium Priority)

def register_object_relation(cls):
    registry = ObjectRelationLibraryRegistry()
    if registry.is_registered(cls.name):
        print(f"WARNING: Object relation {cls.name} is already registered. Doing nothing.")
    else:
        registry.register(cls, cls.name)
    return cls

Issue: Using print() for warnings can be easily missed and does not follow Python logging conventions.

Suggestion: Consider using warnings.warn() or the logging module for consistency with production code:

import warnings
warnings.warn(f"Object relation {cls.name} is already registered. Skipping.", UserWarning, stacklevel=2)

This mirrors the pattern in other Isaac Lab modules and enables proper warning filtering.

2. Missing name Attribute Validation (Medium Priority)

The decorator assumes cls.name exists. If a class is decorated without defining name, this will raise an AttributeError with a potentially confusing traceback.

Suggestion: Add explicit validation:

def register_object_relation(cls):
    if not hasattr(cls, 'name'):
        raise TypeError(f"Cannot register {cls.__name__}: missing required 'name' class attribute")
    # ... rest of implementation

3. IsAnchor Class Body Change (Minor)

class IsAnchor(RelationBase):
    # was: pass
    name = "is_anchor"

The removal of pass is correct since name now provides a class body, but ensure this does not break any introspection or serialization logic that may have relied on IsAnchor having no attributes.

4. Type Hint Return Consistency (Minor)

In relation_class_for_spatial_constraint_type:

) -> "type[RelationBase] | None":

The string-quoted forward reference is fine, but consider importing RelationBase at runtime (guarded by TYPE_CHECKING) to keep the return type consistent with other functions in the module.

---

🔍 Edge Cases & Error Handling

1. Thread Safety: The singleton registry pattern uses a class-level _instances dict. If Isaac Lab ever runs with multi-threaded asset registration, this could cause race conditions. Not a concern for current usage patterns, but worth noting.

2. Registry State Persistence: The registry is populated via imports in ensure_assets_registered(). If modules are reloaded (e.g., during development/testing), the _assets_registered flag may not reflect reality. This is a pre-existing pattern in the codebase, not introduced by this PR.

---

📋 Test Coverage Assessment

The added tests appropriately cover:

• ✅ Registry lookup for registered types (IsAnchor, PositionLimits)
• ✅ None return for unregistered types (AT_POSE, IN)
• ✅ Integration with the existing ArenaEnvGraphSpec test case

Consider adding (in a follow-up PR):

• Unit test for duplicate registration warning behavior
• Unit test for the decorator on a class missing the name attribute

---

📌 Verdict

Approve with suggestions. This is a well-structured addition that follows established patterns and includes appropriate tests. The suggestions above are improvements rather than blockers.

The TODO comment for AT_POSE and IN support is appropriately documented for future work.

---

_{🤖 Automated review by Isaac Lab Review Bot | Guidelines}

---

Update (6a83dbe): New commits add lightwheel_utils.py with acquire_lightwheel_asset() for scoped timeout and retry handling of Lightwheel asset downloads. This is a good addition that addresses download reliability concerns. The implementation is clean with proper timeout restoration in the finally block and the _looks_like_timeout() helper correctly traverses exception chains. Test coverage is adequate. No new issues found — previous suggestions remain valid but are non-blocking. ✅

---

Update (70ff4c2): Significant expansion — this commit adds a TaskRegistry alongside the existing ObjectRelationLibraryRegistry, plus a @register_task decorator. All 14 concrete task classes (NoTask, PickAndPlaceTask, LiftObjectTask, LiftObjectTaskRL, DexsuiteLiftTask, GoalPoseTask, PressButtonTask, RotateRevoluteJointTask, CloseDoorTask, OpenDoorTask, TurnKnobTask, AssemblyTask, SortMultiObjectTask, PlaceUprightTask) are now decorated and auto-registered.

Key observations:

1. ✅ Consistent pattern: TaskRegistry follows the same singleton approach as other registries. The @register_task decorator uses cls.__name__ as the key (appropriate since YAML configs use type: PascalCase), unlike @register_object_relation which uses cls.name.

2. ✅ Good test coverage: test_task_registry.py verifies all expected tasks can be resolved by name.

3. ✅ Proper lazy loading integration: ObjectRelationLibraryRegistry and TaskRegistry are added to the isinstance() checks in is_registered(), get_component_by_name(), and get_all_keys() — ensuring ensure_assets_registered() fires before lookups.

4. ⚠️ Same print() warning pattern: The @register_task decorator uses the same print(f"WARNING: ...") approach noted earlier. Consider using warnings.warn() for consistency.

5. ℹ️ PR title outdated: The PR is titled "Add ObjectRelation Registry" but now also adds TaskRegistry. Consider updating the title/description to reflect the broader scope.

No blocking issues. The expansion is well-structured and follows established conventions. ✅

---

Update (820ae6d): Added re-entry protection to ensure_assets_registered() using _registration_in_progress flag. This prevents circular import issues that could occur when registration decorators call is_registered() mid-import (which would trigger another ensure_assets_registered() call before the first one completes). The fix is clean with proper try/finally to ensure the flag is cleared. No new issues. ✅

---

Update (5df5010): Code cleanup and refactoring. Key changes:

1. ✅ Registry isinstance checks refactored: Replaced repeated isinstance(self, (AssetRegistry, DeviceRegistry, RetargeterRegistry, PolicyRegistry, HDRImageRegistry)) tuples with a single REGISTRIES constant. This is cleaner and ensures consistency when new registries are added.

2. ✅ REGISTRIES tuple excludes EnvironmentRegistry: The docstring correctly explains the exclusion — triggering ensure_assets_registered() during environment registration causes an env↔tasks cycle. Good design decision.

3. ✅ tasks/__init__.py explicit imports: All concrete task modules are now explicitly imported in the package __init__.py so decorators fire when the package is loaded. This ensures tasks are registered when ensure_assets_registered() imports isaaclab_arena.tasks.

4. ✅ Imports cleaned up: ObjectRelationLibraryRegistry and TaskRegistry properly added to register.py imports.

5. ℹ️ Sequential tasks not registered: sequential_task_base is imported but SequentialTaskBase does not appear in the expected list in test_task_registry.py. This appears intentional — it's a base class, not a concrete task.

No new issues introduced. The refactoring improves maintainability. ✅

[greptile-apps]

Greptile Summary

This PR adds ObjectRelationLibraryRegistry and TaskRegistry singletons backed by @register_object_relation and @register_task decorators, centralising the registration pattern already used for assets, devices, and policies. A _registration_in_progress re-entrancy guard and a new ensure_loaded flag on is_registered prevent circular import / segfault during decorator execution, and a REGISTRIES tuple replaces the hard-coded isinstance tuple in the three lazy-loading guard sites.

• Adds ObjectRelationLibraryRegistry and TaskRegistry to registries.py, updates REGISTRIES, and adds an ensure_assets_registered re-entrancy guard; all existing register_* decorators are updated to pass ensure_loaded=False.
• Adds register_object_relation and register_task decorators in register.py; annotates all concrete RelationBase/TaskBase subclasses and adds a name attribute where missing.
• Introduces task_library.py as the lazy-load entry point for task modules, adds relation_class_for_spatial_constraint_type utility, and extends tests for both new registries.

Confidence Score: 5/5

The change is safe to merge; the re-entrancy guard, lazy-loading refactor, and new registry classes are well-reasoned and covered by tests.

The registration mechanism correctly threads ensure_loaded=False through all decorator call sites, the _registration_in_progress guard prevents the circular-import/segfault that motivated the PR, and the REGISTRIES tuple cleanly replaces the three hard-coded isinstance checks. The only finding is a dead import of sequential_task_base in task_library.py, which has no runtime impact.

isaaclab_arena/tasks/task_library.py — the sequential_task_base import registers nothing and can be removed.

Important Files Changed

Filename	Overview
isaaclab_arena/assets/registries.py	Core change: adds ObjectRelationLibraryRegistry, TaskRegistry, REGISTRIES tuple, ensure_loaded param, and _registration_in_progress re-entrancy guard. Logic is sound.
isaaclab_arena/assets/register.py	Adds register_object_relation and register_task decorators; updates all existing decorators to pass ensure_loaded=False. Pattern is consistent with existing decorators.
isaaclab_arena/tasks/task_library.py	New lazy-load entry point that imports all concrete task modules. Includes sequential_task_base which has no @register_task decorator, making that import a no-op from a registration perspective.
isaaclab_arena/environments/graph_spec_utils.py	Adds relation_class_for_spatial_constraint_type utility; correctly uses is_registered with default ensure_loaded=True so lazy loading fires before the lookup.
isaaclab_arena/relations/relations.py	Adds @register_object_relation and a name attribute to all concrete RelationBase subclasses. Clean and consistent.
isaaclab_arena/tests/test_task_registry.py	New test covering TaskRegistry resolution of all 14 concrete task classes; correctly uses run_simulation_app_function to avoid segfault outside SimulationApp.
isaaclab_arena/tests/test_arena_env_graph_spec.py	Adds tests for relation_class_for_spatial_constraint_type and a sync-check between registered relations and the spatial-constraint enum.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A["Caller: get_task_by_name / get_object_relation_by_name\nor is_registered(ensure_loaded=True)"] --> B{{"_assets_registered?"}}
    B -- Yes --> Z["Return from registry"]
    B -- No --> C{{"_registration_in_progress?"}}
    C -- Yes --> D["Return early (re-entrant call)"]
    C -- No --> E["Set _registration_in_progress = True"]
    E --> F["Import library modules\n(relations.py, task_library.py, ...)"]
    F --> G["@register_object_relation / @register_task\ndecorators fire during import"]
    G --> H["is_registered(ensure_loaded=False)\n→ no re-entrancy, just checks _components"]
    H -- Not duplicate --> I["registry.register(cls, cls.name)"]
    H -- Duplicate --> J["print WARNING, skip"]
    I --> K["Set _assets_registered = True\n(finally: _registration_in_progress = False)"]
    K --> Z

Loading

_{Reviews (6): Last reviewed commit: "fix import" | Re-trigger Greptile}

[alexmillane]

Looks great. Thanks for fixing an long-standing issue while doing this!

A few small comments. One requested test addition. Merge when you think you've addressed those.

[alexmillane]

It's not your code (it's my bad), but shall we take this chance to clean this up and define

REGISTRIES = (
    AssetRegistry,
    DeviceRegistry,
    RetargeterRegistry,
    PolicyRegistry,
    HDRImageRegistry,
    ObjectRelationLibraryRegistry,
    TaskRegistry,
)

then

        if isinstance(self, REGISTRIES):
            ensure_assets_registered()

[alexmillane]

Nice. This issue has been hanging around since Arena day 1.

Thanks for fixing it.

[alexmillane]

👍

[alexmillane]

We have these two objects that mirror each other:

• subclasses of RelationBase, that are registered with the RelationRegistry
• members of the ArenaEnvGraphSpatialConstraintType enum

Right now we have this constraint that all registered relations should have a corresponding enum item with the same name.

I feel it makes sense to test that that condition holds true over time. Could you add a test for that (obviously we'll have to exclude (manually) at_pose and in (for now)).