From 0cf1f5fcaf10a9f8f9e8b5e5ced8b6ea40e1d1b7 Mon Sep 17 00:00:00 2001 From: Peter Du Date: Tue, 2 Jun 2026 11:18:07 -0700 Subject: [PATCH 01/10] add fine grained subtask data class and state machine --- .../environments/arena_env_builder.py | 2 + isaaclab_arena/tasks/composite_task_base.py | 38 + .../tasks/fine_grained_state_machine.py | 370 +++++++++ isaaclab_arena/tasks/fine_grained_subtask.py | 174 ++++ isaaclab_arena/tasks/task_base.py | 27 + .../tests/test_fine_grained_subtask.py | 756 ++++++++++++++++++ 6 files changed, 1367 insertions(+) create mode 100644 isaaclab_arena/tasks/fine_grained_state_machine.py create mode 100644 isaaclab_arena/tasks/fine_grained_subtask.py create mode 100644 isaaclab_arena/tests/test_fine_grained_subtask.py diff --git a/isaaclab_arena/environments/arena_env_builder.py b/isaaclab_arena/environments/arena_env_builder.py index 26e6722cb..5c237ac10 100644 --- a/isaaclab_arena/environments/arena_env_builder.py +++ b/isaaclab_arena/environments/arena_env_builder.py @@ -125,12 +125,14 @@ def compose_manager_cfg(self) -> IsaacLabArenaManagerBasedRLEnvCfg: self.arena_env.scene.get_events_cfg(), task.get_events_cfg(), placement_event_cfg, + task.get_fine_grained_subtask_events_cfg(), ) termination_cfg = combine_configclass_instances( "TerminationCfg", task.get_termination_cfg(), self.arena_env.scene.get_termination_cfg(), embodiment.get_termination_cfg(), + task.get_fine_grained_subtask_termination_cfg(), ) actions_cfg = embodiment.get_action_cfg() xr_cfg = embodiment.get_xr_cfg() diff --git a/isaaclab_arena/tasks/composite_task_base.py b/isaaclab_arena/tasks/composite_task_base.py index 2ce5182c2..df6dec41e 100644 --- a/isaaclab_arena/tasks/composite_task_base.py +++ b/isaaclab_arena/tasks/composite_task_base.py @@ -4,6 +4,7 @@ # SPDX-License-Identifier: Apache-2.0 import copy +import dataclasses import numpy as np import torch import warnings @@ -20,6 +21,7 @@ from isaaclab_arena.metrics.metric_base import MetricBase from isaaclab_arena.metrics.metric_term_cfg import MetricTermCfg from isaaclab_arena.tasks.common.mimic_default_params import MIMIC_DATAGEN_CONFIG_DEFAULTS +from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask from isaaclab_arena.tasks.task_base import TaskBase from isaaclab_arena.utils.configclass import ( check_configclass_field_duplicates, @@ -360,6 +362,42 @@ def get_metrics(self) -> list[MetricBase]: return subtask_metrics + def get_own_fine_grained_subtasks(self) -> list[FineGrainedSubtask]: + """Composite-level fine-grained recipes (e.g. cross-child conditions). + + These are added on top of whatever recipes the children declare. They + carry ``parent_subtask_idx = None`` and therefore have no activation + gating — they remain active for the full episode. + """ + return [] + + def get_fine_grained_subtasks(self) -> list[FineGrainedSubtask]: + """Concatenate children's fine-grained recipes with namespace prefixes. + + Each child's recipe gets a new name (``f"subtask_{i}/{original_name}"``) + and a ``parent_subtask_idx = i`` tag. For ``SequentialTaskBase`` the + state machine reads ``parent_subtask_idx`` to gate advancement against + ``env._current_subtask_idx`` so a child's recipes only advance while + that child is the active parent subtask. For unordered + ``CompositeTaskBase`` (no ``_current_subtask_idx``) the gate is a no-op + and every recipe is always active. + + Composite-level recipes from ``get_own_fine_grained_subtasks`` are + appended afterward and are not gated. + """ + out: list[FineGrainedSubtask] = [] + for i, child in enumerate(self.subtasks): + for fgs in child.get_fine_grained_subtasks(): + out.append( + dataclasses.replace( + fgs, + name=f"subtask_{i}/{fgs.name}", + parent_subtask_idx=i, + ) + ) + out.extend(self.get_own_fine_grained_subtasks()) + return out + def _validate_consistent_mimic_eef_names(self, arm_mode: ArmMode) -> set[str]: "Check that all subtasks have the same Mimic eef_names." mimic_eef_names = set(self.subtasks[0].get_mimic_env_cfg(arm_mode).subtask_configs.keys()) diff --git a/isaaclab_arena/tasks/fine_grained_state_machine.py b/isaaclab_arena/tasks/fine_grained_state_machine.py new file mode 100644 index 000000000..b99f5948b --- /dev/null +++ b/isaaclab_arena/tasks/fine_grained_state_machine.py @@ -0,0 +1,370 @@ +# Copyright (c) 2025-2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import torch +from dataclasses import MISSING +from typing import Any + +from isaaclab.managers import EventTermCfg, TerminationTermCfg +from isaaclab.utils import configclass + +from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask + +_STATE_MACHINE_ATTR = "_fine_grained_subtask_state_machine" + + +def _predicate_repr(pred) -> str: + """Generate human-readable string representation for a predicate.""" + + fn = getattr(pred, "func", pred) + name = getattr(fn, "__name__", repr(fn)) + kwargs = getattr(pred, "keywords", None) or {} + args = getattr(pred, "args", ()) or () + parts = [repr(a) for a in args] + for key, value in kwargs.items(): + if isinstance(value, (str, int, float, bool)): + parts.append(f"{key}={value!r}") + return f"{name}({', '.join(parts)})" if parts else name + + +class FineGrainedSubtaskRunner: + """State machine runner for a single FineGrainedSubtask object. + + Each runner is responsible for tracking the progress of all predicate_groups + within a FineGrainedSubtask object across all parallelenvironments. + """ + + def __init__(self, fine_grained_subtask: FineGrainedSubtask, num_envs: int, device): + self.fine_grained_subtask = fine_grained_subtask + self.num_envs = num_envs + self.device = device + + # Initialize the state machine's internal state. + self.current_index: dict[str, torch.Tensor] = {} + self.group_score: dict[str, torch.Tensor] = {} + self.group_complete: dict[str, torch.Tensor] = {} + + for group_name in fine_grained_subtask.group_names: + self.current_index[group_name] = torch.zeros(num_envs, dtype=torch.long, device=device) + self.group_score[group_name] = torch.zeros(num_envs, dtype=torch.float32, device=device) + self.group_complete[group_name] = torch.zeros(num_envs, dtype=torch.bool, device=device) + + def _compute_composite_task_gating_mask(self, env) -> torch.Tensor: + """Per-env mask of whether the FineGrainedSubtask is active. + + The gating is used to determine when tracking of predicates should + be active for composite tasks. + """ + + # If no parent_subtask_idx -> always active (returns all True). + if self.fine_grained_subtask.parent_subtask_idx is None: + return torch.ones(self.num_envs, dtype=torch.bool, device=self.device) + + # If no env._current_subtask_idx -> composite task is not sequential (returns all True). + current_idx = getattr(env, "_current_subtask_idx", None) + if current_idx is None: + return torch.ones(self.num_envs, dtype=torch.bool, device=self.device) + + # Otherwise return True only for envs whose current + # parent-subtask index matches this FineGrainedSubtask's parent_subtask_idx. + if torch.is_tensor(current_idx): + ci = current_idx.to(self.device) + else: + ci = torch.as_tensor(current_idx, device=self.device) + return ci == int(self.fine_grained_subtask.parent_subtask_idx) + + def step(self, env, step_index: torch.Tensor | None) -> list[dict]: + """Step the state machine runner for a single env.step. + + Check each group's current predicate, move the state machine to the next predicate if + the current predicate is True. Emit an event for each env where a predicate was advanced. + """ + + # List of state transition events (events are emitted for an env when a predicate flips True) + events: list[dict] = [] + + # If the FineGrainedSubtask is not active for the composite task, return. + composite_task_gating_mask = self._compute_composite_task_gating_mask(env) + if not bool(composite_task_gating_mask.any().item()): + return events + + # Step through each group of the FineGrainedSubtask. + for group_name, predicate_chain in self.fine_grained_subtask.canonical_predicate_groups.items(): + chain_length = len(predicate_chain) + # Mask for which envs have advanced this step. + advanced = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device) + + for chain_idx, (predicate, score_weight) in enumerate(predicate_chain): + # Compute mask for which envs that should evaluate the predicate. + # Envs should only be evaluated if: + # 1) They are at the current predicate position + # 2) They have not yet advanced this step + # 3) The FineGrainedSubtask is active for the composite task + at_position = (self.current_index[group_name] == chain_idx) & ~advanced & composite_task_gating_mask + if not bool(at_position.any().item()): + continue + + # Evaluate the predicate for all envs, reshaped to a flat (num_envs,) bool tensor. + result = torch.as_tensor(predicate(env), dtype=torch.bool, device=self.device).reshape(-1) + if result.shape[0] != self.num_envs: + raise RuntimeError( + f"Predicate {_predicate_repr(predicate)} returned shape {tuple(result.shape)};" + f" expected ({self.num_envs},)" + ) + + # Compute mask for which envs need to be advanced to the next predicate. + advance_mask = at_position & result + if not bool(advance_mask.any().item()): + continue + + # Advance the state machine to the next predicates. + self.current_index[group_name] = torch.where( + advance_mask, + self.current_index[group_name] + 1, + self.current_index[group_name], + ) + # Update the group score for the envs that were advanced. + self.group_score[group_name] = self.group_score[group_name] + advance_mask.float() * float(score_weight) + # Update the advanced mask for the envs that were advanced. + advanced = advanced | advance_mask + + # Emit an event for each env where a predicate was advanced. + pred_name = _predicate_repr(predicate) + for eid in torch.nonzero(advance_mask, as_tuple=False).flatten().tolist(): + events.append({ + "env_idx": int(eid), + "step": int(step_index[eid].item()) if step_index is not None else -1, + "fine_grained_subtask": self.fine_grained_subtask.name, + "group": group_name, + "predicate_index": chain_idx, + "predicate_name": pred_name, + "score_delta": float(score_weight), + }) + + # Update the group complete mask for the envs that have completed the group. + self.group_complete[group_name] = self.current_index[group_name] >= chain_length + + return events + + def reset(self, env_ids) -> None: + """Reset the state machine runner for the provided envs.""" + + for group_name in self.fine_grained_subtask.group_names: + for eid in env_ids: + self.current_index[group_name][eid] = 0 + self.group_score[group_name][eid] = 0.0 + self.group_complete[group_name][eid] = False + + def is_complete(self) -> torch.Tensor: + """Check if the FineGrainedSubtask is complete for all envs.""" + + groups = self.fine_grained_subtask.group_names + stacked = torch.stack([self.group_complete[g] for g in groups], dim=1) + if self.fine_grained_subtask.logical == "all": + return stacked.all(dim=1) + if self.fine_grained_subtask.logical == "any": + return stacked.any(dim=1) + return stacked.sum(dim=1) >= int(self.fine_grained_subtask.K or 1) + + def overall_score_per_env(self) -> torch.Tensor: + """Compute mean group score within this FineGrainedSubtask (in [0, 1]).""" + + groups = self.fine_grained_subtask.group_names + stacked = torch.stack([self.group_score[g] for g in groups], dim=1) + return stacked.mean(dim=1) + + +class FineGrainedStateMachine: + """State machine that manages runners for all FineGrainedSubtasks. + + Attributes: + fine_grained_subtasks: List of FineGrainedSubtasks to manage. + num_envs: Number of parallelenvironments. + device: Device to manage the state machine on. + runners: List of runners for each FineGrainedSubtask. + _events: List of events for each environment. + """ + + def __init__(self, fine_grained_subtasks: list[FineGrainedSubtask], num_envs: int, device): + self.fine_grained_subtasks = fine_grained_subtasks + self.num_envs = num_envs + self.device = device + self.runners = [FineGrainedSubtaskRunner(s, num_envs, device) for s in fine_grained_subtasks] + self._events: list[list[dict]] = [[] for _ in range(num_envs)] + + def step(self, env, step_index: torch.Tensor | None) -> None: + """Step each runner for a single env.step.""" + + for runner in self.runners: + for event in runner.step(env, step_index): + eid = event.pop("env_idx") + self._events[eid].append(event) + + def reset(self, env_ids) -> None: + """Reset the runners for the provided envs.""" + + for runner in self.runners: + runner.reset(env_ids) + for eid in env_ids: + self._events[eid] = [] + + def get_state(self) -> list[dict]: + """Get the state of each FineGrainedSubtask for all envs.""" + + output: list[dict] = [] + for env_idx in range(self.num_envs): + # Build a per-env dict from each runner's state. + fine_grained_subtask_states: dict[str, dict] = {} + overall_score = 0.0 + all_complete = True + + for runner in self.runners: + fine_grained_subtask = runner.fine_grained_subtask + completed_groups = 0 + total_groups = len(fine_grained_subtask.group_names) + active_predicates: dict[str, str | None] = {} + + # Compute the active predicates and completed groups. + for group_name in fine_grained_subtask.group_names: + cur_group_index = int(runner.current_index[group_name][env_idx].item()) + predicate_chain = fine_grained_subtask.canonical_predicate_groups[group_name] + if cur_group_index >= len(predicate_chain): + active_predicates[group_name] = None + completed_groups += 1 + else: + active_predicates[group_name] = _predicate_repr(predicate_chain[cur_group_index][0]) + + # Compute the overall score and completeness. + fine_grained_subtask_score = float(runner.overall_score_per_env()[env_idx].item()) + is_complete = bool(runner.is_complete()[env_idx].item()) + fine_grained_subtask_states[fine_grained_subtask.name] = { + "completed_groups": completed_groups, + "total_groups": total_groups, + "score": fine_grained_subtask_score, + "is_complete": is_complete, + "active_predicates": active_predicates, + } + overall_score += fine_grained_subtask.score * fine_grained_subtask_score + all_complete = all_complete and is_complete + + # Add the per-env state dict to the output. + output.append({ + "fine_grained_subtasks": fine_grained_subtask_states, + "overall_score": overall_score, + "all_complete": all_complete, + }) + return output + + def get_events(self) -> list[list[dict]]: + """Get all events for all envs.""" + + return [list(e) for e in self._events] + + +def _ensure_state_machine(env, fine_grained_subtasks: list[FineGrainedSubtask]) -> FineGrainedStateMachine: + """Return the env's FineGrainedStateMachine, lazily creating and caching it on first call.""" + + sm: FineGrainedStateMachine | None = getattr(env, _STATE_MACHINE_ATTR, None) + if sm is None: + sm = FineGrainedStateMachine( + fine_grained_subtasks=fine_grained_subtasks, num_envs=env.num_envs, device=env.device + ) + setattr(env, _STATE_MACHINE_ATTR, sm) + return sm + + +def fine_grained_subtask_step_func(env, fine_grained_subtasks: list[FineGrainedSubtask]) -> torch.Tensor: + """Termination-term entry point. + + Ticks the state machine, writes events and states to env.extras["fine_grained_subtask"], + and returns all-False so it does not contribute to termination. + """ + + sm = _ensure_state_machine(env, fine_grained_subtasks) + step_index = getattr(env, "episode_length_buf", None) + sm.step(env, step_index=step_index) + + """ + User-facing event/state information format: + + env.extras["fine_grained_subtask"] = { + "states": [ + { + "fine_grained_subtasks": { + "": { + "completed_groups": int, + "total_groups": int, + "score": float, # 0..1, normalized within subtask + "is_complete": bool, + "active_predicates": {group: str | None}, + }, + ... + }, + "overall_score": float, # weighted by FineGrainedSubtask.score + "all_complete": bool, + }, + ... # one entry per env + ], + "events": [ + [{"step": int, "fine_grained_subtask": str, "group": str, + "predicate_index": int, "predicate_name": str, + "score_delta": float}, ...], + ... # one list per env + ], + } + """ + + env.extras["fine_grained_subtask"] = { + "states": sm.get_state(), + "events": sm.get_events(), + } + + # Return all-False so it does not contribute to termination. + return torch.zeros(env.num_envs, dtype=torch.bool, device=env.device) + + +def fine_grained_subtask_reset_func(env, env_ids, fine_grained_subtasks: list[FineGrainedSubtask]) -> None: + """Reset-event entry point. + + Resets the state machine whenever the Lab env is reset. + """ + + sm = _ensure_state_machine(env, fine_grained_subtasks) + if env_ids is None: + env_ids = list(range(env.num_envs)) + elif torch.is_tensor(env_ids): + env_ids = env_ids.tolist() + sm.reset(env_ids) + + +@configclass +class FineGrainedSubtaskEventsCfg: + reset_fine_grained_subtasks: EventTermCfg = MISSING + + +@configclass +class FineGrainedSubtaskTerminationsCfg: + fine_grained_subtask_step: TerminationTermCfg = MISSING + + +def make_fine_grained_subtask_events_cfg(fine_grained_subtasks: list[FineGrainedSubtask]) -> Any: + return FineGrainedSubtaskEventsCfg( + reset_fine_grained_subtasks=EventTermCfg( + func=fine_grained_subtask_reset_func, + mode="reset", + params={"fine_grained_subtasks": fine_grained_subtasks}, + ) + ) + + +def make_fine_grained_subtask_termination_cfg(fine_grained_subtasks: list[FineGrainedSubtask]) -> Any: + return FineGrainedSubtaskTerminationsCfg( + fine_grained_subtask_step=TerminationTermCfg( + func=fine_grained_subtask_step_func, + params={"fine_grained_subtasks": fine_grained_subtasks}, + ) + ) diff --git a/isaaclab_arena/tasks/fine_grained_subtask.py b/isaaclab_arena/tasks/fine_grained_subtask.py new file mode 100644 index 000000000..eea31a80c --- /dev/null +++ b/isaaclab_arena/tasks/fine_grained_subtask.py @@ -0,0 +1,174 @@ +# Copyright (c) 2025-2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass, field +from typing import Literal, Union + +PredicateGroups = Union[ + Callable, + list[Callable], + list[tuple[Callable, float]], + dict[str, Callable], + dict[str, list[Callable]], + dict[str, list[tuple[Callable, float]]], +] + + +DEFAULT_GROUP_NAME = "default_group" + + +def format_predicate_groups(predicate_groups: PredicateGroups) -> dict[str, list[tuple[Callable, float]]]: + """Format predicate_groups into the canonical form. + + Canonical form: ``dict[group_name: list[(callable, score)]]``. + + Accepted input shapes: + 1. func (single callable) one group with one predicate + 2. [func, func, ...] one group, sequential chain + 3. [(func, score), ...] one group, sequential chain, weighted + 4. {group: func} multiple groups, one predicate each + 5. {group: [func, ...]} multiple groups, sequential chains + 6. {group: [(func, score), ...]} multiple groups, sequential chains, weighted + """ + + if callable(predicate_groups): + return {DEFAULT_GROUP_NAME: [(predicate_groups, 1.0)]} + + if isinstance(predicate_groups, list): + if len(predicate_groups) == 0: + raise ValueError("FineGrainedSubtask.predicate_groups list cannot be empty") + return {DEFAULT_GROUP_NAME: _format_group_chain(predicate_groups, group_name=DEFAULT_GROUP_NAME)} + + if isinstance(predicate_groups, dict): + if len(predicate_groups) == 0: + raise ValueError("FineGrainedSubtask.predicate_groups dict cannot be empty") + return { + group_name: _format_group_chain(value, group_name=group_name) + for group_name, value in predicate_groups.items() + } + + raise TypeError( + f"FineGrainedSubtask.predicate_groups must be a callable, list, or dict; got {type(predicate_groups).__name__}" + ) + + +def _format_group_chain(value, group_name: str) -> list[tuple[Callable, float]]: + if callable(value): + return [(value, 1.0)] + if not isinstance(value, list): + raise TypeError( + f"Predicate chain for group '{group_name}' must be a callable or a list; got {type(value).__name__}" + ) + if len(value) == 0: + raise ValueError(f"Predicate chain for group '{group_name}' cannot be empty") + + first = value[0] + if isinstance(first, tuple): + chain = [] + for i, item in enumerate(value): + if not (isinstance(item, tuple) and len(item) == 2): + raise TypeError(f"Group '{group_name}' index {i}: expected (callable, score) tuple, got {item!r}") + fn, score = item + if not callable(fn): + raise TypeError(f"Group '{group_name}' index {i}: first tuple element must be callable") + if not isinstance(score, (int, float)): + raise TypeError(f"Group '{group_name}' index {i}: score must be a number") + chain.append((fn, float(score))) + return chain + + if callable(first): + equal = 1.0 / len(value) + chain = [] + for i, fn in enumerate(value): + if not callable(fn): + raise TypeError(f"Group '{group_name}' index {i}: expected callable, got {type(fn).__name__}") + chain.append((fn, equal)) + return chain + + raise TypeError( + f"Group '{group_name}' elements must be callables or (callable, score) tuples; got {type(first).__name__}" + ) + + +def normalize_scores( + predicate_groups: dict[str, list[tuple[Callable, float]]], +) -> dict[str, list[tuple[Callable, float]]]: + """Scale each group's scores to sum to 1.0. Zero and negative-sum groups are left untouched.""" + + out: dict[str, list[tuple[Callable, float]]] = {} + for group, chain in predicate_groups.items(): + total = sum(score for _, score in chain) + if total <= 0: + out[group] = list(chain) + continue + out[group] = [(fn, score / total) for fn, score in chain] + return out + + +@dataclass +class FineGrainedSubtask: + """Configuration object that defines a scored predicate sequence to track progress within a task. + + A FineGrainedSubtask specifies what the predicate state machine should track. + Each FineGrainedSubtask holds one or more sequential predicate chains (groups). + Within a group, predicates run in order. Across groups, predicates run in parallel. + + Args: + name: Identifies the FineGrainedSubtask within the TaskBase. + predicate_groups: The sequential predicate chains that define the FineGrainedSubtask. + score: Weight of the FineGrainedSubtask in the TaskBase-level overall_score. + logical: How completed groups combine to determine if the FineGrainedSubtask is complete. + Can be "all", "any", or "choose" + K: Required when logical == "choose". Specifies the number of groups that must be completed + to consider the FineGrainedSubtask complete. + description: An optional description of the FineGrainedSubtask. + """ + + name: str + predicate_groups: PredicateGroups + score: float = 1.0 + logical: Literal["all", "any", "choose"] = "all" + K: int | None = None + description: str | None = None + + canonical_predicate_groups: dict[str, list[tuple[Callable, float]]] = field(init=False, repr=False) + + # Index of the parent TaskBase this recipe belongs to. Set automatically by + # CompositeTaskBase.get_fine_grained_subtasks() when used with composite tasks. + parent_subtask_idx: int | None = None + + def __post_init__(self): + if not (0.0 <= self.score <= 1.0): + raise ValueError(f"FineGrainedSubtask '{self.name}': score must be in [0, 1], got {self.score}") + if self.logical not in ("all", "any", "choose"): + raise ValueError( + f"FineGrainedSubtask '{self.name}': logical must be in ['all', 'any', 'choose'], got {self.logical}" + ) + + # Format the predicate groups into the canonical form and normalize the scores. + formatted = format_predicate_groups(self.predicate_groups) + normalized = normalize_scores(formatted) + self.canonical_predicate_groups = normalized + + # Validate the logical and K parameters. + num_groups = len(self.canonical_predicate_groups) + if self.logical == "choose": + if self.K is None: + raise ValueError(f"FineGrainedSubtask '{self.name}': K is required when logical='choose'") + if not (1 <= self.K <= num_groups): + raise ValueError(f"FineGrainedSubtask '{self.name}': K={self.K} but must be in [1, {num_groups}]") + + @property + def group_names(self) -> list[str]: + """Returns the names of the groups in the FineGrainedSubtask.""" + return list(self.canonical_predicate_groups.keys()) + + def get_chain(self, group_name: str) -> list[tuple[Callable, float]]: + """Returns the chain of predicates for a given group.""" + return self.canonical_predicate_groups[group_name] diff --git a/isaaclab_arena/tasks/task_base.py b/isaaclab_arena/tasks/task_base.py index 85d3a3ab2..3e4d5ef97 100644 --- a/isaaclab_arena/tasks/task_base.py +++ b/isaaclab_arena/tasks/task_base.py @@ -11,6 +11,11 @@ from isaaclab_arena.embodiments.common.arm_mode import ArmMode from isaaclab_arena.metrics.metric_base import MetricBase +from isaaclab_arena.tasks.fine_grained_state_machine import ( + make_fine_grained_subtask_events_cfg, + make_fine_grained_subtask_termination_cfg, +) +from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask class TaskBase(ABC): @@ -62,3 +67,25 @@ def get_episode_length_s(self) -> float | None: def get_task_description(self) -> str | None: return self.task_description + + def get_fine_grained_subtasks(self) -> list[FineGrainedSubtask]: + """Return per-step fine-grained subtasks the state machine should track. Default: none. + + Override in subclasses to opt in to fine-grained subtask tracking. When + the returned list is non-empty, the env builder will automatically + register a reset event and a per-step termination term that tick the + state machine and publish ``env.extras["fine_grained_subtask"]``. + """ + return [] + + def get_fine_grained_subtask_events_cfg(self) -> Any: + fine_grained_subtasks = self.get_fine_grained_subtasks() + if not fine_grained_subtasks: + return None + return make_fine_grained_subtask_events_cfg(fine_grained_subtasks) + + def get_fine_grained_subtask_termination_cfg(self) -> Any: + fine_grained_subtasks = self.get_fine_grained_subtasks() + if not fine_grained_subtasks: + return None + return make_fine_grained_subtask_termination_cfg(fine_grained_subtasks) diff --git a/isaaclab_arena/tests/test_fine_grained_subtask.py b/isaaclab_arena/tests/test_fine_grained_subtask.py new file mode 100644 index 000000000..13366b1c6 --- /dev/null +++ b/isaaclab_arena/tests/test_fine_grained_subtask.py @@ -0,0 +1,756 @@ +# Copyright (c) 2025-2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +import traceback + +from isaaclab_arena.tests.utils.subprocess import run_simulation_app_function + +HEADLESS = True + + +class _MockPredicate: + """Callable predicate that returns a controlled per-env bool tensor.""" + + def __init__(self, num_envs: int, name: str = "mock_predicate"): + import torch + + self.num_envs = num_envs + self.return_value = torch.tensor([False] * num_envs) + self.__name__ = name + + def set(self, values: list[bool]): + import torch + + assert len(values) == self.num_envs + self.return_value = torch.tensor(values) + + def __call__(self, env, **kwargs): + return self.return_value + + +class _MockEnv: + def __init__(self, num_envs: int = 1, device: str = "cpu"): + import torch + + self.num_envs = num_envs + self.device = device + self.extras = {} + self.episode_length_buf = torch.zeros(num_envs, dtype=torch.long) + + +def _advance_step(env, n: int = 1): + env.episode_length_buf = env.episode_length_buf + n + + +def _test_format_single_callable(simulation_app) -> bool: + """A bare predicate becomes a default-named group with weight 1.0.""" + from isaaclab_arena.tasks.fine_grained_subtask import DEFAULT_GROUP_NAME, FineGrainedSubtask + + try: + pred = _MockPredicate(num_envs=1) + fgs = FineGrainedSubtask(name="t", predicate_groups=pred) + assert fgs.group_names == [DEFAULT_GROUP_NAME] + chain = fgs.get_chain(DEFAULT_GROUP_NAME) + assert len(chain) == 1 + assert chain[0][0] is pred + assert abs(chain[0][1] - 1.0) < 1e-6 + except Exception as e: + print(f"Error: {e}") + traceback.print_exc() + return False + return True + + +def _test_format_list_of_callables(simulation_app) -> bool: + """A list of callables becomes a single group with normalized equal scores.""" + from isaaclab_arena.tasks.fine_grained_subtask import DEFAULT_GROUP_NAME, FineGrainedSubtask + + try: + preds = [_MockPredicate(num_envs=1, name=f"p{i}") for i in range(3)] + fgs = FineGrainedSubtask(name="t", predicate_groups=preds) + chain = fgs.get_chain(DEFAULT_GROUP_NAME) + assert [c[0] for c in chain] == preds + # Equal scores normalize to 1/3 each, summing to 1.0. + for _, score in chain: + assert abs(score - 1.0 / 3.0) < 1e-6 + assert abs(sum(s for _, s in chain) - 1.0) < 1e-6 + except Exception as e: + print(f"Error: {e}") + traceback.print_exc() + return False + return True + + +def _test_format_weighted_tuples(simulation_app) -> bool: + """Explicit (callable, score) tuples are normalized to sum to 1.0 within a group.""" + from isaaclab_arena.tasks.fine_grained_subtask import DEFAULT_GROUP_NAME, FineGrainedSubtask + + try: + p1 = _MockPredicate(num_envs=1, name="p1") + p2 = _MockPredicate(num_envs=1, name="p2") + fgs = FineGrainedSubtask(name="t", predicate_groups=[(p1, 1.0), (p2, 3.0)]) + chain = fgs.get_chain(DEFAULT_GROUP_NAME) + # 1.0/4.0 = 0.25, 3.0/4.0 = 0.75 + assert abs(chain[0][1] - 0.25) < 1e-6 + assert abs(chain[1][1] - 0.75) < 1e-6 + except Exception as e: + print(f"Error: {e}") + traceback.print_exc() + return False + return True + + +def _test_format_dict_groups(simulation_app) -> bool: + """Dict input gives one group per key; each group's scores are normalized independently.""" + from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask + + try: + p_a1 = _MockPredicate(num_envs=1, name="a1") + p_a2 = _MockPredicate(num_envs=1, name="a2") + p_b = _MockPredicate(num_envs=1, name="b") + fgs = FineGrainedSubtask( + name="t", + predicate_groups={ + "obj_a": [p_a1, p_a2], + "obj_b": p_b, + }, + logical="all", + ) + assert set(fgs.group_names) == {"obj_a", "obj_b"} + a_chain = fgs.get_chain("obj_a") + b_chain = fgs.get_chain("obj_b") + assert len(a_chain) == 2 + assert len(b_chain) == 1 + # obj_a's equal scores sum to 1.0. + assert abs(sum(s for _, s in a_chain) - 1.0) < 1e-6 + # obj_b's single-element group sums to 1.0. + assert abs(b_chain[0][1] - 1.0) < 1e-6 + except Exception as e: + print(f"Error: {e}") + traceback.print_exc() + return False + return True + + +def _test_format_rejects_invalid_inputs(simulation_app) -> bool: + """Empty containers and non-callable entries should raise.""" + from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask + + try: + for bad in ([], {}, 42, "string"): + try: + FineGrainedSubtask(name="t", predicate_groups=bad) + except (ValueError, TypeError): + continue + print(f"Expected error for input {bad!r}") + return False + # logical=choose without K should raise. + try: + FineGrainedSubtask( + name="t", + predicate_groups=_MockPredicate(num_envs=1), + logical="choose", + ) + except ValueError: + pass + else: + print("Expected ValueError for logical='choose' without K") + return False + except Exception as e: + print(f"Error: {e}") + traceback.print_exc() + return False + return True + + +def _test_state_machine_advances_sequentially(simulation_app) -> bool: + """A single subtask with a 3-predicate chain advances one step per satisfied predicate.""" + from isaaclab_arena.tasks.fine_grained_state_machine import FineGrainedStateMachine + from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask + + try: + env = _MockEnv(num_envs=1) + preds = [_MockPredicate(num_envs=1, name=f"p{i}") for i in range(3)] + fgs = FineGrainedSubtask(name="lift", predicate_groups=preds) + sm = FineGrainedStateMachine(fine_grained_subtasks=[fgs], num_envs=1, device="cpu") + sm.reset([0]) + + # Step 1: p0 fires; p1, p2 still False. Advance to index 1. + preds[0].set([True]) + _advance_step(env) + sm.step(env, step_index=env.episode_length_buf) + state = sm.get_state()[0]["fine_grained_subtasks"]["lift"] + assert state["completed_groups"] == 0 # 3-predicate chain not done until all 3 + assert not state["is_complete"] + events = sm.get_events()[0] + assert len(events) == 1 and events[0]["predicate_index"] == 0 + + # Step 2: p0 reverts (irrelevant; latched forward), p1 fires. + preds[0].set([False]) + preds[1].set([True]) + _advance_step(env) + sm.step(env, step_index=env.episode_length_buf) + events = sm.get_events()[0] + assert len(events) == 2 and events[-1]["predicate_index"] == 1 + + # Step 3: p2 fires; subtask complete. + preds[2].set([True]) + _advance_step(env) + sm.step(env, step_index=env.episode_length_buf) + state = sm.get_state()[0]["fine_grained_subtasks"]["lift"] + assert state["is_complete"] + assert state["completed_groups"] == 1 + assert abs(state["score"] - 1.0) < 1e-6 + events = sm.get_events()[0] + assert len(events) == 3 and events[-1]["predicate_index"] == 2 + except Exception as e: + print(f"Error: {e}") + traceback.print_exc() + return False + return True + + +def _test_state_machine_ignores_out_of_order_success(simulation_app) -> bool: + """If a later predicate fires first, it's ignored until preceding ones have advanced.""" + from isaaclab_arena.tasks.fine_grained_state_machine import FineGrainedStateMachine + from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask + + try: + env = _MockEnv(num_envs=1) + preds = [_MockPredicate(num_envs=1, name=f"p{i}") for i in range(3)] + fgs = FineGrainedSubtask(name="lift", predicate_groups=preds) + sm = FineGrainedStateMachine(fine_grained_subtasks=[fgs], num_envs=1, device="cpu") + sm.reset([0]) + + # p0 stays False; p1 and p2 fire — no progress should be made. + preds[0].set([False]) + preds[1].set([True]) + preds[2].set([True]) + _advance_step(env) + sm.step(env, step_index=env.episode_length_buf) + state = sm.get_state()[0]["fine_grained_subtasks"]["lift"] + assert state["completed_groups"] == 0 + assert not state["is_complete"] + assert state["score"] == 0.0 + assert len(sm.get_events()[0]) == 0 + + # Now p0 fires; chain catches up: 0, 1, 2 should all advance over subsequent + # steps (one per step, in order). + preds[0].set([True]) + for _ in range(3): + _advance_step(env) + sm.step(env, step_index=env.episode_length_buf) + state = sm.get_state()[0]["fine_grained_subtasks"]["lift"] + assert state["is_complete"] + assert state["completed_groups"] == 1 + except Exception as e: + print(f"Error: {e}") + traceback.print_exc() + return False + return True + + +def _test_state_machine_logical_any(simulation_app) -> bool: + """Two parallel groups with logical='any' complete as soon as either one finishes.""" + from isaaclab_arena.tasks.fine_grained_state_machine import FineGrainedStateMachine + from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask + + try: + env = _MockEnv(num_envs=1) + p_a = _MockPredicate(num_envs=1, name="a") + p_b = _MockPredicate(num_envs=1, name="b") + fgs = FineGrainedSubtask( + name="either", + predicate_groups={"a": p_a, "b": p_b}, + logical="any", + ) + sm = FineGrainedStateMachine(fine_grained_subtasks=[fgs], num_envs=1, device="cpu") + sm.reset([0]) + + # Neither group complete -> not done. + _advance_step(env) + sm.step(env, step_index=env.episode_length_buf) + assert not sm.get_state()[0]["fine_grained_subtasks"]["either"]["is_complete"] + + # Group "a" completes -> done. + p_a.set([True]) + _advance_step(env) + sm.step(env, step_index=env.episode_length_buf) + assert sm.get_state()[0]["fine_grained_subtasks"]["either"]["is_complete"] + except Exception as e: + print(f"Error: {e}") + traceback.print_exc() + return False + return True + + +def _test_state_machine_logical_all(simulation_app) -> bool: + """logical='all' requires every group to complete its chain.""" + from isaaclab_arena.tasks.fine_grained_state_machine import FineGrainedStateMachine + from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask + + try: + env = _MockEnv(num_envs=1) + p_a = _MockPredicate(num_envs=1, name="a") + p_b = _MockPredicate(num_envs=1, name="b") + fgs = FineGrainedSubtask( + name="both", + predicate_groups={"a": p_a, "b": p_b}, + logical="all", + ) + sm = FineGrainedStateMachine(fine_grained_subtasks=[fgs], num_envs=1, device="cpu") + sm.reset([0]) + + # Only "a" completes -> still not done. + p_a.set([True]) + _advance_step(env) + sm.step(env, step_index=env.episode_length_buf) + assert not sm.get_state()[0]["fine_grained_subtasks"]["both"]["is_complete"] + + # "b" also completes -> done. + p_b.set([True]) + _advance_step(env) + sm.step(env, step_index=env.episode_length_buf) + assert sm.get_state()[0]["fine_grained_subtasks"]["both"]["is_complete"] + except Exception as e: + print(f"Error: {e}") + traceback.print_exc() + return False + return True + + +def _test_state_machine_logical_choose(simulation_app) -> bool: + """logical='choose' with K=2 requires any two of three groups to complete.""" + from isaaclab_arena.tasks.fine_grained_state_machine import FineGrainedStateMachine + from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask + + try: + env = _MockEnv(num_envs=1) + p_a = _MockPredicate(num_envs=1, name="a") + p_b = _MockPredicate(num_envs=1, name="b") + p_c = _MockPredicate(num_envs=1, name="c") + fgs = FineGrainedSubtask( + name="any_two", + predicate_groups={"a": p_a, "b": p_b, "c": p_c}, + logical="choose", + K=2, + ) + sm = FineGrainedStateMachine(fine_grained_subtasks=[fgs], num_envs=1, device="cpu") + sm.reset([0]) + + # Only one group complete -> not done. + p_a.set([True]) + _advance_step(env) + sm.step(env, step_index=env.episode_length_buf) + assert not sm.get_state()[0]["fine_grained_subtasks"]["any_two"]["is_complete"] + + # Two complete -> done. + p_b.set([True]) + _advance_step(env) + sm.step(env, step_index=env.episode_length_buf) + assert sm.get_state()[0]["fine_grained_subtasks"]["any_two"]["is_complete"] + except Exception as e: + print(f"Error: {e}") + traceback.print_exc() + return False + return True + + +def _test_state_machine_reset_clears_state(simulation_app) -> bool: + """Resetting an env_id zeroes its progress and event log, but leaves other envs alone.""" + from isaaclab_arena.tasks.fine_grained_state_machine import FineGrainedStateMachine + from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask + + try: + env = _MockEnv(num_envs=2) + preds = [_MockPredicate(num_envs=2, name=f"p{i}") for i in range(2)] + fgs = FineGrainedSubtask(name="t", predicate_groups=preds) + sm = FineGrainedStateMachine(fine_grained_subtasks=[fgs], num_envs=2, device="cpu") + sm.reset([0, 1]) + + # Drive env 0 to fully complete, env 1 halfway. + preds[0].set([True, True]) + preds[1].set([True, False]) + _advance_step(env) + sm.step(env, step_index=env.episode_length_buf) + _advance_step(env) + sm.step(env, step_index=env.episode_length_buf) + + state = sm.get_state() + assert state[0]["fine_grained_subtasks"]["t"]["is_complete"] + assert not state[1]["fine_grained_subtasks"]["t"]["is_complete"] + assert len(sm.get_events()[0]) >= 2 + assert len(sm.get_events()[1]) >= 1 + + # Reset only env 0. + sm.reset([0]) + state = sm.get_state() + assert not state[0]["fine_grained_subtasks"]["t"]["is_complete"] + assert state[0]["fine_grained_subtasks"]["t"]["score"] == 0.0 + assert sm.get_events()[0] == [] + # env 1 untouched. + assert len(sm.get_events()[1]) >= 1 + except Exception as e: + print(f"Error: {e}") + traceback.print_exc() + return False + return True + + +def _test_gating_active_when_parent_subtask_idx_matches(simulation_app) -> bool: + """A recipe with ``parent_subtask_idx=N`` advances normally when the env's + ``_current_subtask_idx`` matches N. This mirrors the sequential-composite + case where subtask N is the currently-active parent subtask.""" + from isaaclab_arena.tasks.fine_grained_state_machine import FineGrainedStateMachine + from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask + + try: + env = _MockEnv(num_envs=1) + env._current_subtask_idx = [1] + + pred = _MockPredicate(num_envs=1, name="p") + fgs = FineGrainedSubtask(name="t", predicate_groups=pred, parent_subtask_idx=1) + sm = FineGrainedStateMachine(fine_grained_subtasks=[fgs], num_envs=1, device="cpu") + sm.reset([0]) + + pred.set([True]) + _advance_step(env) + sm.step(env, step_index=env.episode_length_buf) + assert sm.get_state()[0]["fine_grained_subtasks"]["t"]["is_complete"] + assert len(sm.get_events()[0]) == 1 + except Exception as e: + print(f"Error: {e}") + traceback.print_exc() + return False + return True + + +def _test_gating_blocked_when_parent_subtask_idx_mismatches(simulation_app) -> bool: + """A recipe with ``parent_subtask_idx=N`` is frozen when the env's + ``_current_subtask_idx`` differs from N — even if the underlying predicate + is True. Once the parent advances to N, the recipe starts advancing too.""" + from isaaclab_arena.tasks.fine_grained_state_machine import FineGrainedStateMachine + from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask + + try: + env = _MockEnv(num_envs=1) + env._current_subtask_idx = [0] # parent on subtask 0, recipe targets 1 + + pred = _MockPredicate(num_envs=1, name="p") + fgs = FineGrainedSubtask(name="t", predicate_groups=pred, parent_subtask_idx=1) + sm = FineGrainedStateMachine(fine_grained_subtasks=[fgs], num_envs=1, device="cpu") + sm.reset([0]) + + # Predicate True, but the parent isn't at this recipe's index yet. + pred.set([True]) + _advance_step(env) + sm.step(env, step_index=env.episode_length_buf) + assert not sm.get_state()[0]["fine_grained_subtasks"]["t"]["is_complete"] + assert sm.get_state()[0]["fine_grained_subtasks"]["t"]["score"] == 0.0 + assert len(sm.get_events()[0]) == 0 + + # Parent advances to this recipe's index; the recipe catches up. + env._current_subtask_idx = [1] + _advance_step(env) + sm.step(env, step_index=env.episode_length_buf) + assert sm.get_state()[0]["fine_grained_subtasks"]["t"]["is_complete"] + assert len(sm.get_events()[0]) == 1 + except Exception as e: + print(f"Error: {e}") + traceback.print_exc() + return False + return True + + +def _test_gating_noop_when_env_has_no_current_subtask_idx(simulation_app) -> bool: + """For unordered composite tasks, ``env._current_subtask_idx`` is absent. + Recipes carry ``parent_subtask_idx`` (for namespacing) but gating is a no-op + and all recipes advance whenever their predicates fire.""" + from isaaclab_arena.tasks.fine_grained_state_machine import FineGrainedStateMachine + from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask + + try: + env = _MockEnv(num_envs=1) + # No _current_subtask_idx attribute on env -> unordered composite path. + + pred = _MockPredicate(num_envs=1, name="p") + fgs = FineGrainedSubtask(name="t", predicate_groups=pred, parent_subtask_idx=1) + sm = FineGrainedStateMachine(fine_grained_subtasks=[fgs], num_envs=1, device="cpu") + sm.reset([0]) + + pred.set([True]) + _advance_step(env) + sm.step(env, step_index=env.episode_length_buf) + assert sm.get_state()[0]["fine_grained_subtasks"]["t"]["is_complete"] + assert len(sm.get_events()[0]) == 1 + except Exception as e: + print(f"Error: {e}") + traceback.print_exc() + return False + return True + + +def _test_sequential_gating_end_to_end(simulation_app) -> bool: + """Two recipes targeting different parent subtask indices. The parent's + ``_current_subtask_idx`` advances over time. Each recipe only progresses + during its parent's active window.""" + from isaaclab_arena.tasks.fine_grained_state_machine import FineGrainedStateMachine + from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask + + try: + env = _MockEnv(num_envs=1) + env._current_subtask_idx = [0] + + pred_a = _MockPredicate(num_envs=1, name="a") + pred_b = _MockPredicate(num_envs=1, name="b") + fgs_a = FineGrainedSubtask(name="a", predicate_groups=pred_a, parent_subtask_idx=0) + fgs_b = FineGrainedSubtask(name="b", predicate_groups=pred_b, parent_subtask_idx=1) + sm = FineGrainedStateMachine(fine_grained_subtasks=[fgs_a, fgs_b], num_envs=1, device="cpu") + sm.reset([0]) + + # Both predicates True, but only "a" is active (parent on subtask 0). + pred_a.set([True]) + pred_b.set([True]) + _advance_step(env) + sm.step(env, step_index=env.episode_length_buf) + assert sm.get_state()[0]["fine_grained_subtasks"]["a"]["is_complete"] + assert not sm.get_state()[0]["fine_grained_subtasks"]["b"]["is_complete"] + + # Parent advances to subtask 1; "b" is now active. + env._current_subtask_idx = [1] + _advance_step(env) + sm.step(env, step_index=env.episode_length_buf) + assert sm.get_state()[0]["fine_grained_subtasks"]["b"]["is_complete"] + except Exception as e: + print(f"Error: {e}") + traceback.print_exc() + return False + return True + + +def _test_step_func_publishes_to_extras_and_returns_no_termination(simulation_app) -> bool: + """fine_grained_subtask_step_func writes env.extras and returns all-False.""" + from isaaclab_arena.tasks.fine_grained_state_machine import ( + fine_grained_subtask_reset_func, + fine_grained_subtask_step_func, + ) + from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask + + try: + env = _MockEnv(num_envs=2) + pred = _MockPredicate(num_envs=2, name="p") + fgs = FineGrainedSubtask(name="t", predicate_groups=pred) + + fine_grained_subtask_reset_func(env, env_ids=[0, 1], fine_grained_subtasks=[fgs]) + + # Step with predicate False: state machine ticks but no transitions. + result = fine_grained_subtask_step_func(env, fine_grained_subtasks=[fgs]) + assert result.tolist() == [False, False] + assert "fine_grained_subtask" in env.extras + assert len(env.extras["fine_grained_subtask"]["states"]) == 2 + assert env.extras["fine_grained_subtask"]["events"] == [[], []] + assert not env.extras["fine_grained_subtask"]["states"][0]["fine_grained_subtasks"]["t"]["is_complete"] + + # Step with env 0 predicate True: env 0 completes, env 1 does not. + pred.set([True, False]) + _advance_step(env) + result = fine_grained_subtask_step_func(env, fine_grained_subtasks=[fgs]) + assert result.tolist() == [False, False] + states = env.extras["fine_grained_subtask"]["states"] + events = env.extras["fine_grained_subtask"]["events"] + assert states[0]["fine_grained_subtasks"]["t"]["is_complete"] + assert not states[1]["fine_grained_subtasks"]["t"]["is_complete"] + assert len(events[0]) == 1 + assert len(events[1]) == 0 + + # Reset env 0; its state should clear, env 1 untouched. We also flip pred + # to False to verify the post-reset step starts from the chain head rather + # than auto-re-completing. + pred.set([False, False]) + fine_grained_subtask_reset_func(env, env_ids=[0], fine_grained_subtasks=[fgs]) + result = fine_grained_subtask_step_func(env, fine_grained_subtasks=[fgs]) + states = env.extras["fine_grained_subtask"]["states"] + assert not states[0]["fine_grained_subtasks"]["t"]["is_complete"] + assert states[0]["fine_grained_subtasks"]["t"]["score"] == 0.0 + except Exception as e: + print(f"Error: {e}") + traceback.print_exc() + return False + return True + + +def _test_task_base_fine_grained_subtask_hooks(simulation_app) -> bool: + """TaskBase's fine-grained-subtask hooks: default is empty/None; overriding + ``get_fine_grained_subtasks`` causes the events/termination helpers to + return real cfgs that the env builder picks up automatically. + + Importing ``task_base`` is non-trivial in the test sandbox (it transitively + pulls in the asset-library network registration). Both default and opt-in + behavior are exercised inside a single test so the module import only + happens once. + """ + from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask + from isaaclab_arena.tasks.task_base import TaskBase + + try: + + class _Base(TaskBase): + def get_scene_cfg(self): + return None + + def get_termination_cfg(self): + return None + + def get_events_cfg(self): + return None + + def get_mimic_env_cfg(self, arm_mode): + return None + + def get_metrics(self): + return [] + + default_task = _Base() + assert default_task.get_fine_grained_subtasks() == [] + assert default_task.get_fine_grained_subtask_events_cfg() is None + assert default_task.get_fine_grained_subtask_termination_cfg() is None + + class _OptIn(_Base): + def get_fine_grained_subtasks(self): + pred = _MockPredicate(num_envs=1, name="p") + return [FineGrainedSubtask(name="lift", predicate_groups=pred)] + + opt_in = _OptIn() + assert len(opt_in.get_fine_grained_subtasks()) == 1 + assert opt_in.get_fine_grained_subtask_events_cfg() is not None + assert opt_in.get_fine_grained_subtask_termination_cfg() is not None + + # ---- CompositeTaskBase: concatenate child recipes with namespace + parent index ---- + from isaaclab_arena.tasks.composite_task_base import CompositeTaskBase + + class _ChildA(_Base): + def get_fine_grained_subtasks(self): + return [FineGrainedSubtask(name="open", predicate_groups=_MockPredicate(1, name="pa"))] + + class _ChildB(_Base): + def get_fine_grained_subtasks(self): + return [FineGrainedSubtask(name="close", predicate_groups=_MockPredicate(1, name="pb"))] + + composite = CompositeTaskBase(subtasks=[_ChildA(), _ChildB()]) + recipes = composite.get_fine_grained_subtasks() + assert len(recipes) == 2 + assert recipes[0].name == "subtask_0/open" + assert recipes[0].parent_subtask_idx == 0 + assert recipes[1].name == "subtask_1/close" + assert recipes[1].parent_subtask_idx == 1 + + # get_own_fine_grained_subtasks adds composite-level recipes (no gating). + class _CompositeWithOwn(CompositeTaskBase): + def get_own_fine_grained_subtasks(self): + return [FineGrainedSubtask(name="both_done", predicate_groups=_MockPredicate(1, name="own"))] + + composite2 = _CompositeWithOwn(subtasks=[_ChildA(), _ChildB()]) + recipes2 = composite2.get_fine_grained_subtasks() + assert len(recipes2) == 3 + assert recipes2[2].name == "both_done" + assert recipes2[2].parent_subtask_idx is None # composite-level: never gated + except Exception as e: + print(f"Error: {e}") + traceback.print_exc() + return False + return True + + +# Pytest entry points ----------------------------------------------------------- + + +def test_format_single_callable(): + assert run_simulation_app_function(_test_format_single_callable, headless=HEADLESS) + + +def test_format_list_of_callables(): + assert run_simulation_app_function(_test_format_list_of_callables, headless=HEADLESS) + + +def test_format_weighted_tuples(): + assert run_simulation_app_function(_test_format_weighted_tuples, headless=HEADLESS) + + +def test_format_dict_groups(): + assert run_simulation_app_function(_test_format_dict_groups, headless=HEADLESS) + + +def test_format_rejects_invalid_inputs(): + assert run_simulation_app_function(_test_format_rejects_invalid_inputs, headless=HEADLESS) + + +def test_state_machine_advances_sequentially(): + assert run_simulation_app_function(_test_state_machine_advances_sequentially, headless=HEADLESS) + + +def test_state_machine_ignores_out_of_order_success(): + assert run_simulation_app_function(_test_state_machine_ignores_out_of_order_success, headless=HEADLESS) + + +def test_state_machine_logical_any(): + assert run_simulation_app_function(_test_state_machine_logical_any, headless=HEADLESS) + + +def test_state_machine_logical_all(): + assert run_simulation_app_function(_test_state_machine_logical_all, headless=HEADLESS) + + +def test_state_machine_logical_choose(): + assert run_simulation_app_function(_test_state_machine_logical_choose, headless=HEADLESS) + + +def test_state_machine_reset_clears_state(): + assert run_simulation_app_function(_test_state_machine_reset_clears_state, headless=HEADLESS) + + +def test_gating_active_when_parent_subtask_idx_matches(): + assert run_simulation_app_function(_test_gating_active_when_parent_subtask_idx_matches, headless=HEADLESS) + + +def test_gating_blocked_when_parent_subtask_idx_mismatches(): + assert run_simulation_app_function(_test_gating_blocked_when_parent_subtask_idx_mismatches, headless=HEADLESS) + + +def test_gating_noop_when_env_has_no_current_subtask_idx(): + assert run_simulation_app_function(_test_gating_noop_when_env_has_no_current_subtask_idx, headless=HEADLESS) + + +def test_sequential_gating_end_to_end(): + assert run_simulation_app_function(_test_sequential_gating_end_to_end, headless=HEADLESS) + + +def test_step_func_publishes_to_extras_and_returns_no_termination(): + assert run_simulation_app_function( + _test_step_func_publishes_to_extras_and_returns_no_termination, headless=HEADLESS + ) + + +def test_task_base_fine_grained_subtask_hooks(): + assert run_simulation_app_function(_test_task_base_fine_grained_subtask_hooks, headless=HEADLESS) + + +if __name__ == "__main__": + test_format_single_callable() + test_format_list_of_callables() + test_format_weighted_tuples() + test_format_dict_groups() + test_format_rejects_invalid_inputs() + test_state_machine_advances_sequentially() + test_state_machine_ignores_out_of_order_success() + test_state_machine_logical_any() + test_state_machine_logical_all() + test_state_machine_logical_choose() + test_state_machine_reset_clears_state() + test_gating_active_when_parent_subtask_idx_matches() + test_gating_blocked_when_parent_subtask_idx_mismatches() + test_gating_noop_when_env_has_no_current_subtask_idx() + test_sequential_gating_end_to_end() + test_step_func_publishes_to_extras_and_returns_no_termination() + test_task_base_fine_grained_subtask_hooks() From a510bf7ec5f3a51f6a8e5b26ffc749f435d4483f Mon Sep 17 00:00:00 2001 From: Peter Du Date: Wed, 3 Jun 2026 10:41:27 -0700 Subject: [PATCH 02/10] update fgs tests --- .../tests/test_fine_grained_subtask.py | 209 ++++++++---------- 1 file changed, 96 insertions(+), 113 deletions(-) diff --git a/isaaclab_arena/tests/test_fine_grained_subtask.py b/isaaclab_arena/tests/test_fine_grained_subtask.py index 13366b1c6..0450e11a5 100644 --- a/isaaclab_arena/tests/test_fine_grained_subtask.py +++ b/isaaclab_arena/tests/test_fine_grained_subtask.py @@ -9,6 +9,9 @@ HEADLESS = True +# Tolerance for floating-point score comparisons. +SCORE_TOL = 1e-6 + class _MockPredicate: """Callable predicate that returns a controlled per-env bool tensor.""" @@ -44,7 +47,7 @@ def _advance_step(env, n: int = 1): env.episode_length_buf = env.episode_length_buf + n -def _test_format_single_callable(simulation_app) -> bool: +def _test_predicate_groups_single_callable(simulation_app) -> bool: """A bare predicate becomes a default-named group with weight 1.0.""" from isaaclab_arena.tasks.fine_grained_subtask import DEFAULT_GROUP_NAME, FineGrainedSubtask @@ -55,7 +58,7 @@ def _test_format_single_callable(simulation_app) -> bool: chain = fgs.get_chain(DEFAULT_GROUP_NAME) assert len(chain) == 1 assert chain[0][0] is pred - assert abs(chain[0][1] - 1.0) < 1e-6 + assert abs(chain[0][1] - 1.0) < SCORE_TOL except Exception as e: print(f"Error: {e}") traceback.print_exc() @@ -63,7 +66,7 @@ def _test_format_single_callable(simulation_app) -> bool: return True -def _test_format_list_of_callables(simulation_app) -> bool: +def _test_predicate_groups_list_of_callables(simulation_app) -> bool: """A list of callables becomes a single group with normalized equal scores.""" from isaaclab_arena.tasks.fine_grained_subtask import DEFAULT_GROUP_NAME, FineGrainedSubtask @@ -72,10 +75,10 @@ def _test_format_list_of_callables(simulation_app) -> bool: fgs = FineGrainedSubtask(name="t", predicate_groups=preds) chain = fgs.get_chain(DEFAULT_GROUP_NAME) assert [c[0] for c in chain] == preds - # Equal scores normalize to 1/3 each, summing to 1.0. + # Equal scores normalize to 0.33 each, summing to 1.0. for _, score in chain: - assert abs(score - 1.0 / 3.0) < 1e-6 - assert abs(sum(s for _, s in chain) - 1.0) < 1e-6 + assert abs(score - 1.0 / 3.0) < SCORE_TOL + assert abs(sum(s for _, s in chain) - 1.0) < SCORE_TOL except Exception as e: print(f"Error: {e}") traceback.print_exc() @@ -83,7 +86,7 @@ def _test_format_list_of_callables(simulation_app) -> bool: return True -def _test_format_weighted_tuples(simulation_app) -> bool: +def _test_predicate_groups_weighted_tuples(simulation_app) -> bool: """Explicit (callable, score) tuples are normalized to sum to 1.0 within a group.""" from isaaclab_arena.tasks.fine_grained_subtask import DEFAULT_GROUP_NAME, FineGrainedSubtask @@ -93,8 +96,8 @@ def _test_format_weighted_tuples(simulation_app) -> bool: fgs = FineGrainedSubtask(name="t", predicate_groups=[(p1, 1.0), (p2, 3.0)]) chain = fgs.get_chain(DEFAULT_GROUP_NAME) # 1.0/4.0 = 0.25, 3.0/4.0 = 0.75 - assert abs(chain[0][1] - 0.25) < 1e-6 - assert abs(chain[1][1] - 0.75) < 1e-6 + assert abs(chain[0][1] - 0.25) < SCORE_TOL + assert abs(chain[1][1] - 0.75) < SCORE_TOL except Exception as e: print(f"Error: {e}") traceback.print_exc() @@ -102,8 +105,8 @@ def _test_format_weighted_tuples(simulation_app) -> bool: return True -def _test_format_dict_groups(simulation_app) -> bool: - """Dict input gives one group per key; each group's scores are normalized independently.""" +def _test_predicate_groups_dict_groups(simulation_app) -> bool: + """Dict input gives one group per key and each group's scores are normalized independently.""" from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask try: @@ -124,9 +127,9 @@ def _test_format_dict_groups(simulation_app) -> bool: assert len(a_chain) == 2 assert len(b_chain) == 1 # obj_a's equal scores sum to 1.0. - assert abs(sum(s for _, s in a_chain) - 1.0) < 1e-6 + assert abs(sum(s for _, s in a_chain) - 1.0) < SCORE_TOL # obj_b's single-element group sums to 1.0. - assert abs(b_chain[0][1] - 1.0) < 1e-6 + assert abs(b_chain[0][1] - 1.0) < SCORE_TOL except Exception as e: print(f"Error: {e}") traceback.print_exc() @@ -134,8 +137,8 @@ def _test_format_dict_groups(simulation_app) -> bool: return True -def _test_format_rejects_invalid_inputs(simulation_app) -> bool: - """Empty containers and non-callable entries should raise.""" +def _test_predicate_groups_rejects_invalid_inputs(simulation_app) -> bool: + """Empty containers and non-callable entries should raise error.""" from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask try: @@ -146,7 +149,7 @@ def _test_format_rejects_invalid_inputs(simulation_app) -> bool: continue print(f"Expected error for input {bad!r}") return False - # logical=choose without K should raise. + # logical=choose without K should raise error. try: FineGrainedSubtask( name="t", @@ -166,7 +169,7 @@ def _test_format_rejects_invalid_inputs(simulation_app) -> bool: def _test_state_machine_advances_sequentially(simulation_app) -> bool: - """A single subtask with a 3-predicate chain advances one step per satisfied predicate.""" + """A single FineGrainedSubtask with a 3 predicate chain advances one step per satisfied predicate.""" from isaaclab_arena.tasks.fine_grained_state_machine import FineGrainedStateMachine from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask @@ -177,7 +180,7 @@ def _test_state_machine_advances_sequentially(simulation_app) -> bool: sm = FineGrainedStateMachine(fine_grained_subtasks=[fgs], num_envs=1, device="cpu") sm.reset([0]) - # Step 1: p0 fires; p1, p2 still False. Advance to index 1. + # Step 1: p0 True while p1, p2 still False. Advance to index 1. preds[0].set([True]) _advance_step(env) sm.step(env, step_index=env.episode_length_buf) @@ -187,7 +190,7 @@ def _test_state_machine_advances_sequentially(simulation_app) -> bool: events = sm.get_events()[0] assert len(events) == 1 and events[0]["predicate_index"] == 0 - # Step 2: p0 reverts (irrelevant; latched forward), p1 fires. + # Step 2: p0 reverts False, p1 True. preds[0].set([False]) preds[1].set([True]) _advance_step(env) @@ -195,14 +198,14 @@ def _test_state_machine_advances_sequentially(simulation_app) -> bool: events = sm.get_events()[0] assert len(events) == 2 and events[-1]["predicate_index"] == 1 - # Step 3: p2 fires; subtask complete. + # Step 3: p2 True, subtask complete. preds[2].set([True]) _advance_step(env) sm.step(env, step_index=env.episode_length_buf) state = sm.get_state()[0]["fine_grained_subtasks"]["lift"] assert state["is_complete"] assert state["completed_groups"] == 1 - assert abs(state["score"] - 1.0) < 1e-6 + assert abs(state["score"] - 1.0) < SCORE_TOL events = sm.get_events()[0] assert len(events) == 3 and events[-1]["predicate_index"] == 2 except Exception as e: @@ -224,7 +227,7 @@ def _test_state_machine_ignores_out_of_order_success(simulation_app) -> bool: sm = FineGrainedStateMachine(fine_grained_subtasks=[fgs], num_envs=1, device="cpu") sm.reset([0]) - # p0 stays False; p1 and p2 fire — no progress should be made. + # p0 stays False and p1, p2 True. No progress should be made. preds[0].set([False]) preds[1].set([True]) preds[2].set([True]) @@ -236,8 +239,7 @@ def _test_state_machine_ignores_out_of_order_success(simulation_app) -> bool: assert state["score"] == 0.0 assert len(sm.get_events()[0]) == 0 - # Now p0 fires; chain catches up: 0, 1, 2 should all advance over subsequent - # steps (one per step, in order). + # Now p0 True, p1, p2 should advance over subsequent steps. preds[0].set([True]) for _ in range(3): _advance_step(env) @@ -253,7 +255,7 @@ def _test_state_machine_ignores_out_of_order_success(simulation_app) -> bool: def _test_state_machine_logical_any(simulation_app) -> bool: - """Two parallel groups with logical='any' complete as soon as either one finishes.""" + """Two parallel groups with logical=any complete as soon as either one finishes.""" from isaaclab_arena.tasks.fine_grained_state_machine import FineGrainedStateMachine from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask @@ -274,7 +276,7 @@ def _test_state_machine_logical_any(simulation_app) -> bool: sm.step(env, step_index=env.episode_length_buf) assert not sm.get_state()[0]["fine_grained_subtasks"]["either"]["is_complete"] - # Group "a" completes -> done. + # Group p_a completes -> done. p_a.set([True]) _advance_step(env) sm.step(env, step_index=env.episode_length_buf) @@ -287,7 +289,7 @@ def _test_state_machine_logical_any(simulation_app) -> bool: def _test_state_machine_logical_all(simulation_app) -> bool: - """logical='all' requires every group to complete its chain.""" + """Two groups with logical=all complete once all groups are complete.""" from isaaclab_arena.tasks.fine_grained_state_machine import FineGrainedStateMachine from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask @@ -303,13 +305,13 @@ def _test_state_machine_logical_all(simulation_app) -> bool: sm = FineGrainedStateMachine(fine_grained_subtasks=[fgs], num_envs=1, device="cpu") sm.reset([0]) - # Only "a" completes -> still not done. + # Only p_a completes -> still not done. p_a.set([True]) _advance_step(env) sm.step(env, step_index=env.episode_length_buf) assert not sm.get_state()[0]["fine_grained_subtasks"]["both"]["is_complete"] - # "b" also completes -> done. + # p_b also completes -> done. p_b.set([True]) _advance_step(env) sm.step(env, step_index=env.episode_length_buf) @@ -322,7 +324,7 @@ def _test_state_machine_logical_all(simulation_app) -> bool: def _test_state_machine_logical_choose(simulation_app) -> bool: - """logical='choose' with K=2 requires any two of three groups to complete.""" + """Three groups with logical=choose and K=2 complete once any two groups are complete.""" from isaaclab_arena.tasks.fine_grained_state_machine import FineGrainedStateMachine from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask @@ -340,13 +342,13 @@ def _test_state_machine_logical_choose(simulation_app) -> bool: sm = FineGrainedStateMachine(fine_grained_subtasks=[fgs], num_envs=1, device="cpu") sm.reset([0]) - # Only one group complete -> not done. + # Only p_a group complete -> not done. p_a.set([True]) _advance_step(env) sm.step(env, step_index=env.episode_length_buf) assert not sm.get_state()[0]["fine_grained_subtasks"]["any_two"]["is_complete"] - # Two complete -> done. + # p_b also complete -> done. p_b.set([True]) _advance_step(env) sm.step(env, step_index=env.episode_length_buf) @@ -370,7 +372,7 @@ def _test_state_machine_reset_clears_state(simulation_app) -> bool: sm = FineGrainedStateMachine(fine_grained_subtasks=[fgs], num_envs=2, device="cpu") sm.reset([0, 1]) - # Drive env 0 to fully complete, env 1 halfway. + # Set env 0 to fully complete. preds[0].set([True, True]) preds[1].set([True, False]) _advance_step(env) @@ -399,10 +401,8 @@ def _test_state_machine_reset_clears_state(simulation_app) -> bool: return True -def _test_gating_active_when_parent_subtask_idx_matches(simulation_app) -> bool: - """A recipe with ``parent_subtask_idx=N`` advances normally when the env's - ``_current_subtask_idx`` matches N. This mirrors the sequential-composite - case where subtask N is the currently-active parent subtask.""" +def _test_gating_advance_when_parent_subtask_idx_matches(simulation_app) -> bool: + """A FineGrainedSubtask with parent_subtask_idx=N advances when the env's _current_subtask_idx=N.""" from isaaclab_arena.tasks.fine_grained_state_machine import FineGrainedStateMachine from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask @@ -428,22 +428,20 @@ def _test_gating_active_when_parent_subtask_idx_matches(simulation_app) -> bool: def _test_gating_blocked_when_parent_subtask_idx_mismatches(simulation_app) -> bool: - """A recipe with ``parent_subtask_idx=N`` is frozen when the env's - ``_current_subtask_idx`` differs from N — even if the underlying predicate - is True. Once the parent advances to N, the recipe starts advancing too.""" + """A FineGrainedSubtask with parent_subtask_idx=N doesn't advance when the env's _current_subtask_idx!=N.""" from isaaclab_arena.tasks.fine_grained_state_machine import FineGrainedStateMachine from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask try: env = _MockEnv(num_envs=1) - env._current_subtask_idx = [0] # parent on subtask 0, recipe targets 1 + env._current_subtask_idx = [0] pred = _MockPredicate(num_envs=1, name="p") fgs = FineGrainedSubtask(name="t", predicate_groups=pred, parent_subtask_idx=1) sm = FineGrainedStateMachine(fine_grained_subtasks=[fgs], num_envs=1, device="cpu") sm.reset([0]) - # Predicate True, but the parent isn't at this recipe's index yet. + # Predicate True, but the parent isn't at this FGS's index yet. pred.set([True]) _advance_step(env) sm.step(env, step_index=env.episode_length_buf) @@ -451,7 +449,7 @@ def _test_gating_blocked_when_parent_subtask_idx_mismatches(simulation_app) -> b assert sm.get_state()[0]["fine_grained_subtasks"]["t"]["score"] == 0.0 assert len(sm.get_events()[0]) == 0 - # Parent advances to this recipe's index; the recipe catches up. + # Parent advances to this FGS's index, state machine advances. env._current_subtask_idx = [1] _advance_step(env) sm.step(env, step_index=env.episode_length_buf) @@ -464,27 +462,37 @@ def _test_gating_blocked_when_parent_subtask_idx_mismatches(simulation_app) -> b return True -def _test_gating_noop_when_env_has_no_current_subtask_idx(simulation_app) -> bool: - """For unordered composite tasks, ``env._current_subtask_idx`` is absent. - Recipes carry ``parent_subtask_idx`` (for namespacing) but gating is a no-op - and all recipes advance whenever their predicates fire.""" +def _test_gating_sequential_task_end_to_end(simulation_app) -> bool: + """Two FGSs with different parent subtask indices. The parent's + _current_subtask_idx advances over time. Each FGS only progresses + during its active window.""" from isaaclab_arena.tasks.fine_grained_state_machine import FineGrainedStateMachine from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask try: env = _MockEnv(num_envs=1) - # No _current_subtask_idx attribute on env -> unordered composite path. + env._current_subtask_idx = [0] - pred = _MockPredicate(num_envs=1, name="p") - fgs = FineGrainedSubtask(name="t", predicate_groups=pred, parent_subtask_idx=1) - sm = FineGrainedStateMachine(fine_grained_subtasks=[fgs], num_envs=1, device="cpu") + pred_a = _MockPredicate(num_envs=1, name="a") + pred_b = _MockPredicate(num_envs=1, name="b") + fgs_a = FineGrainedSubtask(name="a", predicate_groups=pred_a, parent_subtask_idx=0) + fgs_b = FineGrainedSubtask(name="b", predicate_groups=pred_b, parent_subtask_idx=1) + sm = FineGrainedStateMachine(fine_grained_subtasks=[fgs_a, fgs_b], num_envs=1, device="cpu") sm.reset([0]) - pred.set([True]) + # Both predicates True, but only pred_a is active. + pred_a.set([True]) + pred_b.set([True]) _advance_step(env) sm.step(env, step_index=env.episode_length_buf) - assert sm.get_state()[0]["fine_grained_subtasks"]["t"]["is_complete"] - assert len(sm.get_events()[0]) == 1 + assert sm.get_state()[0]["fine_grained_subtasks"]["a"]["is_complete"] + assert not sm.get_state()[0]["fine_grained_subtasks"]["b"]["is_complete"] + + # Advances to subtask 1 so pred_b is now active. + env._current_subtask_idx = [1] + _advance_step(env) + sm.step(env, step_index=env.episode_length_buf) + assert sm.get_state()[0]["fine_grained_subtasks"]["b"]["is_complete"] except Exception as e: print(f"Error: {e}") traceback.print_exc() @@ -492,37 +500,24 @@ def _test_gating_noop_when_env_has_no_current_subtask_idx(simulation_app) -> boo return True -def _test_sequential_gating_end_to_end(simulation_app) -> bool: - """Two recipes targeting different parent subtask indices. The parent's - ``_current_subtask_idx`` advances over time. Each recipe only progresses - during its parent's active window.""" +def _test_gating_noop_when_env_has_no_current_subtask_idx(simulation_app) -> bool: + """For unordered composite tasks gating is a no-op and all FGSs advance whenever their predicates are True.""" from isaaclab_arena.tasks.fine_grained_state_machine import FineGrainedStateMachine from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask try: env = _MockEnv(num_envs=1) - env._current_subtask_idx = [0] - pred_a = _MockPredicate(num_envs=1, name="a") - pred_b = _MockPredicate(num_envs=1, name="b") - fgs_a = FineGrainedSubtask(name="a", predicate_groups=pred_a, parent_subtask_idx=0) - fgs_b = FineGrainedSubtask(name="b", predicate_groups=pred_b, parent_subtask_idx=1) - sm = FineGrainedStateMachine(fine_grained_subtasks=[fgs_a, fgs_b], num_envs=1, device="cpu") + pred = _MockPredicate(num_envs=1, name="p") + fgs = FineGrainedSubtask(name="t", predicate_groups=pred, parent_subtask_idx=1) + sm = FineGrainedStateMachine(fine_grained_subtasks=[fgs], num_envs=1, device="cpu") sm.reset([0]) - # Both predicates True, but only "a" is active (parent on subtask 0). - pred_a.set([True]) - pred_b.set([True]) - _advance_step(env) - sm.step(env, step_index=env.episode_length_buf) - assert sm.get_state()[0]["fine_grained_subtasks"]["a"]["is_complete"] - assert not sm.get_state()[0]["fine_grained_subtasks"]["b"]["is_complete"] - - # Parent advances to subtask 1; "b" is now active. - env._current_subtask_idx = [1] + pred.set([True]) _advance_step(env) sm.step(env, step_index=env.episode_length_buf) - assert sm.get_state()[0]["fine_grained_subtasks"]["b"]["is_complete"] + assert sm.get_state()[0]["fine_grained_subtasks"]["t"]["is_complete"] + assert len(sm.get_events()[0]) == 1 except Exception as e: print(f"Error: {e}") traceback.print_exc() @@ -545,7 +540,7 @@ def _test_step_func_publishes_to_extras_and_returns_no_termination(simulation_ap fine_grained_subtask_reset_func(env, env_ids=[0, 1], fine_grained_subtasks=[fgs]) - # Step with predicate False: state machine ticks but no transitions. + # Step with predicate=False, state machine ticks but no transitions. result = fine_grained_subtask_step_func(env, fine_grained_subtasks=[fgs]) assert result.tolist() == [False, False] assert "fine_grained_subtask" in env.extras @@ -553,7 +548,7 @@ def _test_step_func_publishes_to_extras_and_returns_no_termination(simulation_ap assert env.extras["fine_grained_subtask"]["events"] == [[], []] assert not env.extras["fine_grained_subtask"]["states"][0]["fine_grained_subtasks"]["t"]["is_complete"] - # Step with env 0 predicate True: env 0 completes, env 1 does not. + # Step with env 0 predicate True, env 0 completes, env 1 does not. pred.set([True, False]) _advance_step(env) result = fine_grained_subtask_step_func(env, fine_grained_subtasks=[fgs]) @@ -565,9 +560,7 @@ def _test_step_func_publishes_to_extras_and_returns_no_termination(simulation_ap assert len(events[0]) == 1 assert len(events[1]) == 0 - # Reset env 0; its state should clear, env 1 untouched. We also flip pred - # to False to verify the post-reset step starts from the chain head rather - # than auto-re-completing. + # Reset env 0, env 1 untouched. pred.set([False, False]) fine_grained_subtask_reset_func(env, env_ids=[0], fine_grained_subtasks=[fgs]) result = fine_grained_subtask_step_func(env, fine_grained_subtasks=[fgs]) @@ -582,14 +575,9 @@ def _test_step_func_publishes_to_extras_and_returns_no_termination(simulation_ap def _test_task_base_fine_grained_subtask_hooks(simulation_app) -> bool: - """TaskBase's fine-grained-subtask hooks: default is empty/None; overriding + """Test TaskBase's fine-grained-subtask hooks. Default is empty/None. Overriding ``get_fine_grained_subtasks`` causes the events/termination helpers to return real cfgs that the env builder picks up automatically. - - Importing ``task_base`` is non-trivial in the test sandbox (it transitively - pulls in the asset-library network registration). Both default and opt-in - behavior are exercised inside a single test so the module import only - happens once. """ from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask from isaaclab_arena.tasks.task_base import TaskBase @@ -627,7 +615,6 @@ def get_fine_grained_subtasks(self): assert opt_in.get_fine_grained_subtask_events_cfg() is not None assert opt_in.get_fine_grained_subtask_termination_cfg() is not None - # ---- CompositeTaskBase: concatenate child recipes with namespace + parent index ---- from isaaclab_arena.tasks.composite_task_base import CompositeTaskBase class _ChildA(_Base): @@ -646,7 +633,6 @@ def get_fine_grained_subtasks(self): assert recipes[1].name == "subtask_1/close" assert recipes[1].parent_subtask_idx == 1 - # get_own_fine_grained_subtasks adds composite-level recipes (no gating). class _CompositeWithOwn(CompositeTaskBase): def get_own_fine_grained_subtasks(self): return [FineGrainedSubtask(name="both_done", predicate_groups=_MockPredicate(1, name="own"))] @@ -655,7 +641,7 @@ def get_own_fine_grained_subtasks(self): recipes2 = composite2.get_fine_grained_subtasks() assert len(recipes2) == 3 assert recipes2[2].name == "both_done" - assert recipes2[2].parent_subtask_idx is None # composite-level: never gated + assert recipes2[2].parent_subtask_idx is None except Exception as e: print(f"Error: {e}") traceback.print_exc() @@ -663,27 +649,24 @@ def get_own_fine_grained_subtasks(self): return True -# Pytest entry points ----------------------------------------------------------- - - -def test_format_single_callable(): - assert run_simulation_app_function(_test_format_single_callable, headless=HEADLESS) +def test_predicate_groups_single_callable(): + assert run_simulation_app_function(_test_predicate_groups_single_callable, headless=HEADLESS) -def test_format_list_of_callables(): - assert run_simulation_app_function(_test_format_list_of_callables, headless=HEADLESS) +def test_predicate_groups_list_of_callables(): + assert run_simulation_app_function(_test_predicate_groups_list_of_callables, headless=HEADLESS) -def test_format_weighted_tuples(): - assert run_simulation_app_function(_test_format_weighted_tuples, headless=HEADLESS) +def test_predicate_groups_weighted_tuples(): + assert run_simulation_app_function(_test_predicate_groups_weighted_tuples, headless=HEADLESS) -def test_format_dict_groups(): - assert run_simulation_app_function(_test_format_dict_groups, headless=HEADLESS) +def test_predicate_groups_dict_groups(): + assert run_simulation_app_function(_test_predicate_groups_dict_groups, headless=HEADLESS) -def test_format_rejects_invalid_inputs(): - assert run_simulation_app_function(_test_format_rejects_invalid_inputs, headless=HEADLESS) +def test_predicate_groups_rejects_invalid_inputs(): + assert run_simulation_app_function(_test_predicate_groups_rejects_invalid_inputs, headless=HEADLESS) def test_state_machine_advances_sequentially(): @@ -710,8 +693,8 @@ def test_state_machine_reset_clears_state(): assert run_simulation_app_function(_test_state_machine_reset_clears_state, headless=HEADLESS) -def test_gating_active_when_parent_subtask_idx_matches(): - assert run_simulation_app_function(_test_gating_active_when_parent_subtask_idx_matches, headless=HEADLESS) +def test_gating_advance_when_parent_subtask_idx_matches(): + assert run_simulation_app_function(_test_gating_advance_when_parent_subtask_idx_matches, headless=HEADLESS) def test_gating_blocked_when_parent_subtask_idx_mismatches(): @@ -722,8 +705,8 @@ def test_gating_noop_when_env_has_no_current_subtask_idx(): assert run_simulation_app_function(_test_gating_noop_when_env_has_no_current_subtask_idx, headless=HEADLESS) -def test_sequential_gating_end_to_end(): - assert run_simulation_app_function(_test_sequential_gating_end_to_end, headless=HEADLESS) +def test_gating_sequential_task_end_to_end(): + assert run_simulation_app_function(_test_gating_sequential_task_end_to_end, headless=HEADLESS) def test_step_func_publishes_to_extras_and_returns_no_termination(): @@ -737,20 +720,20 @@ def test_task_base_fine_grained_subtask_hooks(): if __name__ == "__main__": - test_format_single_callable() - test_format_list_of_callables() - test_format_weighted_tuples() - test_format_dict_groups() - test_format_rejects_invalid_inputs() + test_predicate_groups_single_callable() + test_predicate_groups_list_of_callables() + test_predicate_groups_weighted_tuples() + test_predicate_groups_dict_groups() + test_predicate_groups_rejects_invalid_inputs() test_state_machine_advances_sequentially() test_state_machine_ignores_out_of_order_success() test_state_machine_logical_any() test_state_machine_logical_all() test_state_machine_logical_choose() test_state_machine_reset_clears_state() - test_gating_active_when_parent_subtask_idx_matches() + test_gating_advance_when_parent_subtask_idx_matches() test_gating_blocked_when_parent_subtask_idx_mismatches() test_gating_noop_when_env_has_no_current_subtask_idx() - test_sequential_gating_end_to_end() + test_gating_sequential_task_end_to_end() test_step_func_publishes_to_extras_and_returns_no_termination() test_task_base_fine_grained_subtask_hooks() From 6b734ff8ae80145deedd5057d027abc6ce6423ca Mon Sep 17 00:00:00 2001 From: Peter Du Date: Wed, 3 Jun 2026 10:46:01 -0700 Subject: [PATCH 03/10] update task base and composite task base with fgs --- isaaclab_arena/tasks/composite_task_base.py | 21 +++++---------------- isaaclab_arena/tasks/task_base.py | 7 ------- 2 files changed, 5 insertions(+), 23 deletions(-) diff --git a/isaaclab_arena/tasks/composite_task_base.py b/isaaclab_arena/tasks/composite_task_base.py index df6dec41e..790b923c6 100644 --- a/isaaclab_arena/tasks/composite_task_base.py +++ b/isaaclab_arena/tasks/composite_task_base.py @@ -363,27 +363,16 @@ def get_metrics(self) -> list[MetricBase]: return subtask_metrics def get_own_fine_grained_subtasks(self) -> list[FineGrainedSubtask]: - """Composite-level fine-grained recipes (e.g. cross-child conditions). + """Composite-level FineGrainedSubtasks. - These are added on top of whatever recipes the children declare. They - carry ``parent_subtask_idx = None`` and therefore have no activation - gating — they remain active for the full episode. + These are added on top of whatever FGSs the child subtasks declare and are not gated. """ return [] def get_fine_grained_subtasks(self) -> list[FineGrainedSubtask]: - """Concatenate children's fine-grained recipes with namespace prefixes. - - Each child's recipe gets a new name (``f"subtask_{i}/{original_name}"``) - and a ``parent_subtask_idx = i`` tag. For ``SequentialTaskBase`` the - state machine reads ``parent_subtask_idx`` to gate advancement against - ``env._current_subtask_idx`` so a child's recipes only advance while - that child is the active parent subtask. For unordered - ``CompositeTaskBase`` (no ``_current_subtask_idx``) the gate is a no-op - and every recipe is always active. - - Composite-level recipes from ``get_own_fine_grained_subtasks`` are - appended afterward and are not gated. + """Concatenate child subtasks's FineGrainedSubtasks with namespace prefixes. + + Each child's FGS gets a new name (subtask_{i}/{original_name}) and a parent_subtask_idx = i tag. """ out: list[FineGrainedSubtask] = [] for i, child in enumerate(self.subtasks): diff --git a/isaaclab_arena/tasks/task_base.py b/isaaclab_arena/tasks/task_base.py index 3e4d5ef97..c0884830b 100644 --- a/isaaclab_arena/tasks/task_base.py +++ b/isaaclab_arena/tasks/task_base.py @@ -69,13 +69,6 @@ def get_task_description(self) -> str | None: return self.task_description def get_fine_grained_subtasks(self) -> list[FineGrainedSubtask]: - """Return per-step fine-grained subtasks the state machine should track. Default: none. - - Override in subclasses to opt in to fine-grained subtask tracking. When - the returned list is non-empty, the env builder will automatically - register a reset event and a per-step termination term that tick the - state machine and publish ``env.extras["fine_grained_subtask"]``. - """ return [] def get_fine_grained_subtask_events_cfg(self) -> Any: From 231db88148a1f471c09fd980ab5b3b8ab7dd6dbc Mon Sep 17 00:00:00 2001 From: Peter Du Date: Wed, 3 Jun 2026 10:46:58 -0700 Subject: [PATCH 04/10] update get_fgs in composite tasks --- isaaclab_arena/tasks/composite_task_base.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/isaaclab_arena/tasks/composite_task_base.py b/isaaclab_arena/tasks/composite_task_base.py index 790b923c6..0323194f5 100644 --- a/isaaclab_arena/tasks/composite_task_base.py +++ b/isaaclab_arena/tasks/composite_task_base.py @@ -374,18 +374,18 @@ def get_fine_grained_subtasks(self) -> list[FineGrainedSubtask]: Each child's FGS gets a new name (subtask_{i}/{original_name}) and a parent_subtask_idx = i tag. """ - out: list[FineGrainedSubtask] = [] + fgs_list: list[FineGrainedSubtask] = [] for i, child in enumerate(self.subtasks): for fgs in child.get_fine_grained_subtasks(): - out.append( + fgs_list.append( dataclasses.replace( fgs, name=f"subtask_{i}/{fgs.name}", parent_subtask_idx=i, ) ) - out.extend(self.get_own_fine_grained_subtasks()) - return out + fgs_list.extend(self.get_own_fine_grained_subtasks()) + return fgs_list def _validate_consistent_mimic_eef_names(self, arm_mode: ArmMode) -> set[str]: "Check that all subtasks have the same Mimic eef_names." From 4c9aa2b89029f611eeb47c6ddb7d58d65ac3e81f Mon Sep 17 00:00:00 2001 From: Peter Du Date: Wed, 3 Jun 2026 10:52:47 -0700 Subject: [PATCH 05/10] update state machine naming --- .../tasks/fine_grained_state_machine.py | 10 ++--- .../tests/test_fine_grained_subtask.py | 40 +++++++++---------- 2 files changed, 25 insertions(+), 25 deletions(-) diff --git a/isaaclab_arena/tasks/fine_grained_state_machine.py b/isaaclab_arena/tasks/fine_grained_state_machine.py index b99f5948b..a650c34db 100644 --- a/isaaclab_arena/tasks/fine_grained_state_machine.py +++ b/isaaclab_arena/tasks/fine_grained_state_machine.py @@ -178,7 +178,7 @@ def overall_score_per_env(self) -> torch.Tensor: return stacked.mean(dim=1) -class FineGrainedStateMachine: +class FineGrainedSubtaskTrackingStateMachine: """State machine that manages runners for all FineGrainedSubtasks. Attributes: @@ -265,12 +265,12 @@ def get_events(self) -> list[list[dict]]: return [list(e) for e in self._events] -def _ensure_state_machine(env, fine_grained_subtasks: list[FineGrainedSubtask]) -> FineGrainedStateMachine: - """Return the env's FineGrainedStateMachine, lazily creating and caching it on first call.""" +def _ensure_state_machine(env, fine_grained_subtasks: list[FineGrainedSubtask]) -> FineGrainedSubtaskTrackingStateMachine: + """Return the env's FineGrainedSubtaskTrackingStateMachine, lazily creating and caching it on first call.""" - sm: FineGrainedStateMachine | None = getattr(env, _STATE_MACHINE_ATTR, None) + sm: FineGrainedSubtaskTrackingStateMachine | None = getattr(env, _STATE_MACHINE_ATTR, None) if sm is None: - sm = FineGrainedStateMachine( + sm = FineGrainedSubtaskTrackingStateMachine( fine_grained_subtasks=fine_grained_subtasks, num_envs=env.num_envs, device=env.device ) setattr(env, _STATE_MACHINE_ATTR, sm) diff --git a/isaaclab_arena/tests/test_fine_grained_subtask.py b/isaaclab_arena/tests/test_fine_grained_subtask.py index 0450e11a5..5fe977cf4 100644 --- a/isaaclab_arena/tests/test_fine_grained_subtask.py +++ b/isaaclab_arena/tests/test_fine_grained_subtask.py @@ -170,14 +170,14 @@ def _test_predicate_groups_rejects_invalid_inputs(simulation_app) -> bool: def _test_state_machine_advances_sequentially(simulation_app) -> bool: """A single FineGrainedSubtask with a 3 predicate chain advances one step per satisfied predicate.""" - from isaaclab_arena.tasks.fine_grained_state_machine import FineGrainedStateMachine + from isaaclab_arena.tasks.fine_grained_state_machine import FineGrainedSubtaskTrackingStateMachine from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask try: env = _MockEnv(num_envs=1) preds = [_MockPredicate(num_envs=1, name=f"p{i}") for i in range(3)] fgs = FineGrainedSubtask(name="lift", predicate_groups=preds) - sm = FineGrainedStateMachine(fine_grained_subtasks=[fgs], num_envs=1, device="cpu") + sm = FineGrainedSubtaskTrackingStateMachine(fine_grained_subtasks=[fgs], num_envs=1, device="cpu") sm.reset([0]) # Step 1: p0 True while p1, p2 still False. Advance to index 1. @@ -217,14 +217,14 @@ def _test_state_machine_advances_sequentially(simulation_app) -> bool: def _test_state_machine_ignores_out_of_order_success(simulation_app) -> bool: """If a later predicate fires first, it's ignored until preceding ones have advanced.""" - from isaaclab_arena.tasks.fine_grained_state_machine import FineGrainedStateMachine + from isaaclab_arena.tasks.fine_grained_state_machine import FineGrainedSubtaskTrackingStateMachine from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask try: env = _MockEnv(num_envs=1) preds = [_MockPredicate(num_envs=1, name=f"p{i}") for i in range(3)] fgs = FineGrainedSubtask(name="lift", predicate_groups=preds) - sm = FineGrainedStateMachine(fine_grained_subtasks=[fgs], num_envs=1, device="cpu") + sm = FineGrainedSubtaskTrackingStateMachine(fine_grained_subtasks=[fgs], num_envs=1, device="cpu") sm.reset([0]) # p0 stays False and p1, p2 True. No progress should be made. @@ -256,7 +256,7 @@ def _test_state_machine_ignores_out_of_order_success(simulation_app) -> bool: def _test_state_machine_logical_any(simulation_app) -> bool: """Two parallel groups with logical=any complete as soon as either one finishes.""" - from isaaclab_arena.tasks.fine_grained_state_machine import FineGrainedStateMachine + from isaaclab_arena.tasks.fine_grained_state_machine import FineGrainedSubtaskTrackingStateMachine from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask try: @@ -268,7 +268,7 @@ def _test_state_machine_logical_any(simulation_app) -> bool: predicate_groups={"a": p_a, "b": p_b}, logical="any", ) - sm = FineGrainedStateMachine(fine_grained_subtasks=[fgs], num_envs=1, device="cpu") + sm = FineGrainedSubtaskTrackingStateMachine(fine_grained_subtasks=[fgs], num_envs=1, device="cpu") sm.reset([0]) # Neither group complete -> not done. @@ -290,7 +290,7 @@ def _test_state_machine_logical_any(simulation_app) -> bool: def _test_state_machine_logical_all(simulation_app) -> bool: """Two groups with logical=all complete once all groups are complete.""" - from isaaclab_arena.tasks.fine_grained_state_machine import FineGrainedStateMachine + from isaaclab_arena.tasks.fine_grained_state_machine import FineGrainedSubtaskTrackingStateMachine from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask try: @@ -302,7 +302,7 @@ def _test_state_machine_logical_all(simulation_app) -> bool: predicate_groups={"a": p_a, "b": p_b}, logical="all", ) - sm = FineGrainedStateMachine(fine_grained_subtasks=[fgs], num_envs=1, device="cpu") + sm = FineGrainedSubtaskTrackingStateMachine(fine_grained_subtasks=[fgs], num_envs=1, device="cpu") sm.reset([0]) # Only p_a completes -> still not done. @@ -325,7 +325,7 @@ def _test_state_machine_logical_all(simulation_app) -> bool: def _test_state_machine_logical_choose(simulation_app) -> bool: """Three groups with logical=choose and K=2 complete once any two groups are complete.""" - from isaaclab_arena.tasks.fine_grained_state_machine import FineGrainedStateMachine + from isaaclab_arena.tasks.fine_grained_state_machine import FineGrainedSubtaskTrackingStateMachine from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask try: @@ -339,7 +339,7 @@ def _test_state_machine_logical_choose(simulation_app) -> bool: logical="choose", K=2, ) - sm = FineGrainedStateMachine(fine_grained_subtasks=[fgs], num_envs=1, device="cpu") + sm = FineGrainedSubtaskTrackingStateMachine(fine_grained_subtasks=[fgs], num_envs=1, device="cpu") sm.reset([0]) # Only p_a group complete -> not done. @@ -362,14 +362,14 @@ def _test_state_machine_logical_choose(simulation_app) -> bool: def _test_state_machine_reset_clears_state(simulation_app) -> bool: """Resetting an env_id zeroes its progress and event log, but leaves other envs alone.""" - from isaaclab_arena.tasks.fine_grained_state_machine import FineGrainedStateMachine + from isaaclab_arena.tasks.fine_grained_state_machine import FineGrainedSubtaskTrackingStateMachine from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask try: env = _MockEnv(num_envs=2) preds = [_MockPredicate(num_envs=2, name=f"p{i}") for i in range(2)] fgs = FineGrainedSubtask(name="t", predicate_groups=preds) - sm = FineGrainedStateMachine(fine_grained_subtasks=[fgs], num_envs=2, device="cpu") + sm = FineGrainedSubtaskTrackingStateMachine(fine_grained_subtasks=[fgs], num_envs=2, device="cpu") sm.reset([0, 1]) # Set env 0 to fully complete. @@ -403,7 +403,7 @@ def _test_state_machine_reset_clears_state(simulation_app) -> bool: def _test_gating_advance_when_parent_subtask_idx_matches(simulation_app) -> bool: """A FineGrainedSubtask with parent_subtask_idx=N advances when the env's _current_subtask_idx=N.""" - from isaaclab_arena.tasks.fine_grained_state_machine import FineGrainedStateMachine + from isaaclab_arena.tasks.fine_grained_state_machine import FineGrainedSubtaskTrackingStateMachine from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask try: @@ -412,7 +412,7 @@ def _test_gating_advance_when_parent_subtask_idx_matches(simulation_app) -> bool pred = _MockPredicate(num_envs=1, name="p") fgs = FineGrainedSubtask(name="t", predicate_groups=pred, parent_subtask_idx=1) - sm = FineGrainedStateMachine(fine_grained_subtasks=[fgs], num_envs=1, device="cpu") + sm = FineGrainedSubtaskTrackingStateMachine(fine_grained_subtasks=[fgs], num_envs=1, device="cpu") sm.reset([0]) pred.set([True]) @@ -429,7 +429,7 @@ def _test_gating_advance_when_parent_subtask_idx_matches(simulation_app) -> bool def _test_gating_blocked_when_parent_subtask_idx_mismatches(simulation_app) -> bool: """A FineGrainedSubtask with parent_subtask_idx=N doesn't advance when the env's _current_subtask_idx!=N.""" - from isaaclab_arena.tasks.fine_grained_state_machine import FineGrainedStateMachine + from isaaclab_arena.tasks.fine_grained_state_machine import FineGrainedSubtaskTrackingStateMachine from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask try: @@ -438,7 +438,7 @@ def _test_gating_blocked_when_parent_subtask_idx_mismatches(simulation_app) -> b pred = _MockPredicate(num_envs=1, name="p") fgs = FineGrainedSubtask(name="t", predicate_groups=pred, parent_subtask_idx=1) - sm = FineGrainedStateMachine(fine_grained_subtasks=[fgs], num_envs=1, device="cpu") + sm = FineGrainedSubtaskTrackingStateMachine(fine_grained_subtasks=[fgs], num_envs=1, device="cpu") sm.reset([0]) # Predicate True, but the parent isn't at this FGS's index yet. @@ -466,7 +466,7 @@ def _test_gating_sequential_task_end_to_end(simulation_app) -> bool: """Two FGSs with different parent subtask indices. The parent's _current_subtask_idx advances over time. Each FGS only progresses during its active window.""" - from isaaclab_arena.tasks.fine_grained_state_machine import FineGrainedStateMachine + from isaaclab_arena.tasks.fine_grained_state_machine import FineGrainedSubtaskTrackingStateMachine from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask try: @@ -477,7 +477,7 @@ def _test_gating_sequential_task_end_to_end(simulation_app) -> bool: pred_b = _MockPredicate(num_envs=1, name="b") fgs_a = FineGrainedSubtask(name="a", predicate_groups=pred_a, parent_subtask_idx=0) fgs_b = FineGrainedSubtask(name="b", predicate_groups=pred_b, parent_subtask_idx=1) - sm = FineGrainedStateMachine(fine_grained_subtasks=[fgs_a, fgs_b], num_envs=1, device="cpu") + sm = FineGrainedSubtaskTrackingStateMachine(fine_grained_subtasks=[fgs_a, fgs_b], num_envs=1, device="cpu") sm.reset([0]) # Both predicates True, but only pred_a is active. @@ -502,7 +502,7 @@ def _test_gating_sequential_task_end_to_end(simulation_app) -> bool: def _test_gating_noop_when_env_has_no_current_subtask_idx(simulation_app) -> bool: """For unordered composite tasks gating is a no-op and all FGSs advance whenever their predicates are True.""" - from isaaclab_arena.tasks.fine_grained_state_machine import FineGrainedStateMachine + from isaaclab_arena.tasks.fine_grained_state_machine import FineGrainedSubtaskTrackingStateMachine from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask try: @@ -510,7 +510,7 @@ def _test_gating_noop_when_env_has_no_current_subtask_idx(simulation_app) -> boo pred = _MockPredicate(num_envs=1, name="p") fgs = FineGrainedSubtask(name="t", predicate_groups=pred, parent_subtask_idx=1) - sm = FineGrainedStateMachine(fine_grained_subtasks=[fgs], num_envs=1, device="cpu") + sm = FineGrainedSubtaskTrackingStateMachine(fine_grained_subtasks=[fgs], num_envs=1, device="cpu") sm.reset([0]) pred.set([True]) From ddbb7304f5213b80a0d1888d93a9f07738ed035f Mon Sep 17 00:00:00 2001 From: Peter Du Date: Wed, 3 Jun 2026 11:02:38 -0700 Subject: [PATCH 06/10] update state machine file name --- isaaclab_arena/tasks/fine_grained_subtask.py | 1 - ...grained_subtask_tracking_state_machine.py} | 4 +++- isaaclab_arena/tasks/task_base.py | 4 ++-- .../tests/test_fine_grained_subtask.py | 24 +++++++++---------- 4 files changed, 17 insertions(+), 16 deletions(-) rename isaaclab_arena/tasks/{fine_grained_state_machine.py => fine_grained_subtask_tracking_state_machine.py} (99%) diff --git a/isaaclab_arena/tasks/fine_grained_subtask.py b/isaaclab_arena/tasks/fine_grained_subtask.py index eea31a80c..b65b28849 100644 --- a/isaaclab_arena/tasks/fine_grained_subtask.py +++ b/isaaclab_arena/tasks/fine_grained_subtask.py @@ -3,7 +3,6 @@ # # SPDX-License-Identifier: Apache-2.0 - from __future__ import annotations from collections.abc import Callable diff --git a/isaaclab_arena/tasks/fine_grained_state_machine.py b/isaaclab_arena/tasks/fine_grained_subtask_tracking_state_machine.py similarity index 99% rename from isaaclab_arena/tasks/fine_grained_state_machine.py rename to isaaclab_arena/tasks/fine_grained_subtask_tracking_state_machine.py index a650c34db..6a142ce3e 100644 --- a/isaaclab_arena/tasks/fine_grained_state_machine.py +++ b/isaaclab_arena/tasks/fine_grained_subtask_tracking_state_machine.py @@ -265,7 +265,9 @@ def get_events(self) -> list[list[dict]]: return [list(e) for e in self._events] -def _ensure_state_machine(env, fine_grained_subtasks: list[FineGrainedSubtask]) -> FineGrainedSubtaskTrackingStateMachine: +def _ensure_state_machine( + env, fine_grained_subtasks: list[FineGrainedSubtask] +) -> FineGrainedSubtaskTrackingStateMachine: """Return the env's FineGrainedSubtaskTrackingStateMachine, lazily creating and caching it on first call.""" sm: FineGrainedSubtaskTrackingStateMachine | None = getattr(env, _STATE_MACHINE_ATTR, None) diff --git a/isaaclab_arena/tasks/task_base.py b/isaaclab_arena/tasks/task_base.py index c0884830b..91b65b328 100644 --- a/isaaclab_arena/tasks/task_base.py +++ b/isaaclab_arena/tasks/task_base.py @@ -11,11 +11,11 @@ from isaaclab_arena.embodiments.common.arm_mode import ArmMode from isaaclab_arena.metrics.metric_base import MetricBase -from isaaclab_arena.tasks.fine_grained_state_machine import ( +from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask +from isaaclab_arena.tasks.fine_grained_subtask_tracking_state_machine import ( make_fine_grained_subtask_events_cfg, make_fine_grained_subtask_termination_cfg, ) -from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask class TaskBase(ABC): diff --git a/isaaclab_arena/tests/test_fine_grained_subtask.py b/isaaclab_arena/tests/test_fine_grained_subtask.py index 5fe977cf4..b976c0918 100644 --- a/isaaclab_arena/tests/test_fine_grained_subtask.py +++ b/isaaclab_arena/tests/test_fine_grained_subtask.py @@ -170,8 +170,8 @@ def _test_predicate_groups_rejects_invalid_inputs(simulation_app) -> bool: def _test_state_machine_advances_sequentially(simulation_app) -> bool: """A single FineGrainedSubtask with a 3 predicate chain advances one step per satisfied predicate.""" - from isaaclab_arena.tasks.fine_grained_state_machine import FineGrainedSubtaskTrackingStateMachine from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask + from isaaclab_arena.tasks.fine_grained_subtask_tracking_state_machine import FineGrainedSubtaskTrackingStateMachine try: env = _MockEnv(num_envs=1) @@ -217,8 +217,8 @@ def _test_state_machine_advances_sequentially(simulation_app) -> bool: def _test_state_machine_ignores_out_of_order_success(simulation_app) -> bool: """If a later predicate fires first, it's ignored until preceding ones have advanced.""" - from isaaclab_arena.tasks.fine_grained_state_machine import FineGrainedSubtaskTrackingStateMachine from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask + from isaaclab_arena.tasks.fine_grained_subtask_tracking_state_machine import FineGrainedSubtaskTrackingStateMachine try: env = _MockEnv(num_envs=1) @@ -256,8 +256,8 @@ def _test_state_machine_ignores_out_of_order_success(simulation_app) -> bool: def _test_state_machine_logical_any(simulation_app) -> bool: """Two parallel groups with logical=any complete as soon as either one finishes.""" - from isaaclab_arena.tasks.fine_grained_state_machine import FineGrainedSubtaskTrackingStateMachine from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask + from isaaclab_arena.tasks.fine_grained_subtask_tracking_state_machine import FineGrainedSubtaskTrackingStateMachine try: env = _MockEnv(num_envs=1) @@ -290,8 +290,8 @@ def _test_state_machine_logical_any(simulation_app) -> bool: def _test_state_machine_logical_all(simulation_app) -> bool: """Two groups with logical=all complete once all groups are complete.""" - from isaaclab_arena.tasks.fine_grained_state_machine import FineGrainedSubtaskTrackingStateMachine from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask + from isaaclab_arena.tasks.fine_grained_subtask_tracking_state_machine import FineGrainedSubtaskTrackingStateMachine try: env = _MockEnv(num_envs=1) @@ -325,8 +325,8 @@ def _test_state_machine_logical_all(simulation_app) -> bool: def _test_state_machine_logical_choose(simulation_app) -> bool: """Three groups with logical=choose and K=2 complete once any two groups are complete.""" - from isaaclab_arena.tasks.fine_grained_state_machine import FineGrainedSubtaskTrackingStateMachine from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask + from isaaclab_arena.tasks.fine_grained_subtask_tracking_state_machine import FineGrainedSubtaskTrackingStateMachine try: env = _MockEnv(num_envs=1) @@ -362,8 +362,8 @@ def _test_state_machine_logical_choose(simulation_app) -> bool: def _test_state_machine_reset_clears_state(simulation_app) -> bool: """Resetting an env_id zeroes its progress and event log, but leaves other envs alone.""" - from isaaclab_arena.tasks.fine_grained_state_machine import FineGrainedSubtaskTrackingStateMachine from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask + from isaaclab_arena.tasks.fine_grained_subtask_tracking_state_machine import FineGrainedSubtaskTrackingStateMachine try: env = _MockEnv(num_envs=2) @@ -403,8 +403,8 @@ def _test_state_machine_reset_clears_state(simulation_app) -> bool: def _test_gating_advance_when_parent_subtask_idx_matches(simulation_app) -> bool: """A FineGrainedSubtask with parent_subtask_idx=N advances when the env's _current_subtask_idx=N.""" - from isaaclab_arena.tasks.fine_grained_state_machine import FineGrainedSubtaskTrackingStateMachine from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask + from isaaclab_arena.tasks.fine_grained_subtask_tracking_state_machine import FineGrainedSubtaskTrackingStateMachine try: env = _MockEnv(num_envs=1) @@ -429,8 +429,8 @@ def _test_gating_advance_when_parent_subtask_idx_matches(simulation_app) -> bool def _test_gating_blocked_when_parent_subtask_idx_mismatches(simulation_app) -> bool: """A FineGrainedSubtask with parent_subtask_idx=N doesn't advance when the env's _current_subtask_idx!=N.""" - from isaaclab_arena.tasks.fine_grained_state_machine import FineGrainedSubtaskTrackingStateMachine from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask + from isaaclab_arena.tasks.fine_grained_subtask_tracking_state_machine import FineGrainedSubtaskTrackingStateMachine try: env = _MockEnv(num_envs=1) @@ -466,8 +466,8 @@ def _test_gating_sequential_task_end_to_end(simulation_app) -> bool: """Two FGSs with different parent subtask indices. The parent's _current_subtask_idx advances over time. Each FGS only progresses during its active window.""" - from isaaclab_arena.tasks.fine_grained_state_machine import FineGrainedSubtaskTrackingStateMachine from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask + from isaaclab_arena.tasks.fine_grained_subtask_tracking_state_machine import FineGrainedSubtaskTrackingStateMachine try: env = _MockEnv(num_envs=1) @@ -502,8 +502,8 @@ def _test_gating_sequential_task_end_to_end(simulation_app) -> bool: def _test_gating_noop_when_env_has_no_current_subtask_idx(simulation_app) -> bool: """For unordered composite tasks gating is a no-op and all FGSs advance whenever their predicates are True.""" - from isaaclab_arena.tasks.fine_grained_state_machine import FineGrainedSubtaskTrackingStateMachine from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask + from isaaclab_arena.tasks.fine_grained_subtask_tracking_state_machine import FineGrainedSubtaskTrackingStateMachine try: env = _MockEnv(num_envs=1) @@ -527,11 +527,11 @@ def _test_gating_noop_when_env_has_no_current_subtask_idx(simulation_app) -> boo def _test_step_func_publishes_to_extras_and_returns_no_termination(simulation_app) -> bool: """fine_grained_subtask_step_func writes env.extras and returns all-False.""" - from isaaclab_arena.tasks.fine_grained_state_machine import ( + from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask + from isaaclab_arena.tasks.fine_grained_subtask_tracking_state_machine import ( fine_grained_subtask_reset_func, fine_grained_subtask_step_func, ) - from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask try: env = _MockEnv(num_envs=2) From b394b0968f3c8701b099cd28bbe01c96cdb56dc0 Mon Sep 17 00:00:00 2001 From: Peter Du Date: Fri, 5 Jun 2026 09:56:16 -0700 Subject: [PATCH 07/10] update naming from subtasks to progress_tracking --- .../environments/arena_env_builder.py | 4 +- isaaclab_arena/tasks/composite_task_base.py | 28 +- ....py => fine_grained_progress_objective.py} | 36 +-- ...ne.py => fine_grained_progress_tracker.py} | 154 +++++----- isaaclab_arena/tasks/task_base.py | 26 +- ...ne_grained_progress_objective_tracking.py} | 270 +++++++++--------- 6 files changed, 259 insertions(+), 259 deletions(-) rename isaaclab_arena/tasks/{fine_grained_subtask.py => fine_grained_progress_objective.py} (78%) rename isaaclab_arena/tasks/{fine_grained_subtask_tracking_state_machine.py => fine_grained_progress_tracker.py} (65%) rename isaaclab_arena/tests/{test_fine_grained_subtask.py => test_fine_grained_progress_objective_tracking.py} (62%) diff --git a/isaaclab_arena/environments/arena_env_builder.py b/isaaclab_arena/environments/arena_env_builder.py index 5c237ac10..fda88dba0 100644 --- a/isaaclab_arena/environments/arena_env_builder.py +++ b/isaaclab_arena/environments/arena_env_builder.py @@ -125,14 +125,14 @@ def compose_manager_cfg(self) -> IsaacLabArenaManagerBasedRLEnvCfg: self.arena_env.scene.get_events_cfg(), task.get_events_cfg(), placement_event_cfg, - task.get_fine_grained_subtask_events_cfg(), + task.get_fine_grained_progress_objective_events_cfg(), ) termination_cfg = combine_configclass_instances( "TerminationCfg", task.get_termination_cfg(), self.arena_env.scene.get_termination_cfg(), embodiment.get_termination_cfg(), - task.get_fine_grained_subtask_termination_cfg(), + task.get_fine_grained_progress_objective_termination_cfg(), ) actions_cfg = embodiment.get_action_cfg() xr_cfg = embodiment.get_xr_cfg() diff --git a/isaaclab_arena/tasks/composite_task_base.py b/isaaclab_arena/tasks/composite_task_base.py index 0323194f5..1e2eaf732 100644 --- a/isaaclab_arena/tasks/composite_task_base.py +++ b/isaaclab_arena/tasks/composite_task_base.py @@ -21,7 +21,7 @@ from isaaclab_arena.metrics.metric_base import MetricBase from isaaclab_arena.metrics.metric_term_cfg import MetricTermCfg from isaaclab_arena.tasks.common.mimic_default_params import MIMIC_DATAGEN_CONFIG_DEFAULTS -from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask +from isaaclab_arena.tasks.fine_grained_progress_objective import FineGrainedProgressObjective from isaaclab_arena.tasks.task_base import TaskBase from isaaclab_arena.utils.configclass import ( check_configclass_field_duplicates, @@ -362,30 +362,30 @@ def get_metrics(self) -> list[MetricBase]: return subtask_metrics - def get_own_fine_grained_subtasks(self) -> list[FineGrainedSubtask]: - """Composite-level FineGrainedSubtasks. + def get_own_fine_grained_progress_objectives(self) -> list[FineGrainedProgressObjective]: + """Composite-level FineGrainedProgressObjectives. - These are added on top of whatever FGSs the child subtasks declare and are not gated. + These are added on top of whatever FGPOs the child subtasks declare and are not gated. """ return [] - def get_fine_grained_subtasks(self) -> list[FineGrainedSubtask]: - """Concatenate child subtasks's FineGrainedSubtasks with namespace prefixes. + def get_fine_grained_progress_objectives(self) -> list[FineGrainedProgressObjective]: + """Concatenate child subtasks's FineGrainedProgressObjectives with namespace prefixes. - Each child's FGS gets a new name (subtask_{i}/{original_name}) and a parent_subtask_idx = i tag. + Each child's FGPO gets a new name (subtask_{i}/{original_name}) and a parent_subtask_idx = i tag. """ - fgs_list: list[FineGrainedSubtask] = [] + fgpo_list: list[FineGrainedProgressObjective] = [] for i, child in enumerate(self.subtasks): - for fgs in child.get_fine_grained_subtasks(): - fgs_list.append( + for fgpo in child.get_fine_grained_progress_objectives(): + fgpo_list.append( dataclasses.replace( - fgs, - name=f"subtask_{i}/{fgs.name}", + fgpo, + name=f"subtask_{i}/{fgpo.name}", parent_subtask_idx=i, ) ) - fgs_list.extend(self.get_own_fine_grained_subtasks()) - return fgs_list + fgpo_list.extend(self.get_own_fine_grained_progress_objectives()) + return fgpo_list def _validate_consistent_mimic_eef_names(self, arm_mode: ArmMode) -> set[str]: "Check that all subtasks have the same Mimic eef_names." diff --git a/isaaclab_arena/tasks/fine_grained_subtask.py b/isaaclab_arena/tasks/fine_grained_progress_objective.py similarity index 78% rename from isaaclab_arena/tasks/fine_grained_subtask.py rename to isaaclab_arena/tasks/fine_grained_progress_objective.py index b65b28849..b0cca726b 100644 --- a/isaaclab_arena/tasks/fine_grained_subtask.py +++ b/isaaclab_arena/tasks/fine_grained_progress_objective.py @@ -41,19 +41,19 @@ def format_predicate_groups(predicate_groups: PredicateGroups) -> dict[str, list if isinstance(predicate_groups, list): if len(predicate_groups) == 0: - raise ValueError("FineGrainedSubtask.predicate_groups list cannot be empty") + raise ValueError("FineGrainedProgressObjective.predicate_groups list cannot be empty") return {DEFAULT_GROUP_NAME: _format_group_chain(predicate_groups, group_name=DEFAULT_GROUP_NAME)} if isinstance(predicate_groups, dict): if len(predicate_groups) == 0: - raise ValueError("FineGrainedSubtask.predicate_groups dict cannot be empty") + raise ValueError("FineGrainedProgressObjective.predicate_groups dict cannot be empty") return { group_name: _format_group_chain(value, group_name=group_name) for group_name, value in predicate_groups.items() } raise TypeError( - f"FineGrainedSubtask.predicate_groups must be a callable, list, or dict; got {type(predicate_groups).__name__}" + f"FineGrainedProgressObjective.predicate_groups must be a callable, list, or dict; got {type(predicate_groups).__name__}" ) @@ -111,22 +111,22 @@ def normalize_scores( @dataclass -class FineGrainedSubtask: +class FineGrainedProgressObjective: """Configuration object that defines a scored predicate sequence to track progress within a task. - A FineGrainedSubtask specifies what the predicate state machine should track. - Each FineGrainedSubtask holds one or more sequential predicate chains (groups). + A FineGrainedProgressObjective specifies what the predicate state machine should track. + Each FineGrainedProgressObjective holds one or more sequential predicate chains (groups). Within a group, predicates run in order. Across groups, predicates run in parallel. Args: - name: Identifies the FineGrainedSubtask within the TaskBase. - predicate_groups: The sequential predicate chains that define the FineGrainedSubtask. - score: Weight of the FineGrainedSubtask in the TaskBase-level overall_score. - logical: How completed groups combine to determine if the FineGrainedSubtask is complete. + name: Identifies the FineGrainedProgressObjective within the TaskBase. + predicate_groups: The sequential predicate chains that define the FineGrainedProgressObjective. + score: Weight of the FineGrainedProgressObjective in the TaskBase-level overall_score. + logical: How completed groups combine to determine if the FineGrainedProgressObjective is complete. Can be "all", "any", or "choose" K: Required when logical == "choose". Specifies the number of groups that must be completed - to consider the FineGrainedSubtask complete. - description: An optional description of the FineGrainedSubtask. + to consider the FineGrainedProgressObjective complete. + description: An optional description of the FineGrainedProgressObjective. """ name: str @@ -139,15 +139,15 @@ class FineGrainedSubtask: canonical_predicate_groups: dict[str, list[tuple[Callable, float]]] = field(init=False, repr=False) # Index of the parent TaskBase this recipe belongs to. Set automatically by - # CompositeTaskBase.get_fine_grained_subtasks() when used with composite tasks. + # CompositeTaskBase.get_fine_grained_progress_objectives() when used with composite tasks. parent_subtask_idx: int | None = None def __post_init__(self): if not (0.0 <= self.score <= 1.0): - raise ValueError(f"FineGrainedSubtask '{self.name}': score must be in [0, 1], got {self.score}") + raise ValueError(f"FineGrainedProgressObjective '{self.name}': score must be in [0, 1], got {self.score}") if self.logical not in ("all", "any", "choose"): raise ValueError( - f"FineGrainedSubtask '{self.name}': logical must be in ['all', 'any', 'choose'], got {self.logical}" + f"FineGrainedProgressObjective '{self.name}': logical must be in ['all', 'any', 'choose'], got {self.logical}" ) # Format the predicate groups into the canonical form and normalize the scores. @@ -159,13 +159,13 @@ def __post_init__(self): num_groups = len(self.canonical_predicate_groups) if self.logical == "choose": if self.K is None: - raise ValueError(f"FineGrainedSubtask '{self.name}': K is required when logical='choose'") + raise ValueError(f"FineGrainedProgressObjective '{self.name}': K is required when logical='choose'") if not (1 <= self.K <= num_groups): - raise ValueError(f"FineGrainedSubtask '{self.name}': K={self.K} but must be in [1, {num_groups}]") + raise ValueError(f"FineGrainedProgressObjective '{self.name}': K={self.K} but must be in [1, {num_groups}]") @property def group_names(self) -> list[str]: - """Returns the names of the groups in the FineGrainedSubtask.""" + """Returns the names of the groups in the FineGrainedProgressObjective.""" return list(self.canonical_predicate_groups.keys()) def get_chain(self, group_name: str) -> list[tuple[Callable, float]]: diff --git a/isaaclab_arena/tasks/fine_grained_subtask_tracking_state_machine.py b/isaaclab_arena/tasks/fine_grained_progress_tracker.py similarity index 65% rename from isaaclab_arena/tasks/fine_grained_subtask_tracking_state_machine.py rename to isaaclab_arena/tasks/fine_grained_progress_tracker.py index 6a142ce3e..360a1c150 100644 --- a/isaaclab_arena/tasks/fine_grained_subtask_tracking_state_machine.py +++ b/isaaclab_arena/tasks/fine_grained_progress_tracker.py @@ -12,9 +12,9 @@ from isaaclab.managers import EventTermCfg, TerminationTermCfg from isaaclab.utils import configclass -from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask +from isaaclab_arena.tasks.fine_grained_progress_objective import FineGrainedProgressObjective -_STATE_MACHINE_ATTR = "_fine_grained_subtask_state_machine" +_STATE_MACHINE_ATTR = "_fine_grained_progress_tracker" def _predicate_repr(pred) -> str: @@ -31,15 +31,15 @@ def _predicate_repr(pred) -> str: return f"{name}({', '.join(parts)})" if parts else name -class FineGrainedSubtaskRunner: - """State machine runner for a single FineGrainedSubtask object. +class FineGrainedProgressObjectiveRunner: + """State machine runner for a single FineGrainedProgressObjective object. Each runner is responsible for tracking the progress of all predicate_groups - within a FineGrainedSubtask object across all parallelenvironments. + within a FineGrainedProgressObjective object across all parallel environments. """ - def __init__(self, fine_grained_subtask: FineGrainedSubtask, num_envs: int, device): - self.fine_grained_subtask = fine_grained_subtask + def __init__(self, fine_grained_progress_objective: FineGrainedProgressObjective, num_envs: int, device): + self.fine_grained_progress_objective = fine_grained_progress_objective self.num_envs = num_envs self.device = device @@ -48,20 +48,20 @@ def __init__(self, fine_grained_subtask: FineGrainedSubtask, num_envs: int, devi self.group_score: dict[str, torch.Tensor] = {} self.group_complete: dict[str, torch.Tensor] = {} - for group_name in fine_grained_subtask.group_names: + for group_name in fine_grained_progress_objective.group_names: self.current_index[group_name] = torch.zeros(num_envs, dtype=torch.long, device=device) self.group_score[group_name] = torch.zeros(num_envs, dtype=torch.float32, device=device) self.group_complete[group_name] = torch.zeros(num_envs, dtype=torch.bool, device=device) def _compute_composite_task_gating_mask(self, env) -> torch.Tensor: - """Per-env mask of whether the FineGrainedSubtask is active. + """Per-env mask of whether the FineGrainedProgressObjective is active. The gating is used to determine when tracking of predicates should be active for composite tasks. """ # If no parent_subtask_idx -> always active (returns all True). - if self.fine_grained_subtask.parent_subtask_idx is None: + if self.fine_grained_progress_objective.parent_subtask_idx is None: return torch.ones(self.num_envs, dtype=torch.bool, device=self.device) # If no env._current_subtask_idx -> composite task is not sequential (returns all True). @@ -70,12 +70,12 @@ def _compute_composite_task_gating_mask(self, env) -> torch.Tensor: return torch.ones(self.num_envs, dtype=torch.bool, device=self.device) # Otherwise return True only for envs whose current - # parent-subtask index matches this FineGrainedSubtask's parent_subtask_idx. + # parent-subtask index matches this FineGrainedProgressObjective's parent_subtask_idx. if torch.is_tensor(current_idx): ci = current_idx.to(self.device) else: ci = torch.as_tensor(current_idx, device=self.device) - return ci == int(self.fine_grained_subtask.parent_subtask_idx) + return ci == int(self.fine_grained_progress_objective.parent_subtask_idx) def step(self, env, step_index: torch.Tensor | None) -> list[dict]: """Step the state machine runner for a single env.step. @@ -87,13 +87,13 @@ def step(self, env, step_index: torch.Tensor | None) -> list[dict]: # List of state transition events (events are emitted for an env when a predicate flips True) events: list[dict] = [] - # If the FineGrainedSubtask is not active for the composite task, return. + # If the FineGrainedProgressObjective is not active for the composite task, return. composite_task_gating_mask = self._compute_composite_task_gating_mask(env) if not bool(composite_task_gating_mask.any().item()): return events - # Step through each group of the FineGrainedSubtask. - for group_name, predicate_chain in self.fine_grained_subtask.canonical_predicate_groups.items(): + # Step through each group of the FineGrainedProgressObjective. + for group_name, predicate_chain in self.fine_grained_progress_objective.canonical_predicate_groups.items(): chain_length = len(predicate_chain) # Mask for which envs have advanced this step. advanced = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device) @@ -103,7 +103,7 @@ def step(self, env, step_index: torch.Tensor | None) -> list[dict]: # Envs should only be evaluated if: # 1) They are at the current predicate position # 2) They have not yet advanced this step - # 3) The FineGrainedSubtask is active for the composite task + # 3) The FineGrainedProgressObjective is active for the composite task at_position = (self.current_index[group_name] == chain_idx) & ~advanced & composite_task_gating_mask if not bool(at_position.any().item()): continue @@ -138,7 +138,7 @@ def step(self, env, step_index: torch.Tensor | None) -> list[dict]: events.append({ "env_idx": int(eid), "step": int(step_index[eid].item()) if step_index is not None else -1, - "fine_grained_subtask": self.fine_grained_subtask.name, + "fine_grained_progress_objective": self.fine_grained_progress_objective.name, "group": group_name, "predicate_index": chain_idx, "predicate_name": pred_name, @@ -153,47 +153,47 @@ def step(self, env, step_index: torch.Tensor | None) -> list[dict]: def reset(self, env_ids) -> None: """Reset the state machine runner for the provided envs.""" - for group_name in self.fine_grained_subtask.group_names: + for group_name in self.fine_grained_progress_objective.group_names: for eid in env_ids: self.current_index[group_name][eid] = 0 self.group_score[group_name][eid] = 0.0 self.group_complete[group_name][eid] = False def is_complete(self) -> torch.Tensor: - """Check if the FineGrainedSubtask is complete for all envs.""" + """Check if the FineGrainedProgressObjective is complete for all envs.""" - groups = self.fine_grained_subtask.group_names + groups = self.fine_grained_progress_objective.group_names stacked = torch.stack([self.group_complete[g] for g in groups], dim=1) - if self.fine_grained_subtask.logical == "all": + if self.fine_grained_progress_objective.logical == "all": return stacked.all(dim=1) - if self.fine_grained_subtask.logical == "any": + if self.fine_grained_progress_objective.logical == "any": return stacked.any(dim=1) - return stacked.sum(dim=1) >= int(self.fine_grained_subtask.K or 1) + return stacked.sum(dim=1) >= int(self.fine_grained_progress_objective.K or 1) def overall_score_per_env(self) -> torch.Tensor: - """Compute mean group score within this FineGrainedSubtask (in [0, 1]).""" + """Compute mean group score within this FineGrainedProgressObjective (in [0, 1]).""" - groups = self.fine_grained_subtask.group_names + groups = self.fine_grained_progress_objective.group_names stacked = torch.stack([self.group_score[g] for g in groups], dim=1) return stacked.mean(dim=1) -class FineGrainedSubtaskTrackingStateMachine: - """State machine that manages runners for all FineGrainedSubtasks. +class FineGrainedProgressTracker: + """State machine that manages runners for all FineGrainedProgressObjectives. Attributes: - fine_grained_subtasks: List of FineGrainedSubtasks to manage. - num_envs: Number of parallelenvironments. + fine_grained_progress_objectives: List of FineGrainedProgressObjectives to manage. + num_envs: Number of parallel environments. device: Device to manage the state machine on. - runners: List of runners for each FineGrainedSubtask. + runners: List of runners for each FineGrainedProgressObjective. _events: List of events for each environment. """ - def __init__(self, fine_grained_subtasks: list[FineGrainedSubtask], num_envs: int, device): - self.fine_grained_subtasks = fine_grained_subtasks + def __init__(self, fine_grained_progress_objectives: list[FineGrainedProgressObjective], num_envs: int, device): + self.fine_grained_progress_objectives = fine_grained_progress_objectives self.num_envs = num_envs self.device = device - self.runners = [FineGrainedSubtaskRunner(s, num_envs, device) for s in fine_grained_subtasks] + self.runners = [FineGrainedProgressObjectiveRunner(s, num_envs, device) for s in fine_grained_progress_objectives] self._events: list[list[dict]] = [[] for _ in range(num_envs)] def step(self, env, step_index: torch.Tensor | None) -> None: @@ -213,25 +213,25 @@ def reset(self, env_ids) -> None: self._events[eid] = [] def get_state(self) -> list[dict]: - """Get the state of each FineGrainedSubtask for all envs.""" + """Get the state of each FineGrainedProgressObjective for all envs.""" output: list[dict] = [] for env_idx in range(self.num_envs): # Build a per-env dict from each runner's state. - fine_grained_subtask_states: dict[str, dict] = {} + progress_objective_states: dict[str, dict] = {} overall_score = 0.0 all_complete = True for runner in self.runners: - fine_grained_subtask = runner.fine_grained_subtask + fine_grained_progress_objective = runner.fine_grained_progress_objective completed_groups = 0 - total_groups = len(fine_grained_subtask.group_names) + total_groups = len(fine_grained_progress_objective.group_names) active_predicates: dict[str, str | None] = {} # Compute the active predicates and completed groups. - for group_name in fine_grained_subtask.group_names: + for group_name in fine_grained_progress_objective.group_names: cur_group_index = int(runner.current_index[group_name][env_idx].item()) - predicate_chain = fine_grained_subtask.canonical_predicate_groups[group_name] + predicate_chain = fine_grained_progress_objective.canonical_predicate_groups[group_name] if cur_group_index >= len(predicate_chain): active_predicates[group_name] = None completed_groups += 1 @@ -239,21 +239,21 @@ def get_state(self) -> list[dict]: active_predicates[group_name] = _predicate_repr(predicate_chain[cur_group_index][0]) # Compute the overall score and completeness. - fine_grained_subtask_score = float(runner.overall_score_per_env()[env_idx].item()) + progress_objective_score = float(runner.overall_score_per_env()[env_idx].item()) is_complete = bool(runner.is_complete()[env_idx].item()) - fine_grained_subtask_states[fine_grained_subtask.name] = { + progress_objective_states[fine_grained_progress_objective.name] = { "completed_groups": completed_groups, "total_groups": total_groups, - "score": fine_grained_subtask_score, + "score": progress_objective_score, "is_complete": is_complete, "active_predicates": active_predicates, } - overall_score += fine_grained_subtask.score * fine_grained_subtask_score + overall_score += fine_grained_progress_objective.score * progress_objective_score all_complete = all_complete and is_complete # Add the per-env state dict to the output. output.append({ - "fine_grained_subtasks": fine_grained_subtask_states, + "fine_grained_progress_objectives": progress_objective_states, "overall_score": overall_score, "all_complete": all_complete, }) @@ -265,54 +265,54 @@ def get_events(self) -> list[list[dict]]: return [list(e) for e in self._events] -def _ensure_state_machine( - env, fine_grained_subtasks: list[FineGrainedSubtask] -) -> FineGrainedSubtaskTrackingStateMachine: - """Return the env's FineGrainedSubtaskTrackingStateMachine, lazily creating and caching it on first call.""" +def _ensure_progress_tracker( + env, fine_grained_progress_objectives: list[FineGrainedProgressObjective] +) -> FineGrainedProgressTracker: + """Return the env's FineGrainedProgressTracker, lazily creating and caching it on first call.""" - sm: FineGrainedSubtaskTrackingStateMachine | None = getattr(env, _STATE_MACHINE_ATTR, None) + sm: FineGrainedProgressTracker | None = getattr(env, _STATE_MACHINE_ATTR, None) if sm is None: - sm = FineGrainedSubtaskTrackingStateMachine( - fine_grained_subtasks=fine_grained_subtasks, num_envs=env.num_envs, device=env.device + sm = FineGrainedProgressTracker( + fine_grained_progress_objectives=fine_grained_progress_objectives, num_envs=env.num_envs, device=env.device ) setattr(env, _STATE_MACHINE_ATTR, sm) return sm -def fine_grained_subtask_step_func(env, fine_grained_subtasks: list[FineGrainedSubtask]) -> torch.Tensor: +def fine_grained_progress_step_func(env, fine_grained_progress_objectives: list[FineGrainedProgressObjective]) -> torch.Tensor: """Termination-term entry point. - Ticks the state machine, writes events and states to env.extras["fine_grained_subtask"], + Ticks the state machine, writes events and states to env.extras["fine_grained_progress"], and returns all-False so it does not contribute to termination. """ - sm = _ensure_state_machine(env, fine_grained_subtasks) + sm = _ensure_progress_tracker(env, fine_grained_progress_objectives) step_index = getattr(env, "episode_length_buf", None) sm.step(env, step_index=step_index) """ User-facing event/state information format: - env.extras["fine_grained_subtask"] = { + env.extras["fine_grained_progress"] = { "states": [ { - "fine_grained_subtasks": { - "": { + "fine_grained_progress_objectives": { + "": { "completed_groups": int, "total_groups": int, - "score": float, # 0..1, normalized within subtask + "score": float, # 0..1, normalized within objective "is_complete": bool, "active_predicates": {group: str | None}, }, ... }, - "overall_score": float, # weighted by FineGrainedSubtask.score + "overall_score": float, # weighted by FineGrainedProgressObjective.score "all_complete": bool, }, ... # one entry per env ], "events": [ - [{"step": int, "fine_grained_subtask": str, "group": str, + [{"step": int, "fine_grained_progress_objective": str, "group": str, "predicate_index": int, "predicate_name": str, "score_delta": float}, ...], ... # one list per env @@ -320,7 +320,7 @@ def fine_grained_subtask_step_func(env, fine_grained_subtasks: list[FineGrainedS } """ - env.extras["fine_grained_subtask"] = { + env.extras["fine_grained_progress"] = { "states": sm.get_state(), "events": sm.get_events(), } @@ -329,13 +329,13 @@ def fine_grained_subtask_step_func(env, fine_grained_subtasks: list[FineGrainedS return torch.zeros(env.num_envs, dtype=torch.bool, device=env.device) -def fine_grained_subtask_reset_func(env, env_ids, fine_grained_subtasks: list[FineGrainedSubtask]) -> None: +def fine_grained_progress_reset_func(env, env_ids, fine_grained_progress_objectives: list[FineGrainedProgressObjective]) -> None: """Reset-event entry point. Resets the state machine whenever the Lab env is reset. """ - sm = _ensure_state_machine(env, fine_grained_subtasks) + sm = _ensure_progress_tracker(env, fine_grained_progress_objectives) if env_ids is None: env_ids = list(range(env.num_envs)) elif torch.is_tensor(env_ids): @@ -344,29 +344,29 @@ def fine_grained_subtask_reset_func(env, env_ids, fine_grained_subtasks: list[Fi @configclass -class FineGrainedSubtaskEventsCfg: - reset_fine_grained_subtasks: EventTermCfg = MISSING +class FineGrainedProgressObjectiveEventsCfg: + reset_fine_grained_progress_objectives: EventTermCfg = MISSING @configclass -class FineGrainedSubtaskTerminationsCfg: - fine_grained_subtask_step: TerminationTermCfg = MISSING +class FineGrainedProgressObjectiveTerminationsCfg: + fine_grained_progress_step: TerminationTermCfg = MISSING -def make_fine_grained_subtask_events_cfg(fine_grained_subtasks: list[FineGrainedSubtask]) -> Any: - return FineGrainedSubtaskEventsCfg( - reset_fine_grained_subtasks=EventTermCfg( - func=fine_grained_subtask_reset_func, +def make_fine_grained_progress_objective_events_cfg(fine_grained_progress_objectives: list[FineGrainedProgressObjective]) -> Any: + return FineGrainedProgressObjectiveEventsCfg( + reset_fine_grained_progress_objectives=EventTermCfg( + func=fine_grained_progress_reset_func, mode="reset", - params={"fine_grained_subtasks": fine_grained_subtasks}, + params={"fine_grained_progress_objectives": fine_grained_progress_objectives}, ) ) -def make_fine_grained_subtask_termination_cfg(fine_grained_subtasks: list[FineGrainedSubtask]) -> Any: - return FineGrainedSubtaskTerminationsCfg( - fine_grained_subtask_step=TerminationTermCfg( - func=fine_grained_subtask_step_func, - params={"fine_grained_subtasks": fine_grained_subtasks}, +def make_fine_grained_progress_objective_termination_cfg(fine_grained_progress_objectives: list[FineGrainedProgressObjective]) -> Any: + return FineGrainedProgressObjectiveTerminationsCfg( + fine_grained_progress_step=TerminationTermCfg( + func=fine_grained_progress_step_func, + params={"fine_grained_progress_objectives": fine_grained_progress_objectives}, ) ) diff --git a/isaaclab_arena/tasks/task_base.py b/isaaclab_arena/tasks/task_base.py index 91b65b328..e2f3bb606 100644 --- a/isaaclab_arena/tasks/task_base.py +++ b/isaaclab_arena/tasks/task_base.py @@ -11,10 +11,10 @@ from isaaclab_arena.embodiments.common.arm_mode import ArmMode from isaaclab_arena.metrics.metric_base import MetricBase -from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask -from isaaclab_arena.tasks.fine_grained_subtask_tracking_state_machine import ( - make_fine_grained_subtask_events_cfg, - make_fine_grained_subtask_termination_cfg, +from isaaclab_arena.tasks.fine_grained_progress_objective import FineGrainedProgressObjective +from isaaclab_arena.tasks.fine_grained_progress_tracker import ( + make_fine_grained_progress_objective_events_cfg, + make_fine_grained_progress_objective_termination_cfg, ) @@ -68,17 +68,17 @@ def get_episode_length_s(self) -> float | None: def get_task_description(self) -> str | None: return self.task_description - def get_fine_grained_subtasks(self) -> list[FineGrainedSubtask]: + def get_fine_grained_progress_objectives(self) -> list[FineGrainedProgressObjective]: return [] - def get_fine_grained_subtask_events_cfg(self) -> Any: - fine_grained_subtasks = self.get_fine_grained_subtasks() - if not fine_grained_subtasks: + def get_fine_grained_progress_objective_events_cfg(self) -> Any: + fine_grained_progress_objectives = self.get_fine_grained_progress_objectives() + if not fine_grained_progress_objectives: return None - return make_fine_grained_subtask_events_cfg(fine_grained_subtasks) + return make_fine_grained_progress_objective_events_cfg(fine_grained_progress_objectives) - def get_fine_grained_subtask_termination_cfg(self) -> Any: - fine_grained_subtasks = self.get_fine_grained_subtasks() - if not fine_grained_subtasks: + def get_fine_grained_progress_objective_termination_cfg(self) -> Any: + fine_grained_progress_objectives = self.get_fine_grained_progress_objectives() + if not fine_grained_progress_objectives: return None - return make_fine_grained_subtask_termination_cfg(fine_grained_subtasks) + return make_fine_grained_progress_objective_termination_cfg(fine_grained_progress_objectives) diff --git a/isaaclab_arena/tests/test_fine_grained_subtask.py b/isaaclab_arena/tests/test_fine_grained_progress_objective_tracking.py similarity index 62% rename from isaaclab_arena/tests/test_fine_grained_subtask.py rename to isaaclab_arena/tests/test_fine_grained_progress_objective_tracking.py index b976c0918..2f3df659b 100644 --- a/isaaclab_arena/tests/test_fine_grained_subtask.py +++ b/isaaclab_arena/tests/test_fine_grained_progress_objective_tracking.py @@ -49,13 +49,13 @@ def _advance_step(env, n: int = 1): def _test_predicate_groups_single_callable(simulation_app) -> bool: """A bare predicate becomes a default-named group with weight 1.0.""" - from isaaclab_arena.tasks.fine_grained_subtask import DEFAULT_GROUP_NAME, FineGrainedSubtask + from isaaclab_arena.tasks.fine_grained_progress_objective import DEFAULT_GROUP_NAME, FineGrainedProgressObjective try: pred = _MockPredicate(num_envs=1) - fgs = FineGrainedSubtask(name="t", predicate_groups=pred) - assert fgs.group_names == [DEFAULT_GROUP_NAME] - chain = fgs.get_chain(DEFAULT_GROUP_NAME) + fgpo = FineGrainedProgressObjective(name="t", predicate_groups=pred) + assert fgpo.group_names == [DEFAULT_GROUP_NAME] + chain = fgpo.get_chain(DEFAULT_GROUP_NAME) assert len(chain) == 1 assert chain[0][0] is pred assert abs(chain[0][1] - 1.0) < SCORE_TOL @@ -68,12 +68,12 @@ def _test_predicate_groups_single_callable(simulation_app) -> bool: def _test_predicate_groups_list_of_callables(simulation_app) -> bool: """A list of callables becomes a single group with normalized equal scores.""" - from isaaclab_arena.tasks.fine_grained_subtask import DEFAULT_GROUP_NAME, FineGrainedSubtask + from isaaclab_arena.tasks.fine_grained_progress_objective import DEFAULT_GROUP_NAME, FineGrainedProgressObjective try: preds = [_MockPredicate(num_envs=1, name=f"p{i}") for i in range(3)] - fgs = FineGrainedSubtask(name="t", predicate_groups=preds) - chain = fgs.get_chain(DEFAULT_GROUP_NAME) + fgpo = FineGrainedProgressObjective(name="t", predicate_groups=preds) + chain = fgpo.get_chain(DEFAULT_GROUP_NAME) assert [c[0] for c in chain] == preds # Equal scores normalize to 0.33 each, summing to 1.0. for _, score in chain: @@ -88,13 +88,13 @@ def _test_predicate_groups_list_of_callables(simulation_app) -> bool: def _test_predicate_groups_weighted_tuples(simulation_app) -> bool: """Explicit (callable, score) tuples are normalized to sum to 1.0 within a group.""" - from isaaclab_arena.tasks.fine_grained_subtask import DEFAULT_GROUP_NAME, FineGrainedSubtask + from isaaclab_arena.tasks.fine_grained_progress_objective import DEFAULT_GROUP_NAME, FineGrainedProgressObjective try: p1 = _MockPredicate(num_envs=1, name="p1") p2 = _MockPredicate(num_envs=1, name="p2") - fgs = FineGrainedSubtask(name="t", predicate_groups=[(p1, 1.0), (p2, 3.0)]) - chain = fgs.get_chain(DEFAULT_GROUP_NAME) + fgpo = FineGrainedProgressObjective(name="t", predicate_groups=[(p1, 1.0), (p2, 3.0)]) + chain = fgpo.get_chain(DEFAULT_GROUP_NAME) # 1.0/4.0 = 0.25, 3.0/4.0 = 0.75 assert abs(chain[0][1] - 0.25) < SCORE_TOL assert abs(chain[1][1] - 0.75) < SCORE_TOL @@ -107,13 +107,13 @@ def _test_predicate_groups_weighted_tuples(simulation_app) -> bool: def _test_predicate_groups_dict_groups(simulation_app) -> bool: """Dict input gives one group per key and each group's scores are normalized independently.""" - from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask + from isaaclab_arena.tasks.fine_grained_progress_objective import FineGrainedProgressObjective try: p_a1 = _MockPredicate(num_envs=1, name="a1") p_a2 = _MockPredicate(num_envs=1, name="a2") p_b = _MockPredicate(num_envs=1, name="b") - fgs = FineGrainedSubtask( + fgpo = FineGrainedProgressObjective( name="t", predicate_groups={ "obj_a": [p_a1, p_a2], @@ -121,9 +121,9 @@ def _test_predicate_groups_dict_groups(simulation_app) -> bool: }, logical="all", ) - assert set(fgs.group_names) == {"obj_a", "obj_b"} - a_chain = fgs.get_chain("obj_a") - b_chain = fgs.get_chain("obj_b") + assert set(fgpo.group_names) == {"obj_a", "obj_b"} + a_chain = fgpo.get_chain("obj_a") + b_chain = fgpo.get_chain("obj_b") assert len(a_chain) == 2 assert len(b_chain) == 1 # obj_a's equal scores sum to 1.0. @@ -139,19 +139,19 @@ def _test_predicate_groups_dict_groups(simulation_app) -> bool: def _test_predicate_groups_rejects_invalid_inputs(simulation_app) -> bool: """Empty containers and non-callable entries should raise error.""" - from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask + from isaaclab_arena.tasks.fine_grained_progress_objective import FineGrainedProgressObjective try: for bad in ([], {}, 42, "string"): try: - FineGrainedSubtask(name="t", predicate_groups=bad) + FineGrainedProgressObjective(name="t", predicate_groups=bad) except (ValueError, TypeError): continue print(f"Expected error for input {bad!r}") return False # logical=choose without K should raise error. try: - FineGrainedSubtask( + FineGrainedProgressObjective( name="t", predicate_groups=_MockPredicate(num_envs=1), logical="choose", @@ -169,22 +169,22 @@ def _test_predicate_groups_rejects_invalid_inputs(simulation_app) -> bool: def _test_state_machine_advances_sequentially(simulation_app) -> bool: - """A single FineGrainedSubtask with a 3 predicate chain advances one step per satisfied predicate.""" - from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask - from isaaclab_arena.tasks.fine_grained_subtask_tracking_state_machine import FineGrainedSubtaskTrackingStateMachine + """A single FineGrainedProgressObjective with a 3 predicate chain advances one step per satisfied predicate.""" + from isaaclab_arena.tasks.fine_grained_progress_objective import FineGrainedProgressObjective + from isaaclab_arena.tasks.fine_grained_progress_tracker import FineGrainedProgressTracker try: env = _MockEnv(num_envs=1) preds = [_MockPredicate(num_envs=1, name=f"p{i}") for i in range(3)] - fgs = FineGrainedSubtask(name="lift", predicate_groups=preds) - sm = FineGrainedSubtaskTrackingStateMachine(fine_grained_subtasks=[fgs], num_envs=1, device="cpu") + fgpo = FineGrainedProgressObjective(name="lift", predicate_groups=preds) + sm = FineGrainedProgressTracker(fine_grained_progress_objectives=[fgpo], num_envs=1, device="cpu") sm.reset([0]) # Step 1: p0 True while p1, p2 still False. Advance to index 1. preds[0].set([True]) _advance_step(env) sm.step(env, step_index=env.episode_length_buf) - state = sm.get_state()[0]["fine_grained_subtasks"]["lift"] + state = sm.get_state()[0]["fine_grained_progress_objectives"]["lift"] assert state["completed_groups"] == 0 # 3-predicate chain not done until all 3 assert not state["is_complete"] events = sm.get_events()[0] @@ -198,11 +198,11 @@ def _test_state_machine_advances_sequentially(simulation_app) -> bool: events = sm.get_events()[0] assert len(events) == 2 and events[-1]["predicate_index"] == 1 - # Step 3: p2 True, subtask complete. + # Step 3: p2 True, objective complete. preds[2].set([True]) _advance_step(env) sm.step(env, step_index=env.episode_length_buf) - state = sm.get_state()[0]["fine_grained_subtasks"]["lift"] + state = sm.get_state()[0]["fine_grained_progress_objectives"]["lift"] assert state["is_complete"] assert state["completed_groups"] == 1 assert abs(state["score"] - 1.0) < SCORE_TOL @@ -217,14 +217,14 @@ def _test_state_machine_advances_sequentially(simulation_app) -> bool: def _test_state_machine_ignores_out_of_order_success(simulation_app) -> bool: """If a later predicate fires first, it's ignored until preceding ones have advanced.""" - from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask - from isaaclab_arena.tasks.fine_grained_subtask_tracking_state_machine import FineGrainedSubtaskTrackingStateMachine + from isaaclab_arena.tasks.fine_grained_progress_objective import FineGrainedProgressObjective + from isaaclab_arena.tasks.fine_grained_progress_tracker import FineGrainedProgressTracker try: env = _MockEnv(num_envs=1) preds = [_MockPredicate(num_envs=1, name=f"p{i}") for i in range(3)] - fgs = FineGrainedSubtask(name="lift", predicate_groups=preds) - sm = FineGrainedSubtaskTrackingStateMachine(fine_grained_subtasks=[fgs], num_envs=1, device="cpu") + fgpo = FineGrainedProgressObjective(name="lift", predicate_groups=preds) + sm = FineGrainedProgressTracker(fine_grained_progress_objectives=[fgpo], num_envs=1, device="cpu") sm.reset([0]) # p0 stays False and p1, p2 True. No progress should be made. @@ -233,7 +233,7 @@ def _test_state_machine_ignores_out_of_order_success(simulation_app) -> bool: preds[2].set([True]) _advance_step(env) sm.step(env, step_index=env.episode_length_buf) - state = sm.get_state()[0]["fine_grained_subtasks"]["lift"] + state = sm.get_state()[0]["fine_grained_progress_objectives"]["lift"] assert state["completed_groups"] == 0 assert not state["is_complete"] assert state["score"] == 0.0 @@ -244,7 +244,7 @@ def _test_state_machine_ignores_out_of_order_success(simulation_app) -> bool: for _ in range(3): _advance_step(env) sm.step(env, step_index=env.episode_length_buf) - state = sm.get_state()[0]["fine_grained_subtasks"]["lift"] + state = sm.get_state()[0]["fine_grained_progress_objectives"]["lift"] assert state["is_complete"] assert state["completed_groups"] == 1 except Exception as e: @@ -256,31 +256,31 @@ def _test_state_machine_ignores_out_of_order_success(simulation_app) -> bool: def _test_state_machine_logical_any(simulation_app) -> bool: """Two parallel groups with logical=any complete as soon as either one finishes.""" - from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask - from isaaclab_arena.tasks.fine_grained_subtask_tracking_state_machine import FineGrainedSubtaskTrackingStateMachine + from isaaclab_arena.tasks.fine_grained_progress_objective import FineGrainedProgressObjective + from isaaclab_arena.tasks.fine_grained_progress_tracker import FineGrainedProgressTracker try: env = _MockEnv(num_envs=1) p_a = _MockPredicate(num_envs=1, name="a") p_b = _MockPredicate(num_envs=1, name="b") - fgs = FineGrainedSubtask( + fgpo = FineGrainedProgressObjective( name="either", predicate_groups={"a": p_a, "b": p_b}, logical="any", ) - sm = FineGrainedSubtaskTrackingStateMachine(fine_grained_subtasks=[fgs], num_envs=1, device="cpu") + sm = FineGrainedProgressTracker(fine_grained_progress_objectives=[fgpo], num_envs=1, device="cpu") sm.reset([0]) # Neither group complete -> not done. _advance_step(env) sm.step(env, step_index=env.episode_length_buf) - assert not sm.get_state()[0]["fine_grained_subtasks"]["either"]["is_complete"] + assert not sm.get_state()[0]["fine_grained_progress_objectives"]["either"]["is_complete"] # Group p_a completes -> done. p_a.set([True]) _advance_step(env) sm.step(env, step_index=env.episode_length_buf) - assert sm.get_state()[0]["fine_grained_subtasks"]["either"]["is_complete"] + assert sm.get_state()[0]["fine_grained_progress_objectives"]["either"]["is_complete"] except Exception as e: print(f"Error: {e}") traceback.print_exc() @@ -290,32 +290,32 @@ def _test_state_machine_logical_any(simulation_app) -> bool: def _test_state_machine_logical_all(simulation_app) -> bool: """Two groups with logical=all complete once all groups are complete.""" - from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask - from isaaclab_arena.tasks.fine_grained_subtask_tracking_state_machine import FineGrainedSubtaskTrackingStateMachine + from isaaclab_arena.tasks.fine_grained_progress_objective import FineGrainedProgressObjective + from isaaclab_arena.tasks.fine_grained_progress_tracker import FineGrainedProgressTracker try: env = _MockEnv(num_envs=1) p_a = _MockPredicate(num_envs=1, name="a") p_b = _MockPredicate(num_envs=1, name="b") - fgs = FineGrainedSubtask( + fgpo = FineGrainedProgressObjective( name="both", predicate_groups={"a": p_a, "b": p_b}, logical="all", ) - sm = FineGrainedSubtaskTrackingStateMachine(fine_grained_subtasks=[fgs], num_envs=1, device="cpu") + sm = FineGrainedProgressTracker(fine_grained_progress_objectives=[fgpo], num_envs=1, device="cpu") sm.reset([0]) # Only p_a completes -> still not done. p_a.set([True]) _advance_step(env) sm.step(env, step_index=env.episode_length_buf) - assert not sm.get_state()[0]["fine_grained_subtasks"]["both"]["is_complete"] + assert not sm.get_state()[0]["fine_grained_progress_objectives"]["both"]["is_complete"] # p_b also completes -> done. p_b.set([True]) _advance_step(env) sm.step(env, step_index=env.episode_length_buf) - assert sm.get_state()[0]["fine_grained_subtasks"]["both"]["is_complete"] + assert sm.get_state()[0]["fine_grained_progress_objectives"]["both"]["is_complete"] except Exception as e: print(f"Error: {e}") traceback.print_exc() @@ -325,34 +325,34 @@ def _test_state_machine_logical_all(simulation_app) -> bool: def _test_state_machine_logical_choose(simulation_app) -> bool: """Three groups with logical=choose and K=2 complete once any two groups are complete.""" - from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask - from isaaclab_arena.tasks.fine_grained_subtask_tracking_state_machine import FineGrainedSubtaskTrackingStateMachine + from isaaclab_arena.tasks.fine_grained_progress_objective import FineGrainedProgressObjective + from isaaclab_arena.tasks.fine_grained_progress_tracker import FineGrainedProgressTracker try: env = _MockEnv(num_envs=1) p_a = _MockPredicate(num_envs=1, name="a") p_b = _MockPredicate(num_envs=1, name="b") p_c = _MockPredicate(num_envs=1, name="c") - fgs = FineGrainedSubtask( + fgpo = FineGrainedProgressObjective( name="any_two", predicate_groups={"a": p_a, "b": p_b, "c": p_c}, logical="choose", K=2, ) - sm = FineGrainedSubtaskTrackingStateMachine(fine_grained_subtasks=[fgs], num_envs=1, device="cpu") + sm = FineGrainedProgressTracker(fine_grained_progress_objectives=[fgpo], num_envs=1, device="cpu") sm.reset([0]) # Only p_a group complete -> not done. p_a.set([True]) _advance_step(env) sm.step(env, step_index=env.episode_length_buf) - assert not sm.get_state()[0]["fine_grained_subtasks"]["any_two"]["is_complete"] + assert not sm.get_state()[0]["fine_grained_progress_objectives"]["any_two"]["is_complete"] # p_b also complete -> done. p_b.set([True]) _advance_step(env) sm.step(env, step_index=env.episode_length_buf) - assert sm.get_state()[0]["fine_grained_subtasks"]["any_two"]["is_complete"] + assert sm.get_state()[0]["fine_grained_progress_objectives"]["any_two"]["is_complete"] except Exception as e: print(f"Error: {e}") traceback.print_exc() @@ -362,14 +362,14 @@ def _test_state_machine_logical_choose(simulation_app) -> bool: def _test_state_machine_reset_clears_state(simulation_app) -> bool: """Resetting an env_id zeroes its progress and event log, but leaves other envs alone.""" - from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask - from isaaclab_arena.tasks.fine_grained_subtask_tracking_state_machine import FineGrainedSubtaskTrackingStateMachine + from isaaclab_arena.tasks.fine_grained_progress_objective import FineGrainedProgressObjective + from isaaclab_arena.tasks.fine_grained_progress_tracker import FineGrainedProgressTracker try: env = _MockEnv(num_envs=2) preds = [_MockPredicate(num_envs=2, name=f"p{i}") for i in range(2)] - fgs = FineGrainedSubtask(name="t", predicate_groups=preds) - sm = FineGrainedSubtaskTrackingStateMachine(fine_grained_subtasks=[fgs], num_envs=2, device="cpu") + fgpo = FineGrainedProgressObjective(name="t", predicate_groups=preds) + sm = FineGrainedProgressTracker(fine_grained_progress_objectives=[fgpo], num_envs=2, device="cpu") sm.reset([0, 1]) # Set env 0 to fully complete. @@ -381,16 +381,16 @@ def _test_state_machine_reset_clears_state(simulation_app) -> bool: sm.step(env, step_index=env.episode_length_buf) state = sm.get_state() - assert state[0]["fine_grained_subtasks"]["t"]["is_complete"] - assert not state[1]["fine_grained_subtasks"]["t"]["is_complete"] + assert state[0]["fine_grained_progress_objectives"]["t"]["is_complete"] + assert not state[1]["fine_grained_progress_objectives"]["t"]["is_complete"] assert len(sm.get_events()[0]) >= 2 assert len(sm.get_events()[1]) >= 1 # Reset only env 0. sm.reset([0]) state = sm.get_state() - assert not state[0]["fine_grained_subtasks"]["t"]["is_complete"] - assert state[0]["fine_grained_subtasks"]["t"]["score"] == 0.0 + assert not state[0]["fine_grained_progress_objectives"]["t"]["is_complete"] + assert state[0]["fine_grained_progress_objectives"]["t"]["score"] == 0.0 assert sm.get_events()[0] == [] # env 1 untouched. assert len(sm.get_events()[1]) >= 1 @@ -402,23 +402,23 @@ def _test_state_machine_reset_clears_state(simulation_app) -> bool: def _test_gating_advance_when_parent_subtask_idx_matches(simulation_app) -> bool: - """A FineGrainedSubtask with parent_subtask_idx=N advances when the env's _current_subtask_idx=N.""" - from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask - from isaaclab_arena.tasks.fine_grained_subtask_tracking_state_machine import FineGrainedSubtaskTrackingStateMachine + """A FineGrainedProgressObjective with parent_subtask_idx=N advances when the env's _current_subtask_idx=N.""" + from isaaclab_arena.tasks.fine_grained_progress_objective import FineGrainedProgressObjective + from isaaclab_arena.tasks.fine_grained_progress_tracker import FineGrainedProgressTracker try: env = _MockEnv(num_envs=1) env._current_subtask_idx = [1] pred = _MockPredicate(num_envs=1, name="p") - fgs = FineGrainedSubtask(name="t", predicate_groups=pred, parent_subtask_idx=1) - sm = FineGrainedSubtaskTrackingStateMachine(fine_grained_subtasks=[fgs], num_envs=1, device="cpu") + fgpo = FineGrainedProgressObjective(name="t", predicate_groups=pred, parent_subtask_idx=1) + sm = FineGrainedProgressTracker(fine_grained_progress_objectives=[fgpo], num_envs=1, device="cpu") sm.reset([0]) pred.set([True]) _advance_step(env) sm.step(env, step_index=env.episode_length_buf) - assert sm.get_state()[0]["fine_grained_subtasks"]["t"]["is_complete"] + assert sm.get_state()[0]["fine_grained_progress_objectives"]["t"]["is_complete"] assert len(sm.get_events()[0]) == 1 except Exception as e: print(f"Error: {e}") @@ -428,32 +428,32 @@ def _test_gating_advance_when_parent_subtask_idx_matches(simulation_app) -> bool def _test_gating_blocked_when_parent_subtask_idx_mismatches(simulation_app) -> bool: - """A FineGrainedSubtask with parent_subtask_idx=N doesn't advance when the env's _current_subtask_idx!=N.""" - from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask - from isaaclab_arena.tasks.fine_grained_subtask_tracking_state_machine import FineGrainedSubtaskTrackingStateMachine + """A FineGrainedProgressObjective with parent_subtask_idx=N doesn't advance when the env's _current_subtask_idx!=N.""" + from isaaclab_arena.tasks.fine_grained_progress_objective import FineGrainedProgressObjective + from isaaclab_arena.tasks.fine_grained_progress_tracker import FineGrainedProgressTracker try: env = _MockEnv(num_envs=1) env._current_subtask_idx = [0] pred = _MockPredicate(num_envs=1, name="p") - fgs = FineGrainedSubtask(name="t", predicate_groups=pred, parent_subtask_idx=1) - sm = FineGrainedSubtaskTrackingStateMachine(fine_grained_subtasks=[fgs], num_envs=1, device="cpu") + fgpo = FineGrainedProgressObjective(name="t", predicate_groups=pred, parent_subtask_idx=1) + sm = FineGrainedProgressTracker(fine_grained_progress_objectives=[fgpo], num_envs=1, device="cpu") sm.reset([0]) - # Predicate True, but the parent isn't at this FGS's index yet. + # Predicate True, but the parent isn't at this FGPO's index yet. pred.set([True]) _advance_step(env) sm.step(env, step_index=env.episode_length_buf) - assert not sm.get_state()[0]["fine_grained_subtasks"]["t"]["is_complete"] - assert sm.get_state()[0]["fine_grained_subtasks"]["t"]["score"] == 0.0 + assert not sm.get_state()[0]["fine_grained_progress_objectives"]["t"]["is_complete"] + assert sm.get_state()[0]["fine_grained_progress_objectives"]["t"]["score"] == 0.0 assert len(sm.get_events()[0]) == 0 - # Parent advances to this FGS's index, state machine advances. + # Parent advances to this FGPO's index, state machine advances. env._current_subtask_idx = [1] _advance_step(env) sm.step(env, step_index=env.episode_length_buf) - assert sm.get_state()[0]["fine_grained_subtasks"]["t"]["is_complete"] + assert sm.get_state()[0]["fine_grained_progress_objectives"]["t"]["is_complete"] assert len(sm.get_events()[0]) == 1 except Exception as e: print(f"Error: {e}") @@ -463,11 +463,11 @@ def _test_gating_blocked_when_parent_subtask_idx_mismatches(simulation_app) -> b def _test_gating_sequential_task_end_to_end(simulation_app) -> bool: - """Two FGSs with different parent subtask indices. The parent's - _current_subtask_idx advances over time. Each FGS only progresses + """Two FGPOs with different parent subtask indices. The parent's + _current_subtask_idx advances over time. Each FGPO only progresses during its active window.""" - from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask - from isaaclab_arena.tasks.fine_grained_subtask_tracking_state_machine import FineGrainedSubtaskTrackingStateMachine + from isaaclab_arena.tasks.fine_grained_progress_objective import FineGrainedProgressObjective + from isaaclab_arena.tasks.fine_grained_progress_tracker import FineGrainedProgressTracker try: env = _MockEnv(num_envs=1) @@ -475,9 +475,9 @@ def _test_gating_sequential_task_end_to_end(simulation_app) -> bool: pred_a = _MockPredicate(num_envs=1, name="a") pred_b = _MockPredicate(num_envs=1, name="b") - fgs_a = FineGrainedSubtask(name="a", predicate_groups=pred_a, parent_subtask_idx=0) - fgs_b = FineGrainedSubtask(name="b", predicate_groups=pred_b, parent_subtask_idx=1) - sm = FineGrainedSubtaskTrackingStateMachine(fine_grained_subtasks=[fgs_a, fgs_b], num_envs=1, device="cpu") + fgpo_a = FineGrainedProgressObjective(name="a", predicate_groups=pred_a, parent_subtask_idx=0) + fgpo_b = FineGrainedProgressObjective(name="b", predicate_groups=pred_b, parent_subtask_idx=1) + sm = FineGrainedProgressTracker(fine_grained_progress_objectives=[fgpo_a, fgpo_b], num_envs=1, device="cpu") sm.reset([0]) # Both predicates True, but only pred_a is active. @@ -485,14 +485,14 @@ def _test_gating_sequential_task_end_to_end(simulation_app) -> bool: pred_b.set([True]) _advance_step(env) sm.step(env, step_index=env.episode_length_buf) - assert sm.get_state()[0]["fine_grained_subtasks"]["a"]["is_complete"] - assert not sm.get_state()[0]["fine_grained_subtasks"]["b"]["is_complete"] + assert sm.get_state()[0]["fine_grained_progress_objectives"]["a"]["is_complete"] + assert not sm.get_state()[0]["fine_grained_progress_objectives"]["b"]["is_complete"] # Advances to subtask 1 so pred_b is now active. env._current_subtask_idx = [1] _advance_step(env) sm.step(env, step_index=env.episode_length_buf) - assert sm.get_state()[0]["fine_grained_subtasks"]["b"]["is_complete"] + assert sm.get_state()[0]["fine_grained_progress_objectives"]["b"]["is_complete"] except Exception as e: print(f"Error: {e}") traceback.print_exc() @@ -501,22 +501,22 @@ def _test_gating_sequential_task_end_to_end(simulation_app) -> bool: def _test_gating_noop_when_env_has_no_current_subtask_idx(simulation_app) -> bool: - """For unordered composite tasks gating is a no-op and all FGSs advance whenever their predicates are True.""" - from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask - from isaaclab_arena.tasks.fine_grained_subtask_tracking_state_machine import FineGrainedSubtaskTrackingStateMachine + """For unordered composite tasks gating is a no-op and all FGPOs advance whenever their predicates are True.""" + from isaaclab_arena.tasks.fine_grained_progress_objective import FineGrainedProgressObjective + from isaaclab_arena.tasks.fine_grained_progress_tracker import FineGrainedProgressTracker try: env = _MockEnv(num_envs=1) pred = _MockPredicate(num_envs=1, name="p") - fgs = FineGrainedSubtask(name="t", predicate_groups=pred, parent_subtask_idx=1) - sm = FineGrainedSubtaskTrackingStateMachine(fine_grained_subtasks=[fgs], num_envs=1, device="cpu") + fgpo = FineGrainedProgressObjective(name="t", predicate_groups=pred, parent_subtask_idx=1) + sm = FineGrainedProgressTracker(fine_grained_progress_objectives=[fgpo], num_envs=1, device="cpu") sm.reset([0]) pred.set([True]) _advance_step(env) sm.step(env, step_index=env.episode_length_buf) - assert sm.get_state()[0]["fine_grained_subtasks"]["t"]["is_complete"] + assert sm.get_state()[0]["fine_grained_progress_objectives"]["t"]["is_complete"] assert len(sm.get_events()[0]) == 1 except Exception as e: print(f"Error: {e}") @@ -526,47 +526,47 @@ def _test_gating_noop_when_env_has_no_current_subtask_idx(simulation_app) -> boo def _test_step_func_publishes_to_extras_and_returns_no_termination(simulation_app) -> bool: - """fine_grained_subtask_step_func writes env.extras and returns all-False.""" - from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask - from isaaclab_arena.tasks.fine_grained_subtask_tracking_state_machine import ( - fine_grained_subtask_reset_func, - fine_grained_subtask_step_func, + """fine_grained_progress_step_func writes env.extras and returns all-False.""" + from isaaclab_arena.tasks.fine_grained_progress_objective import FineGrainedProgressObjective + from isaaclab_arena.tasks.fine_grained_progress_tracker import ( + fine_grained_progress_reset_func, + fine_grained_progress_step_func, ) try: env = _MockEnv(num_envs=2) pred = _MockPredicate(num_envs=2, name="p") - fgs = FineGrainedSubtask(name="t", predicate_groups=pred) + fgpo = FineGrainedProgressObjective(name="t", predicate_groups=pred) - fine_grained_subtask_reset_func(env, env_ids=[0, 1], fine_grained_subtasks=[fgs]) + fine_grained_progress_reset_func(env, env_ids=[0, 1], fine_grained_progress_objectives=[fgpo]) # Step with predicate=False, state machine ticks but no transitions. - result = fine_grained_subtask_step_func(env, fine_grained_subtasks=[fgs]) + result = fine_grained_progress_step_func(env, fine_grained_progress_objectives=[fgpo]) assert result.tolist() == [False, False] - assert "fine_grained_subtask" in env.extras - assert len(env.extras["fine_grained_subtask"]["states"]) == 2 - assert env.extras["fine_grained_subtask"]["events"] == [[], []] - assert not env.extras["fine_grained_subtask"]["states"][0]["fine_grained_subtasks"]["t"]["is_complete"] + assert "fine_grained_progress" in env.extras + assert len(env.extras["fine_grained_progress"]["states"]) == 2 + assert env.extras["fine_grained_progress"]["events"] == [[], []] + assert not env.extras["fine_grained_progress"]["states"][0]["fine_grained_progress_objectives"]["t"]["is_complete"] # Step with env 0 predicate True, env 0 completes, env 1 does not. pred.set([True, False]) _advance_step(env) - result = fine_grained_subtask_step_func(env, fine_grained_subtasks=[fgs]) + result = fine_grained_progress_step_func(env, fine_grained_progress_objectives=[fgpo]) assert result.tolist() == [False, False] - states = env.extras["fine_grained_subtask"]["states"] - events = env.extras["fine_grained_subtask"]["events"] - assert states[0]["fine_grained_subtasks"]["t"]["is_complete"] - assert not states[1]["fine_grained_subtasks"]["t"]["is_complete"] + states = env.extras["fine_grained_progress"]["states"] + events = env.extras["fine_grained_progress"]["events"] + assert states[0]["fine_grained_progress_objectives"]["t"]["is_complete"] + assert not states[1]["fine_grained_progress_objectives"]["t"]["is_complete"] assert len(events[0]) == 1 assert len(events[1]) == 0 # Reset env 0, env 1 untouched. pred.set([False, False]) - fine_grained_subtask_reset_func(env, env_ids=[0], fine_grained_subtasks=[fgs]) - result = fine_grained_subtask_step_func(env, fine_grained_subtasks=[fgs]) - states = env.extras["fine_grained_subtask"]["states"] - assert not states[0]["fine_grained_subtasks"]["t"]["is_complete"] - assert states[0]["fine_grained_subtasks"]["t"]["score"] == 0.0 + fine_grained_progress_reset_func(env, env_ids=[0], fine_grained_progress_objectives=[fgpo]) + result = fine_grained_progress_step_func(env, fine_grained_progress_objectives=[fgpo]) + states = env.extras["fine_grained_progress"]["states"] + assert not states[0]["fine_grained_progress_objectives"]["t"]["is_complete"] + assert states[0]["fine_grained_progress_objectives"]["t"]["score"] == 0.0 except Exception as e: print(f"Error: {e}") traceback.print_exc() @@ -574,12 +574,12 @@ def _test_step_func_publishes_to_extras_and_returns_no_termination(simulation_ap return True -def _test_task_base_fine_grained_subtask_hooks(simulation_app) -> bool: - """Test TaskBase's fine-grained-subtask hooks. Default is empty/None. Overriding - ``get_fine_grained_subtasks`` causes the events/termination helpers to +def _test_task_base_fine_grained_progress_objective_hooks(simulation_app) -> bool: + """Test TaskBase's fine-grained-progress-objective hooks. Default is empty/None. Overriding + ``get_fine_grained_progress_objectives`` causes the events/termination helpers to return real cfgs that the env builder picks up automatically. """ - from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask + from isaaclab_arena.tasks.fine_grained_progress_objective import FineGrainedProgressObjective from isaaclab_arena.tasks.task_base import TaskBase try: @@ -601,32 +601,32 @@ def get_metrics(self): return [] default_task = _Base() - assert default_task.get_fine_grained_subtasks() == [] - assert default_task.get_fine_grained_subtask_events_cfg() is None - assert default_task.get_fine_grained_subtask_termination_cfg() is None + assert default_task.get_fine_grained_progress_objectives() == [] + assert default_task.get_fine_grained_progress_objective_events_cfg() is None + assert default_task.get_fine_grained_progress_objective_termination_cfg() is None class _OptIn(_Base): - def get_fine_grained_subtasks(self): + def get_fine_grained_progress_objectives(self): pred = _MockPredicate(num_envs=1, name="p") - return [FineGrainedSubtask(name="lift", predicate_groups=pred)] + return [FineGrainedProgressObjective(name="lift", predicate_groups=pred)] opt_in = _OptIn() - assert len(opt_in.get_fine_grained_subtasks()) == 1 - assert opt_in.get_fine_grained_subtask_events_cfg() is not None - assert opt_in.get_fine_grained_subtask_termination_cfg() is not None + assert len(opt_in.get_fine_grained_progress_objectives()) == 1 + assert opt_in.get_fine_grained_progress_objective_events_cfg() is not None + assert opt_in.get_fine_grained_progress_objective_termination_cfg() is not None from isaaclab_arena.tasks.composite_task_base import CompositeTaskBase class _ChildA(_Base): - def get_fine_grained_subtasks(self): - return [FineGrainedSubtask(name="open", predicate_groups=_MockPredicate(1, name="pa"))] + def get_fine_grained_progress_objectives(self): + return [FineGrainedProgressObjective(name="open", predicate_groups=_MockPredicate(1, name="pa"))] class _ChildB(_Base): - def get_fine_grained_subtasks(self): - return [FineGrainedSubtask(name="close", predicate_groups=_MockPredicate(1, name="pb"))] + def get_fine_grained_progress_objectives(self): + return [FineGrainedProgressObjective(name="close", predicate_groups=_MockPredicate(1, name="pb"))] composite = CompositeTaskBase(subtasks=[_ChildA(), _ChildB()]) - recipes = composite.get_fine_grained_subtasks() + recipes = composite.get_fine_grained_progress_objectives() assert len(recipes) == 2 assert recipes[0].name == "subtask_0/open" assert recipes[0].parent_subtask_idx == 0 @@ -634,11 +634,11 @@ def get_fine_grained_subtasks(self): assert recipes[1].parent_subtask_idx == 1 class _CompositeWithOwn(CompositeTaskBase): - def get_own_fine_grained_subtasks(self): - return [FineGrainedSubtask(name="both_done", predicate_groups=_MockPredicate(1, name="own"))] + def get_own_fine_grained_progress_objectives(self): + return [FineGrainedProgressObjective(name="both_done", predicate_groups=_MockPredicate(1, name="own"))] composite2 = _CompositeWithOwn(subtasks=[_ChildA(), _ChildB()]) - recipes2 = composite2.get_fine_grained_subtasks() + recipes2 = composite2.get_fine_grained_progress_objectives() assert len(recipes2) == 3 assert recipes2[2].name == "both_done" assert recipes2[2].parent_subtask_idx is None @@ -715,8 +715,8 @@ def test_step_func_publishes_to_extras_and_returns_no_termination(): ) -def test_task_base_fine_grained_subtask_hooks(): - assert run_simulation_app_function(_test_task_base_fine_grained_subtask_hooks, headless=HEADLESS) +def test_task_base_fine_grained_progress_objective_hooks(): + assert run_simulation_app_function(_test_task_base_fine_grained_progress_objective_hooks, headless=HEADLESS) if __name__ == "__main__": @@ -736,4 +736,4 @@ def test_task_base_fine_grained_subtask_hooks(): test_gating_noop_when_env_has_no_current_subtask_idx() test_gating_sequential_task_end_to_end() test_step_func_publishes_to_extras_and_returns_no_termination() - test_task_base_fine_grained_subtask_hooks() + test_task_base_fine_grained_progress_objective_hooks() From de73e3dcf6c8213c4f9627c3bd4038037676ba57 Mon Sep 17 00:00:00 2001 From: Peter Du Date: Fri, 5 Jun 2026 11:05:39 -0700 Subject: [PATCH 08/10] move from step hook from terminations mgr to recorder mgr --- .../environments/arena_env_builder.py | 2 +- .../tasks/fine_grained_progress_tracker.py | 127 ++++++++++-------- isaaclab_arena/tasks/task_base.py | 6 +- ...ine_grained_progress_objective_tracking.py | 42 +++--- 4 files changed, 101 insertions(+), 76 deletions(-) diff --git a/isaaclab_arena/environments/arena_env_builder.py b/isaaclab_arena/environments/arena_env_builder.py index e64288928..47739712c 100644 --- a/isaaclab_arena/environments/arena_env_builder.py +++ b/isaaclab_arena/environments/arena_env_builder.py @@ -190,7 +190,6 @@ def compose_manager_cfg(self) -> IsaacLabArenaManagerBasedRLEnvCfg: task.get_termination_cfg(), self.arena_env.scene.get_termination_cfg(), embodiment.get_termination_cfg(), - task.get_fine_grained_progress_objective_termination_cfg(), ) actions_cfg = embodiment.get_action_cfg() xr_cfg = embodiment.get_xr_cfg() @@ -213,6 +212,7 @@ def compose_manager_cfg(self) -> IsaacLabArenaManagerBasedRLEnvCfg: metrics_recorder_manager_cfg, task.get_recorder_term_cfg(), embodiment.get_recorder_term_cfg(), + task.get_fine_grained_progress_objective_recorder_cfg(), bases=(RecorderManagerBaseCfg,), ) recorder_manager_cfg = self._modify_recorder_cfg_dataset_filename(recorder_manager_cfg) diff --git a/isaaclab_arena/tasks/fine_grained_progress_tracker.py b/isaaclab_arena/tasks/fine_grained_progress_tracker.py index 360a1c150..c41125142 100644 --- a/isaaclab_arena/tasks/fine_grained_progress_tracker.py +++ b/isaaclab_arena/tasks/fine_grained_progress_tracker.py @@ -9,7 +9,8 @@ from dataclasses import MISSING from typing import Any -from isaaclab.managers import EventTermCfg, TerminationTermCfg +from isaaclab.managers import EventTermCfg +from isaaclab.managers.recorder_manager import RecorderManagerBaseCfg, RecorderTerm, RecorderTermCfg from isaaclab.utils import configclass from isaaclab_arena.tasks.fine_grained_progress_objective import FineGrainedProgressObjective @@ -193,7 +194,9 @@ def __init__(self, fine_grained_progress_objectives: list[FineGrainedProgressObj self.fine_grained_progress_objectives = fine_grained_progress_objectives self.num_envs = num_envs self.device = device - self.runners = [FineGrainedProgressObjectiveRunner(s, num_envs, device) for s in fine_grained_progress_objectives] + self.runners = [ + FineGrainedProgressObjectiveRunner(s, num_envs, device) for s in fine_grained_progress_objectives + ] self._events: list[list[dict]] = [[] for _ in range(num_envs)] def step(self, env, step_index: torch.Tensor | None) -> None: @@ -279,57 +282,64 @@ def _ensure_progress_tracker( return sm -def fine_grained_progress_step_func(env, fine_grained_progress_objectives: list[FineGrainedProgressObjective]) -> torch.Tensor: - """Termination-term entry point. - - Ticks the state machine, writes events and states to env.extras["fine_grained_progress"], - and returns all-False so it does not contribute to termination. - """ - - sm = _ensure_progress_tracker(env, fine_grained_progress_objectives) - step_index = getattr(env, "episode_length_buf", None) - sm.step(env, step_index=step_index) - - """ - User-facing event/state information format: - - env.extras["fine_grained_progress"] = { - "states": [ - { - "fine_grained_progress_objectives": { - "": { - "completed_groups": int, - "total_groups": int, - "score": float, # 0..1, normalized within objective - "is_complete": bool, - "active_predicates": {group: str | None}, +class FineGrainedProgressRecorder(RecorderTerm): + """Per-step hook that ticks the FineGrainedProgressTracker. Records nothing. + + Registered as a recorder term so it runs once per env.step via + record_post_step. It advances the state nachine and publishes the per-step state/events to + env.extras["fine_grained_progress"], then returns + (None, None) so nothing is written to the recorded episode data. + + env.extras["fine_grained_progress"] format: + + { + "states": [ # one entry per env + { + "fine_grained_progress_objectives": { + "": { + "completed_groups": int, + "total_groups": int, + "score": float, # 0..1, normalized within objective + "is_complete": bool, + "active_predicates": {group: str | None}, + }, + ... }, - ... + "overall_score": float, # weighted by FineGrainedProgressObjective.score + "all_complete": bool, }, - "overall_score": float, # weighted by FineGrainedProgressObjective.score - "all_complete": bool, - }, - ... # one entry per env - ], - "events": [ - [{"step": int, "fine_grained_progress_objective": str, "group": str, - "predicate_index": int, "predicate_name": str, - "score_delta": float}, ...], - ... # one list per env - ], - } + ... + ], + "events": [ # one list per env + [{"step": int, "fine_grained_progress_objective": str, "group": str, + "predicate_index": int, "predicate_name": str, + "score_delta": float}, ...], + ... + ], + } """ - env.extras["fine_grained_progress"] = { - "states": sm.get_state(), - "events": sm.get_events(), - } + def __init__(self, cfg: FineGrainedProgressObjectiveRecorderCfg, env): + super().__init__(cfg, env) + self._fine_grained_progress_objectives = cfg.fine_grained_progress_objectives - # Return all-False so it does not contribute to termination. - return torch.zeros(env.num_envs, dtype=torch.bool, device=env.device) + def record_post_step(self): + """Ticks the state machine, writes events and states to env.extras["fine_grained_progress"]""" + sm = _ensure_progress_tracker(self._env, self._fine_grained_progress_objectives) + step_index = getattr(self._env, "episode_length_buf", None) + sm.step(self._env, step_index=step_index) + self._env.extras["fine_grained_progress"] = { + "states": sm.get_state(), + "events": sm.get_events(), + } + # This term is a per-step hook only — record nothing. + return None, None -def fine_grained_progress_reset_func(env, env_ids, fine_grained_progress_objectives: list[FineGrainedProgressObjective]) -> None: + +def fine_grained_progress_reset_func( + env, env_ids, fine_grained_progress_objectives: list[FineGrainedProgressObjective] +) -> None: """Reset-event entry point. Resets the state machine whenever the Lab env is reset. @@ -349,11 +359,19 @@ class FineGrainedProgressObjectiveEventsCfg: @configclass -class FineGrainedProgressObjectiveTerminationsCfg: - fine_grained_progress_step: TerminationTermCfg = MISSING +class FineGrainedProgressObjectiveRecorderCfg(RecorderTermCfg): + class_type: type[RecorderTerm] = FineGrainedProgressRecorder + fine_grained_progress_objectives: list[FineGrainedProgressObjective] = MISSING + +@configclass +class FineGrainedProgressObjectiveRecorderManagerCfg(RecorderManagerBaseCfg): + fine_grained_progress: FineGrainedProgressObjectiveRecorderCfg = MISSING -def make_fine_grained_progress_objective_events_cfg(fine_grained_progress_objectives: list[FineGrainedProgressObjective]) -> Any: + +def make_fine_grained_progress_objective_events_cfg( + fine_grained_progress_objectives: list[FineGrainedProgressObjective], +) -> Any: return FineGrainedProgressObjectiveEventsCfg( reset_fine_grained_progress_objectives=EventTermCfg( func=fine_grained_progress_reset_func, @@ -363,10 +381,11 @@ def make_fine_grained_progress_objective_events_cfg(fine_grained_progress_object ) -def make_fine_grained_progress_objective_termination_cfg(fine_grained_progress_objectives: list[FineGrainedProgressObjective]) -> Any: - return FineGrainedProgressObjectiveTerminationsCfg( - fine_grained_progress_step=TerminationTermCfg( - func=fine_grained_progress_step_func, - params={"fine_grained_progress_objectives": fine_grained_progress_objectives}, +def make_fine_grained_progress_objective_recorder_cfg( + fine_grained_progress_objectives: list[FineGrainedProgressObjective], +) -> Any: + return FineGrainedProgressObjectiveRecorderManagerCfg( + fine_grained_progress=FineGrainedProgressObjectiveRecorderCfg( + fine_grained_progress_objectives=fine_grained_progress_objectives, ) ) diff --git a/isaaclab_arena/tasks/task_base.py b/isaaclab_arena/tasks/task_base.py index e2f3bb606..62bc88fd2 100644 --- a/isaaclab_arena/tasks/task_base.py +++ b/isaaclab_arena/tasks/task_base.py @@ -14,7 +14,7 @@ from isaaclab_arena.tasks.fine_grained_progress_objective import FineGrainedProgressObjective from isaaclab_arena.tasks.fine_grained_progress_tracker import ( make_fine_grained_progress_objective_events_cfg, - make_fine_grained_progress_objective_termination_cfg, + make_fine_grained_progress_objective_recorder_cfg, ) @@ -77,8 +77,8 @@ def get_fine_grained_progress_objective_events_cfg(self) -> Any: return None return make_fine_grained_progress_objective_events_cfg(fine_grained_progress_objectives) - def get_fine_grained_progress_objective_termination_cfg(self) -> Any: + def get_fine_grained_progress_objective_recorder_cfg(self) -> Any: fine_grained_progress_objectives = self.get_fine_grained_progress_objectives() if not fine_grained_progress_objectives: return None - return make_fine_grained_progress_objective_termination_cfg(fine_grained_progress_objectives) + return make_fine_grained_progress_objective_recorder_cfg(fine_grained_progress_objectives) diff --git a/isaaclab_arena/tests/test_fine_grained_progress_objective_tracking.py b/isaaclab_arena/tests/test_fine_grained_progress_objective_tracking.py index 2f3df659b..d1d16799d 100644 --- a/isaaclab_arena/tests/test_fine_grained_progress_objective_tracking.py +++ b/isaaclab_arena/tests/test_fine_grained_progress_objective_tracking.py @@ -525,12 +525,17 @@ def _test_gating_noop_when_env_has_no_current_subtask_idx(simulation_app) -> boo return True -def _test_step_func_publishes_to_extras_and_returns_no_termination(simulation_app) -> bool: - """fine_grained_progress_step_func writes env.extras and returns all-False.""" +def _test_recorder_publishes_to_extras_and_records_nothing(simulation_app) -> bool: + """FineGrainedProgressRecorder.record_post_step writes env.extras and records nothing. + + ``record_post_step`` returns ``(None, None)`` (so nothing is added to the recorded + episode data) while still ticking the tracker and publishing the per-step state to + ``env.extras["fine_grained_progress"]``. + """ from isaaclab_arena.tasks.fine_grained_progress_objective import FineGrainedProgressObjective from isaaclab_arena.tasks.fine_grained_progress_tracker import ( + FineGrainedProgressObjectiveRecorderCfg, fine_grained_progress_reset_func, - fine_grained_progress_step_func, ) try: @@ -538,21 +543,24 @@ def _test_step_func_publishes_to_extras_and_returns_no_termination(simulation_ap pred = _MockPredicate(num_envs=2, name="p") fgpo = FineGrainedProgressObjective(name="t", predicate_groups=pred) + recorder_cfg = FineGrainedProgressObjectiveRecorderCfg(fine_grained_progress_objectives=[fgpo]) + recorder = recorder_cfg.class_type(recorder_cfg, env) + fine_grained_progress_reset_func(env, env_ids=[0, 1], fine_grained_progress_objectives=[fgpo]) - # Step with predicate=False, state machine ticks but no transitions. - result = fine_grained_progress_step_func(env, fine_grained_progress_objectives=[fgpo]) - assert result.tolist() == [False, False] + # Step with predicate=False, state machine ticks but no transitions. Records nothing. + assert recorder.record_post_step() == (None, None) assert "fine_grained_progress" in env.extras assert len(env.extras["fine_grained_progress"]["states"]) == 2 assert env.extras["fine_grained_progress"]["events"] == [[], []] - assert not env.extras["fine_grained_progress"]["states"][0]["fine_grained_progress_objectives"]["t"]["is_complete"] + assert not env.extras["fine_grained_progress"]["states"][0]["fine_grained_progress_objectives"]["t"][ + "is_complete" + ] # Step with env 0 predicate True, env 0 completes, env 1 does not. pred.set([True, False]) _advance_step(env) - result = fine_grained_progress_step_func(env, fine_grained_progress_objectives=[fgpo]) - assert result.tolist() == [False, False] + assert recorder.record_post_step() == (None, None) states = env.extras["fine_grained_progress"]["states"] events = env.extras["fine_grained_progress"]["events"] assert states[0]["fine_grained_progress_objectives"]["t"]["is_complete"] @@ -563,7 +571,7 @@ def _test_step_func_publishes_to_extras_and_returns_no_termination(simulation_ap # Reset env 0, env 1 untouched. pred.set([False, False]) fine_grained_progress_reset_func(env, env_ids=[0], fine_grained_progress_objectives=[fgpo]) - result = fine_grained_progress_step_func(env, fine_grained_progress_objectives=[fgpo]) + assert recorder.record_post_step() == (None, None) states = env.extras["fine_grained_progress"]["states"] assert not states[0]["fine_grained_progress_objectives"]["t"]["is_complete"] assert states[0]["fine_grained_progress_objectives"]["t"]["score"] == 0.0 @@ -576,7 +584,7 @@ def _test_step_func_publishes_to_extras_and_returns_no_termination(simulation_ap def _test_task_base_fine_grained_progress_objective_hooks(simulation_app) -> bool: """Test TaskBase's fine-grained-progress-objective hooks. Default is empty/None. Overriding - ``get_fine_grained_progress_objectives`` causes the events/termination helpers to + ``get_fine_grained_progress_objectives`` causes the events/recorder helpers to return real cfgs that the env builder picks up automatically. """ from isaaclab_arena.tasks.fine_grained_progress_objective import FineGrainedProgressObjective @@ -603,7 +611,7 @@ def get_metrics(self): default_task = _Base() assert default_task.get_fine_grained_progress_objectives() == [] assert default_task.get_fine_grained_progress_objective_events_cfg() is None - assert default_task.get_fine_grained_progress_objective_termination_cfg() is None + assert default_task.get_fine_grained_progress_objective_recorder_cfg() is None class _OptIn(_Base): def get_fine_grained_progress_objectives(self): @@ -613,7 +621,7 @@ def get_fine_grained_progress_objectives(self): opt_in = _OptIn() assert len(opt_in.get_fine_grained_progress_objectives()) == 1 assert opt_in.get_fine_grained_progress_objective_events_cfg() is not None - assert opt_in.get_fine_grained_progress_objective_termination_cfg() is not None + assert opt_in.get_fine_grained_progress_objective_recorder_cfg() is not None from isaaclab_arena.tasks.composite_task_base import CompositeTaskBase @@ -709,10 +717,8 @@ def test_gating_sequential_task_end_to_end(): assert run_simulation_app_function(_test_gating_sequential_task_end_to_end, headless=HEADLESS) -def test_step_func_publishes_to_extras_and_returns_no_termination(): - assert run_simulation_app_function( - _test_step_func_publishes_to_extras_and_returns_no_termination, headless=HEADLESS - ) +def test_recorder_publishes_to_extras_and_records_nothing(): + assert run_simulation_app_function(_test_recorder_publishes_to_extras_and_records_nothing, headless=HEADLESS) def test_task_base_fine_grained_progress_objective_hooks(): @@ -735,5 +741,5 @@ def test_task_base_fine_grained_progress_objective_hooks(): test_gating_blocked_when_parent_subtask_idx_mismatches() test_gating_noop_when_env_has_no_current_subtask_idx() test_gating_sequential_task_end_to_end() - test_step_func_publishes_to_extras_and_returns_no_termination() + test_recorder_publishes_to_extras_and_records_nothing() test_task_base_fine_grained_progress_objective_hooks() From 680fa6d4643fe21a432518665ea4818293e8e220 Mon Sep 17 00:00:00 2001 From: Peter Du Date: Fri, 5 Jun 2026 11:19:08 -0700 Subject: [PATCH 09/10] address comments for readbility in state machine --- .../tasks/fine_grained_progress_objective.py | 10 +- .../tasks/fine_grained_progress_tracker.py | 208 ++++++++++-------- 2 files changed, 125 insertions(+), 93 deletions(-) diff --git a/isaaclab_arena/tasks/fine_grained_progress_objective.py b/isaaclab_arena/tasks/fine_grained_progress_objective.py index b0cca726b..7b432f33c 100644 --- a/isaaclab_arena/tasks/fine_grained_progress_objective.py +++ b/isaaclab_arena/tasks/fine_grained_progress_objective.py @@ -53,7 +53,8 @@ def format_predicate_groups(predicate_groups: PredicateGroups) -> dict[str, list } raise TypeError( - f"FineGrainedProgressObjective.predicate_groups must be a callable, list, or dict; got {type(predicate_groups).__name__}" + "FineGrainedProgressObjective.predicate_groups must be a callable, list, or dict; got" + f" {type(predicate_groups).__name__}" ) @@ -147,7 +148,8 @@ def __post_init__(self): raise ValueError(f"FineGrainedProgressObjective '{self.name}': score must be in [0, 1], got {self.score}") if self.logical not in ("all", "any", "choose"): raise ValueError( - f"FineGrainedProgressObjective '{self.name}': logical must be in ['all', 'any', 'choose'], got {self.logical}" + f"FineGrainedProgressObjective '{self.name}': logical must be in ['all', 'any', 'choose'], got" + f" {self.logical}" ) # Format the predicate groups into the canonical form and normalize the scores. @@ -161,7 +163,9 @@ def __post_init__(self): if self.K is None: raise ValueError(f"FineGrainedProgressObjective '{self.name}': K is required when logical='choose'") if not (1 <= self.K <= num_groups): - raise ValueError(f"FineGrainedProgressObjective '{self.name}': K={self.K} but must be in [1, {num_groups}]") + raise ValueError( + f"FineGrainedProgressObjective '{self.name}': K={self.K} but must be in [1, {num_groups}]" + ) @property def group_names(self) -> list[str]: diff --git a/isaaclab_arena/tasks/fine_grained_progress_tracker.py b/isaaclab_arena/tasks/fine_grained_progress_tracker.py index c41125142..08677ca84 100644 --- a/isaaclab_arena/tasks/fine_grained_progress_tracker.py +++ b/isaaclab_arena/tasks/fine_grained_progress_tracker.py @@ -81,74 +81,91 @@ def _compute_composite_task_gating_mask(self, env) -> torch.Tensor: def step(self, env, step_index: torch.Tensor | None) -> list[dict]: """Step the state machine runner for a single env.step. - Check each group's current predicate, move the state machine to the next predicate if - the current predicate is True. Emit an event for each env where a predicate was advanced. + Advance each group's predicate chain by at most one position per env and return a + transition event for every env/group that advanced this step. """ - # List of state transition events (events are emitted for an env when a predicate flips True) + # If the FineGrainedProgressObjective is not active for the composite task, there is + # nothing to advance for any env. + gating_mask = self._compute_composite_task_gating_mask(env) + if not bool(gating_mask.any().item()): + return [] + events: list[dict] = [] + for group_name, predicate_chain in self.fine_grained_progress_objective.canonical_predicate_groups.items(): + events += self._step_group(env, group_name, predicate_chain, gating_mask, step_index) + return events - # If the FineGrainedProgressObjective is not active for the composite task, return. - composite_task_gating_mask = self._compute_composite_task_gating_mask(env) - if not bool(composite_task_gating_mask.any().item()): - return events + def _step_group( + self, + env, + group_name: str, + predicate_chain: list[tuple], + gating_mask: torch.Tensor, + step_index: torch.Tensor | None, + ) -> list[dict]: + """Advance a single group's predicate chain by at most one position per env. + + Evaluates the current predicate for the envs sitting at each chain position, advances + those whose predicate is satisfied, updates the group's score and completion mask, and + returns one transition event per env that advanced. + """ - # Step through each group of the FineGrainedProgressObjective. - for group_name, predicate_chain in self.fine_grained_progress_objective.canonical_predicate_groups.items(): - chain_length = len(predicate_chain) - # Mask for which envs have advanced this step. - advanced = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device) - - for chain_idx, (predicate, score_weight) in enumerate(predicate_chain): - # Compute mask for which envs that should evaluate the predicate. - # Envs should only be evaluated if: - # 1) They are at the current predicate position - # 2) They have not yet advanced this step - # 3) The FineGrainedProgressObjective is active for the composite task - at_position = (self.current_index[group_name] == chain_idx) & ~advanced & composite_task_gating_mask - if not bool(at_position.any().item()): - continue - - # Evaluate the predicate for all envs, reshaped to a flat (num_envs,) bool tensor. - result = torch.as_tensor(predicate(env), dtype=torch.bool, device=self.device).reshape(-1) - if result.shape[0] != self.num_envs: - raise RuntimeError( - f"Predicate {_predicate_repr(predicate)} returned shape {tuple(result.shape)};" - f" expected ({self.num_envs},)" - ) - - # Compute mask for which envs need to be advanced to the next predicate. - advance_mask = at_position & result - if not bool(advance_mask.any().item()): - continue - - # Advance the state machine to the next predicates. - self.current_index[group_name] = torch.where( - advance_mask, - self.current_index[group_name] + 1, - self.current_index[group_name], + # List of state transition events (events are emitted for an env when a predicate flips True) + events: list[dict] = [] + chain_length = len(predicate_chain) + # Mask for which envs have advanced this step (at most one advance per env per group). + advanced = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device) + + for chain_idx, (predicate, score_weight) in enumerate(predicate_chain): + # Compute mask for which envs that should evaluate the predicate. + # Envs should only be evaluated if: + # 1) They are at the current predicate position + # 2) They have not yet advanced this step + # 3) The FineGrainedProgressObjective is active for the composite task + at_position = (self.current_index[group_name] == chain_idx) & ~advanced & gating_mask + if not bool(at_position.any().item()): + continue + + # Evaluate the predicate for all envs, reshaped to a flat (num_envs,) bool tensor. + result = torch.as_tensor(predicate(env), dtype=torch.bool, device=self.device).reshape(-1) + if result.shape[0] != self.num_envs: + raise RuntimeError( + f"Predicate {_predicate_repr(predicate)} returned shape {tuple(result.shape)};" + f" expected ({self.num_envs},)" ) - # Update the group score for the envs that were advanced. - self.group_score[group_name] = self.group_score[group_name] + advance_mask.float() * float(score_weight) - # Update the advanced mask for the envs that were advanced. - advanced = advanced | advance_mask - - # Emit an event for each env where a predicate was advanced. - pred_name = _predicate_repr(predicate) - for eid in torch.nonzero(advance_mask, as_tuple=False).flatten().tolist(): - events.append({ - "env_idx": int(eid), - "step": int(step_index[eid].item()) if step_index is not None else -1, - "fine_grained_progress_objective": self.fine_grained_progress_objective.name, - "group": group_name, - "predicate_index": chain_idx, - "predicate_name": pred_name, - "score_delta": float(score_weight), - }) - - # Update the group complete mask for the envs that have completed the group. - self.group_complete[group_name] = self.current_index[group_name] >= chain_length + # Compute mask for which envs need to be advanced to the next predicate. + advance_mask = at_position & result + if not bool(advance_mask.any().item()): + continue + + # Advance the state machine to the next predicates. + self.current_index[group_name] = torch.where( + advance_mask, + self.current_index[group_name] + 1, + self.current_index[group_name], + ) + # Update the group score for the envs that were advanced. + self.group_score[group_name] = self.group_score[group_name] + advance_mask.float() * float(score_weight) + # Update the advanced mask for the envs that were advanced. + advanced = advanced | advance_mask + + # Emit an event for each env where a predicate was advanced. + pred_name = _predicate_repr(predicate) + for eid in torch.nonzero(advance_mask, as_tuple=False).flatten().tolist(): + events.append({ + "env_idx": int(eid), + "step": int(step_index[eid].item()) if step_index is not None else -1, + "fine_grained_progress_objective": self.fine_grained_progress_objective.name, + "group": group_name, + "predicate_index": chain_idx, + "predicate_name": pred_name, + "score_delta": float(score_weight), + }) + + # Update the group complete mask for the envs that have completed the group. + self.group_complete[group_name] = self.current_index[group_name] >= chain_length return events def reset(self, env_ids) -> None: @@ -178,6 +195,36 @@ def overall_score_per_env(self) -> torch.Tensor: stacked = torch.stack([self.group_score[g] for g in groups], dim=1) return stacked.mean(dim=1) + def get_state_for_env(self, env_idx: int, is_complete, score) -> dict: + """Per-env view of this objective's progress. + + is_complete and score are passed in (rather than recomputed here) so the full + (num_envs,) tensor reductions run once per runner in + FineGrainedProgressTracker, instead of once per env. + """ + + objective = self.fine_grained_progress_objective + completed_groups = 0 + active_predicates: dict[str, str | None] = {} + # The active predicate for a group is the one at its current chain position. Any group + # whose pointer has run off the end of the chain is complete (no active predicate). + for group_name in objective.group_names: + predicate_chain = objective.canonical_predicate_groups[group_name] + cur_group_index = int(self.current_index[group_name][env_idx].item()) + if cur_group_index >= len(predicate_chain): + active_predicates[group_name] = None + completed_groups += 1 + else: + active_predicates[group_name] = _predicate_repr(predicate_chain[cur_group_index][0]) + + return { + "completed_groups": completed_groups, + "total_groups": len(objective.group_names), + "score": float(score), + "is_complete": bool(is_complete), + "active_predicates": active_predicates, + } + class FineGrainedProgressTracker: """State machine that manages runners for all FineGrainedProgressObjectives. @@ -218,41 +265,22 @@ def reset(self, env_ids) -> None: def get_state(self) -> list[dict]: """Get the state of each FineGrainedProgressObjective for all envs.""" + # Compute the per-runner (num_envs,) tensors once + completeness = [runner.is_complete() for runner in self.runners] + scores = [runner.overall_score_per_env() for runner in self.runners] + output: list[dict] = [] for env_idx in range(self.num_envs): # Build a per-env dict from each runner's state. progress_objective_states: dict[str, dict] = {} overall_score = 0.0 all_complete = True - - for runner in self.runners: - fine_grained_progress_objective = runner.fine_grained_progress_objective - completed_groups = 0 - total_groups = len(fine_grained_progress_objective.group_names) - active_predicates: dict[str, str | None] = {} - - # Compute the active predicates and completed groups. - for group_name in fine_grained_progress_objective.group_names: - cur_group_index = int(runner.current_index[group_name][env_idx].item()) - predicate_chain = fine_grained_progress_objective.canonical_predicate_groups[group_name] - if cur_group_index >= len(predicate_chain): - active_predicates[group_name] = None - completed_groups += 1 - else: - active_predicates[group_name] = _predicate_repr(predicate_chain[cur_group_index][0]) - - # Compute the overall score and completeness. - progress_objective_score = float(runner.overall_score_per_env()[env_idx].item()) - is_complete = bool(runner.is_complete()[env_idx].item()) - progress_objective_states[fine_grained_progress_objective.name] = { - "completed_groups": completed_groups, - "total_groups": total_groups, - "score": progress_objective_score, - "is_complete": is_complete, - "active_predicates": active_predicates, - } - overall_score += fine_grained_progress_objective.score * progress_objective_score - all_complete = all_complete and is_complete + for i, runner in enumerate(self.runners): + objective = runner.fine_grained_progress_objective + state = runner.get_state_for_env(env_idx, completeness[i][env_idx], scores[i][env_idx]) + progress_objective_states[objective.name] = state + overall_score += objective.score * state["score"] + all_complete = all_complete and state["is_complete"] # Add the per-env state dict to the output. output.append({ From 24807a493ffbb447863619145cb728abe05f965f Mon Sep 17 00:00:00 2001 From: Peter Du Date: Fri, 5 Jun 2026 12:19:48 -0700 Subject: [PATCH 10/10] perf improvement in reset and resolve object lists --- isaaclab_arena/tasks/fine_grained_progress_tracker.py | 10 +++++----- isaaclab_arena/tasks/task_base.py | 10 ++++++++-- 2 files changed, 13 insertions(+), 7 deletions(-) diff --git a/isaaclab_arena/tasks/fine_grained_progress_tracker.py b/isaaclab_arena/tasks/fine_grained_progress_tracker.py index 08677ca84..5d30e91aa 100644 --- a/isaaclab_arena/tasks/fine_grained_progress_tracker.py +++ b/isaaclab_arena/tasks/fine_grained_progress_tracker.py @@ -171,11 +171,11 @@ def _step_group( def reset(self, env_ids) -> None: """Reset the state machine runner for the provided envs.""" + env_ids = torch.as_tensor(env_ids, dtype=torch.long, device=self.device) for group_name in self.fine_grained_progress_objective.group_names: - for eid in env_ids: - self.current_index[group_name][eid] = 0 - self.group_score[group_name][eid] = 0.0 - self.group_complete[group_name][eid] = False + self.current_index[group_name][env_ids] = 0 + self.group_score[group_name][env_ids] = 0.0 + self.group_complete[group_name][env_ids] = False def is_complete(self) -> torch.Tensor: """Check if the FineGrainedProgressObjective is complete for all envs.""" @@ -314,7 +314,7 @@ class FineGrainedProgressRecorder(RecorderTerm): """Per-step hook that ticks the FineGrainedProgressTracker. Records nothing. Registered as a recorder term so it runs once per env.step via - record_post_step. It advances the state nachine and publishes the per-step state/events to + record_post_step. It advances the state machine and publishes the per-step state/events to env.extras["fine_grained_progress"], then returns (None, None) so nothing is written to the recorded episode data. diff --git a/isaaclab_arena/tasks/task_base.py b/isaaclab_arena/tasks/task_base.py index 62bc88fd2..6189a6f03 100644 --- a/isaaclab_arena/tasks/task_base.py +++ b/isaaclab_arena/tasks/task_base.py @@ -71,14 +71,20 @@ def get_task_description(self) -> str | None: def get_fine_grained_progress_objectives(self) -> list[FineGrainedProgressObjective]: return [] + def _resolve_fine_grained_progress_objectives(self) -> list[FineGrainedProgressObjective]: + # Resolve once and cache so the reset (events) and step (recorder) cfgs share the same objective objects. + if not hasattr(self, "_resolved_fine_grained_progress_objectives"): + self._resolved_fine_grained_progress_objectives = self.get_fine_grained_progress_objectives() + return self._resolved_fine_grained_progress_objectives + def get_fine_grained_progress_objective_events_cfg(self) -> Any: - fine_grained_progress_objectives = self.get_fine_grained_progress_objectives() + fine_grained_progress_objectives = self._resolve_fine_grained_progress_objectives() if not fine_grained_progress_objectives: return None return make_fine_grained_progress_objective_events_cfg(fine_grained_progress_objectives) def get_fine_grained_progress_objective_recorder_cfg(self) -> Any: - fine_grained_progress_objectives = self.get_fine_grained_progress_objectives() + fine_grained_progress_objectives = self._resolve_fine_grained_progress_objectives() if not fine_grained_progress_objectives: return None return make_fine_grained_progress_objective_recorder_cfg(fine_grained_progress_objectives)