xlab-api 0.1.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
xlab_api/__init__.py ADDED
@@ -0,0 +1,66 @@
1
+ """Public engine-neutral XLab schemas and type declarations."""
2
+
3
+ from __future__ import annotations
4
+
5
+ from typing import Any
6
+
7
+ from . import spec
8
+
9
+ API_VERSION = "0.1"
10
+
11
+
12
+ def __getattr__(name: str) -> Any:
13
+ try:
14
+ value = getattr(spec, name)
15
+ except AttributeError:
16
+ raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from None
17
+ globals()[name] = value
18
+ return value
19
+
20
+
21
+ __all__ = [
22
+ "API_VERSION",
23
+ "UNIVERSAL_KINDS",
24
+ "ActionTermSpec",
25
+ "AgentSpec",
26
+ "ArticulationRef",
27
+ "BackendAsset",
28
+ "ClockDomainSpec",
29
+ "CollisionExclusionRef",
30
+ "CommandTermSpec",
31
+ "ComponentLifecycleSpec",
32
+ "ContactSensorRef",
33
+ "CurriculumTermSpec",
34
+ "DoneTermSpec",
35
+ "EntityRef",
36
+ "EventTermSpec",
37
+ "Grid3dPointsRef",
38
+ "JointProperties",
39
+ "LifecycleSpec",
40
+ "MdpSpec",
41
+ "MotionReferenceRef",
42
+ "NativeSensorRef",
43
+ "NoiseSpec",
44
+ "ObsGroupSpec",
45
+ "ObsTermSpec",
46
+ "RayCasterRef",
47
+ "RayPatternRef",
48
+ "Requirement",
49
+ "ResolvedClockDomain",
50
+ "RewardTermSpec",
51
+ "RigidObjectRef",
52
+ "RobotSpec",
53
+ "SceneSpec",
54
+ "SimSpec",
55
+ "SubTerrainSpec",
56
+ "SymmetricAugmentationSpec",
57
+ "TaskSpec",
58
+ "TermSpec",
59
+ "TerrainGeneratorSpec",
60
+ "TerrainSpec",
61
+ "VirtualObstacleRef",
62
+ "VolumePointsRef",
63
+ "freeze_task_spec",
64
+ "portability_report",
65
+ "resolve_lifecycle_contract",
66
+ ]
xlab_api/data.py ADDED
@@ -0,0 +1,90 @@
1
+ """Portable dataset URIs and local data-root resolution."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import os
6
+ from pathlib import Path, PurePosixPath
7
+ from urllib.parse import unquote, urlsplit
8
+
9
+ DATA_ROOT_ENV = "XLAB_DATA_ROOT"
10
+ DATASET_SCHEME = "dataset"
11
+
12
+
13
+ def is_dataset_uri(value: str | os.PathLike[str]) -> bool:
14
+ """Return whether ``value`` uses XLab's portable dataset URI scheme."""
15
+ return os.fspath(value).startswith(f"{DATASET_SCHEME}://")
16
+
17
+
18
+ def dataset_root() -> Path:
19
+ """Return the configured local root for portable dataset URIs."""
20
+ declared = os.environ.get(DATA_ROOT_ENV, "~/Datasets")
21
+ return Path(declared).expanduser().resolve()
22
+
23
+
24
+ def _dataset_relative_path(value: str) -> Path:
25
+ parsed = urlsplit(value)
26
+ if parsed.scheme != DATASET_SCHEME or not parsed.netloc:
27
+ raise ValueError(
28
+ f"Dataset URI must have the form dataset://collection/path, got {value!r}."
29
+ )
30
+ if (
31
+ parsed.query
32
+ or parsed.fragment
33
+ or parsed.username
34
+ or parsed.password
35
+ or parsed.port
36
+ ):
37
+ raise ValueError(f"Dataset URI contains unsupported URL fields: {value!r}.")
38
+ authority = unquote(parsed.netloc)
39
+ decoded_path = unquote(parsed.path)
40
+ if any(separator in authority for separator in ("/", "\\")) or "\\" in decoded_path:
41
+ raise ValueError(
42
+ f"Dataset URI contains an encoded or non-portable path separator: {value!r}."
43
+ )
44
+ components = (authority, *PurePosixPath(decoded_path).parts)
45
+ relative_parts = tuple(part for part in components if part not in {"", "/"})
46
+ if not relative_parts or any(part in {".", ".."} for part in relative_parts):
47
+ raise ValueError(
48
+ f"Dataset URI must not contain traversal components: {value!r}."
49
+ )
50
+ return Path(*relative_parts)
51
+
52
+
53
+ def resolve_data_path(
54
+ value: str | os.PathLike[str],
55
+ *,
56
+ relative_to: str | os.PathLike[str] | None = None,
57
+ ) -> Path:
58
+ """Resolve a portable dataset URI or a readable legacy filesystem path.
59
+
60
+ ``dataset://collection/path`` is rooted at :envvar:`XLAB_DATA_ROOT`,
61
+ which defaults to ``~/Datasets``. Plain paths retain their historical
62
+ ``~`` behavior. Relative plain paths can be anchored with ``relative_to``.
63
+ Resolution does not require the target to exist so preflight can report
64
+ missing optional resources without changing the declaration.
65
+ """
66
+ declared = os.fspath(value)
67
+ if is_dataset_uri(declared):
68
+ root = dataset_root()
69
+ resolved = (root / _dataset_relative_path(declared)).resolve()
70
+ try:
71
+ resolved.relative_to(root)
72
+ except ValueError as exc:
73
+ raise ValueError(
74
+ f"Dataset URI resolves outside {DATA_ROOT_ENV}: {declared!r}."
75
+ ) from exc
76
+ return resolved
77
+
78
+ path = Path(declared).expanduser()
79
+ if relative_to is not None and not path.is_absolute():
80
+ path = Path(relative_to).expanduser() / path
81
+ return path.resolve()
82
+
83
+
84
+ __all__ = [
85
+ "DATASET_SCHEME",
86
+ "DATA_ROOT_ENV",
87
+ "dataset_root",
88
+ "is_dataset_uri",
89
+ "resolve_data_path",
90
+ ]
xlab_api/py.typed ADDED
@@ -0,0 +1 @@
1
+
@@ -0,0 +1,118 @@
1
+ """Engine-agnostic declaration layer.
2
+
3
+ Nothing here may import a physics engine, directly or transitively. A task is declared once in
4
+ these types and compiled to a native environment by the backend for whichever engine is running,
5
+ so anything that reaches for an engine at this level has already lost the property that makes the
6
+ declaration portable. Engine isolation is an explicit package boundary.
7
+
8
+ :class:`~xlab_api.spec.task.TaskSpec` is the entry point: a whole task, produced by a frontend
9
+ that read some project's native definition and consumed by the backend for the engine in use.
10
+ """
11
+
12
+ from __future__ import annotations
13
+
14
+ from importlib import import_module
15
+ from typing import Any
16
+
17
+ _EXPORTS = {
18
+ "ArticulationRef": ("articulation", "ArticulationRef"),
19
+ "Requirement": ("capability", "Requirement"),
20
+ "UNIVERSAL_KINDS": ("entity", "UNIVERSAL_KINDS"),
21
+ "EntityRef": ("entity", "EntityRef"),
22
+ "freeze_task_spec": ("freeze", "freeze_task_spec"),
23
+ "ClockDomainSpec": ("lifecycle", "ClockDomainSpec"),
24
+ "ComponentLifecycleSpec": ("lifecycle", "ComponentLifecycleSpec"),
25
+ "LifecycleSpec": ("lifecycle", "LifecycleSpec"),
26
+ "ResolvedClockDomain": ("lifecycle", "ResolvedClockDomain"),
27
+ "resolve_lifecycle_contract": ("lifecycle", "resolve_lifecycle_contract"),
28
+ "ActionTermSpec": ("mdp", "ActionTermSpec"),
29
+ "CommandTermSpec": ("mdp", "CommandTermSpec"),
30
+ "CurriculumTermSpec": ("mdp", "CurriculumTermSpec"),
31
+ "DoneTermSpec": ("mdp", "DoneTermSpec"),
32
+ "EventTermSpec": ("mdp", "EventTermSpec"),
33
+ "MdpSpec": ("mdp", "MdpSpec"),
34
+ "NoiseSpec": ("mdp", "NoiseSpec"),
35
+ "ObsGroupSpec": ("mdp", "ObsGroupSpec"),
36
+ "ObsTermSpec": ("mdp", "ObsTermSpec"),
37
+ "RewardTermSpec": ("mdp", "RewardTermSpec"),
38
+ "TermSpec": ("mdp", "TermSpec"),
39
+ "portability_report": ("portability", "portability_report"),
40
+ "RigidObjectRef": ("rigid_object", "RigidObjectRef"),
41
+ "CollisionExclusionRef": ("relations", "CollisionExclusionRef"),
42
+ "BackendAsset": ("robot", "BackendAsset"),
43
+ "JointProperties": ("robot", "JointProperties"),
44
+ "RobotSpec": ("robot", "RobotSpec"),
45
+ "ContactSensorRef": ("sensor", "ContactSensorRef"),
46
+ "Grid3dPointsRef": ("sensor", "Grid3dPointsRef"),
47
+ "MotionReferenceRef": ("sensor", "MotionReferenceRef"),
48
+ "NativeSensorRef": ("sensor", "NativeSensorRef"),
49
+ "RayCasterRef": ("sensor", "RayCasterRef"),
50
+ "RayPatternRef": ("sensor", "RayPatternRef"),
51
+ "SymmetricAugmentationSpec": ("sensor", "SymmetricAugmentationSpec"),
52
+ "VirtualObstacleRef": ("sensor", "VirtualObstacleRef"),
53
+ "VolumePointsRef": ("sensor", "VolumePointsRef"),
54
+ "AgentSpec": ("task", "AgentSpec"),
55
+ "SceneSpec": ("task", "SceneSpec"),
56
+ "SimSpec": ("task", "SimSpec"),
57
+ "SubTerrainSpec": ("task", "SubTerrainSpec"),
58
+ "TaskSpec": ("task", "TaskSpec"),
59
+ "TerrainGeneratorSpec": ("task", "TerrainGeneratorSpec"),
60
+ "TerrainSpec": ("task", "TerrainSpec"),
61
+ }
62
+
63
+
64
+ def __getattr__(name: str) -> Any:
65
+ try:
66
+ module_name, attribute = _EXPORTS[name]
67
+ except KeyError:
68
+ raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from None
69
+ value = getattr(import_module(f"{__name__}.{module_name}"), attribute)
70
+ globals()[name] = value
71
+ return value
72
+
73
+
74
+ __all__ = [
75
+ "UNIVERSAL_KINDS",
76
+ "ActionTermSpec",
77
+ "AgentSpec",
78
+ "ArticulationRef",
79
+ "BackendAsset",
80
+ "ClockDomainSpec",
81
+ "CollisionExclusionRef",
82
+ "CommandTermSpec",
83
+ "ComponentLifecycleSpec",
84
+ "ContactSensorRef",
85
+ "CurriculumTermSpec",
86
+ "DoneTermSpec",
87
+ "EntityRef",
88
+ "EventTermSpec",
89
+ "Grid3dPointsRef",
90
+ "JointProperties",
91
+ "LifecycleSpec",
92
+ "MdpSpec",
93
+ "MotionReferenceRef",
94
+ "NativeSensorRef",
95
+ "NoiseSpec",
96
+ "ObsGroupSpec",
97
+ "ObsTermSpec",
98
+ "RayCasterRef",
99
+ "RayPatternRef",
100
+ "Requirement",
101
+ "ResolvedClockDomain",
102
+ "RewardTermSpec",
103
+ "RigidObjectRef",
104
+ "RobotSpec",
105
+ "SceneSpec",
106
+ "SimSpec",
107
+ "SubTerrainSpec",
108
+ "SymmetricAugmentationSpec",
109
+ "TaskSpec",
110
+ "TermSpec",
111
+ "TerrainGeneratorSpec",
112
+ "TerrainSpec",
113
+ "VirtualObstacleRef",
114
+ "VolumePointsRef",
115
+ "freeze_task_spec",
116
+ "portability_report",
117
+ "resolve_lifecycle_contract",
118
+ ]
@@ -0,0 +1,15 @@
1
+ """Small name-collection helpers shared by declaration modules."""
2
+
3
+ from __future__ import annotations
4
+
5
+ from collections.abc import Sequence
6
+
7
+
8
+ def as_name_tuple(names: str | Sequence[str]) -> tuple[str, ...]:
9
+ """Normalize one name or an ordered name sequence to an immutable tuple."""
10
+ if isinstance(names, str):
11
+ return (names,)
12
+ return tuple(names)
13
+
14
+
15
+ __all__ = ["as_name_tuple"]
@@ -0,0 +1,34 @@
1
+ """Additional articulated scene entities with canonical joint schemas."""
2
+
3
+ from __future__ import annotations
4
+
5
+ from dataclasses import dataclass
6
+
7
+ from .robot import RobotSpec
8
+
9
+
10
+ @dataclass(frozen=True)
11
+ class ArticulationRef:
12
+ """A named non-primary articulation materialized from a :class:`RobotSpec`.
13
+
14
+ ``RobotSpec.joint_names`` remains the canonical DFS axis. Reusing that
15
+ schema makes policy and observation selectors independent of each native
16
+ simulator's articulation order.
17
+ """
18
+
19
+ name: str
20
+ schema: RobotSpec
21
+
22
+ def __post_init__(self) -> None:
23
+ if not self.name or not self.name.isidentifier():
24
+ raise ValueError(
25
+ f"Additional articulation names must be Python identifiers, got {self.name!r}."
26
+ )
27
+ if self.name in {"robot", "terrain"}:
28
+ raise ValueError(
29
+ f"Additional articulation name {self.name!r} is reserved by the scene."
30
+ )
31
+ self.schema.validate()
32
+
33
+
34
+ __all__ = ["ArticulationRef"]
@@ -0,0 +1,198 @@
1
+ """Capability vocabulary and task requirements shared by tasks and engines.
2
+
3
+ Capabilities are open, namespaced strings. Engine plugins register what they
4
+ provide, while task terms declare both a capability and how strongly they
5
+ require it. Keeping both sides of that protocol here prevents the declaration
6
+ layer from depending on the engine implementation package.
7
+
8
+ "Skip what the engine cannot do" is only safe when the task gets to say which things it can afford
9
+ to lose. Dropping a friction randomisation costs some robustness; dropping an observation changes
10
+ the shape of the policy input, and dropping a reward changes what is being optimised while the run
11
+ still looks healthy. One level cannot cover both, so terms carry a :class:`Requirement`, and the
12
+ compiler acts on it.
13
+
14
+ The defaults follow from that, and are set on each term class rather than chosen per task:
15
+
16
+ =================== ========== ==========================================================
17
+ family default why
18
+ =================== ========== ==========================================================
19
+ observation REQUIRED absence changes the network's input width and meaning
20
+ action REQUIRED absence means the policy cannot act
21
+ termination REQUIRED absence changes the episode structure
22
+ command REQUIRED an observation term reads it
23
+ reward OPTIONAL losing a regulariser is survivable -- but must be recorded
24
+ event / DR OPTIONAL this is where engine capability actually differs
25
+ curriculum OPTIONAL --
26
+ =================== ========== ==========================================================
27
+
28
+ A task overrides per term where its own judgement differs: a locomotion task that is only stable
29
+ because of one particular reward should mark that reward REQUIRED and find out at startup.
30
+
31
+ OPTIONAL does not mean silent. Every skip is recorded in the compilation's ``Resolution`` and
32
+ printed once as a table at startup, because a silently dropped reward term is a changed objective,
33
+ and the resulting policy is otherwise indistinguishable from a healthy one. ``--strict-capabilities``
34
+ promotes every OPTIONAL to REQUIRED for CI and for runs that are meant to be comparable.
35
+ """
36
+
37
+ from __future__ import annotations
38
+
39
+ from collections.abc import Iterable, Mapping
40
+ from dataclasses import dataclass
41
+ from enum import Enum
42
+ from types import MappingProxyType
43
+
44
+ _REGISTRY: dict[str, str] = {}
45
+
46
+
47
+ class UnknownCapability(KeyError):
48
+ """Raised when a capability identifier was never registered."""
49
+
50
+
51
+ def capability(identifier: str, description: str) -> str:
52
+ """Register one namespaced capability and return its identifier."""
53
+ if "." not in identifier:
54
+ raise ValueError(f"{identifier!r} needs a namespace, as in 'contact.air_time'")
55
+ if not description.strip():
56
+ raise ValueError(
57
+ f"{identifier!r} was registered without saying what providing it means"
58
+ )
59
+ existing = _REGISTRY.get(identifier)
60
+ if existing is not None and existing != description:
61
+ raise ValueError(f"{identifier!r} is already registered as {existing!r}")
62
+ _REGISTRY[identifier] = description
63
+ return identifier
64
+
65
+
66
+ def known() -> Mapping[str, str]:
67
+ """Return every registered capability and its meaning."""
68
+ return MappingProxyType(dict(_REGISTRY))
69
+
70
+
71
+ def check_known(values: Iterable[str]) -> None:
72
+ """Reject identifiers that no provider or shared vocabulary registered."""
73
+ unknown = sorted(set(values) - set(_REGISTRY))
74
+ if unknown:
75
+ raise UnknownCapability(
76
+ f"{unknown} are not registered capabilities. An engine package registers what it can do "
77
+ f"when it is imported; registered so far: {sorted(_REGISTRY)}"
78
+ )
79
+
80
+
81
+ BATCHED_SIMULATION = capability(
82
+ "sim.batched", "Many environments stepped together in one call."
83
+ )
84
+ GPU_SIMULATION = capability(
85
+ "sim.gpu", "Physics runs on the GPU with state left in device memory."
86
+ )
87
+ PLANE_TERRAIN = capability("terrain.plane", "An infinite ground plane.")
88
+ ROOT_STATE = capability("state.root", "Reading and writing a body's root pose.")
89
+ ROOT_VELOCITY_WRITE = capability(
90
+ "state.root_velocity", "Writing a root velocity, frame qualified."
91
+ )
92
+ JOINT_STATE = capability(
93
+ "state.joint", "Reading and writing joint positions and velocities."
94
+ )
95
+ BODY_STATE = capability(
96
+ "state.body", "Per-body poses and velocities of an articulation."
97
+ )
98
+ IMPLICIT_POSITION_CONTROL = capability(
99
+ "control.position_implicit", "Joint position targets tracked by the solver."
100
+ )
101
+ EFFORT_CONTROL = capability("control.effort", "Direct joint torque commands.")
102
+ CONTACT_ACTIVE = capability(
103
+ "contact.active", "Whether a body is currently touching something."
104
+ )
105
+ CONTACT_HISTORY = capability(
106
+ "contact.history", "Contact readings kept for several past steps."
107
+ )
108
+ CONTACT_AIR_TIME = capability(
109
+ "contact.air_time", "Durations a body has been in contact or in the air."
110
+ )
111
+ CONTACT_FORCE_VECTOR = capability(
112
+ "contact.force_vector", "Contact force as a vector rather than a magnitude."
113
+ )
114
+ DR_SLIDING_FRICTION = capability(
115
+ "dr.friction.sliding", "Randomising the sliding friction coefficient."
116
+ )
117
+ DR_RESTITUTION = capability("dr.restitution", "Randomising restitution.")
118
+ BODY_MASS_PROPERTIES = capability(
119
+ "body.mass_properties", "Changing a body's mass or inertia after load."
120
+ )
121
+ EXTERNAL_WRENCH = capability(
122
+ "body.external_wrench", "Applying an external force or torque to a body."
123
+ )
124
+ HUMAN_VIEWER = capability("render.human", "An interactive viewer window.")
125
+ RGB_ARRAY = capability("render.rgb_array", "Rendering frames to arrays.")
126
+
127
+
128
+ @dataclass(frozen=True)
129
+ class CapabilitySet:
130
+ """The validated capabilities provided by one engine plugin."""
131
+
132
+ values: frozenset[str]
133
+
134
+ @classmethod
135
+ def of(cls, values: Iterable[str]) -> CapabilitySet:
136
+ collected = frozenset(values)
137
+ check_known(collected)
138
+ return cls(collected)
139
+
140
+ def supports(self, capability: str) -> bool:
141
+ return capability in self.values
142
+
143
+ def require(self, required: Iterable[str], *, context: str) -> None:
144
+ required = frozenset(required)
145
+ check_known(required)
146
+ missing = required.difference(self.values)
147
+ if missing:
148
+ raise RuntimeError(
149
+ f"{context} requires unsupported engine capabilities: {', '.join(sorted(missing))}"
150
+ )
151
+
152
+
153
+ class Requirement(str, Enum):
154
+ """What the compiler does when the engine cannot provide a term."""
155
+
156
+ REQUIRED = "required"
157
+ """Fail at startup. The task is not runnable on this engine and should say so immediately."""
158
+
159
+ OPTIONAL = "optional"
160
+ """Skip it, record it in the resolution, and report it in the startup summary."""
161
+
162
+ EMULATE = "emulate"
163
+ """Substitute the adapter's registered stand-in; fall back to OPTIONAL when it has none.
164
+
165
+ For terms whose effect can be approximated by other means -- a push event realised by writing
166
+ root velocity where an engine has no external-wrench API, say. The substitution is recorded
167
+ separately from a skip, because an emulated term is running *something*, and a later comparison
168
+ between engines needs to know which.
169
+ """
170
+
171
+
172
+ __all__ = [
173
+ "BATCHED_SIMULATION",
174
+ "BODY_MASS_PROPERTIES",
175
+ "BODY_STATE",
176
+ "CONTACT_ACTIVE",
177
+ "CONTACT_AIR_TIME",
178
+ "CONTACT_FORCE_VECTOR",
179
+ "CONTACT_HISTORY",
180
+ "DR_RESTITUTION",
181
+ "DR_SLIDING_FRICTION",
182
+ "EFFORT_CONTROL",
183
+ "EXTERNAL_WRENCH",
184
+ "GPU_SIMULATION",
185
+ "HUMAN_VIEWER",
186
+ "IMPLICIT_POSITION_CONTROL",
187
+ "JOINT_STATE",
188
+ "PLANE_TERRAIN",
189
+ "RGB_ARRAY",
190
+ "ROOT_STATE",
191
+ "ROOT_VELOCITY_WRITE",
192
+ "CapabilitySet",
193
+ "Requirement",
194
+ "UnknownCapability",
195
+ "capability",
196
+ "check_known",
197
+ "known",
198
+ ]
@@ -0,0 +1,156 @@
1
+ """References to a subset of a scene entity, stated without naming an engine.
2
+
3
+ An :class:`EntityRef` says *which parts of which entity* a term acts on -- these joints, those
4
+ bodies -- and leaves it to the backend to turn that into the engine's own selector configuration.
5
+ The translation happens once at compile time, so the runtime cost is nil and terms keep reading
6
+ tensors directly.
7
+
8
+ Selector kinds are open on purpose. Isaac Lab and mjlab agree on only two of them:
9
+
10
+ =============== ==========================================================================
11
+ engine kinds its ``SceneEntityCfg`` accepts
12
+ =============== ==========================================================================
13
+ both ``joint``, ``body``
14
+ Isaac Lab only ``fixed_tendon``, ``object_collection``
15
+ mjlab only ``actuator``, ``camera``, ``geom``, ``light``, ``material``, ``pair``,
16
+ ``site``, ``tendon``
17
+ =============== ==========================================================================
18
+
19
+ Two kinds out of twelve. A fixed pair of ``joints`` / ``bodies`` fields would therefore be able to
20
+ express Isaac Lab tasks and quietly lose everything an mjlab task says about geoms and sites, which
21
+ is the direction this project has to support as well. So the common two get named fields for
22
+ legibility and everything else goes in :attr:`other`, where the backend can either translate it or
23
+ refuse loudly.
24
+
25
+ Note that Isaac Lab's ``fixed_tendon`` and mjlab's ``tendon`` are related but not the same kind,
26
+ and are deliberately not unified here; naming them apart keeps a backend from silently accepting a
27
+ selector it cannot honour.
28
+ """
29
+
30
+ from __future__ import annotations
31
+
32
+ import re
33
+ from collections.abc import Mapping, Sequence
34
+ from dataclasses import dataclass, field
35
+
36
+ __all__ = ["UNIVERSAL_KINDS", "EntityRef", "resolve_entity_names"]
37
+
38
+ UNIVERSAL_KINDS: tuple[str, ...] = ("joint", "body")
39
+ """The only selector kinds every supported engine can express. Checked in the tests."""
40
+
41
+
42
+ def _normalise(patterns: str | Sequence[str] | None) -> tuple[str, ...] | None:
43
+ """Accept a bare string the way both engines' ``SceneEntityCfg`` does."""
44
+ if patterns is None:
45
+ return None
46
+ if isinstance(patterns, str):
47
+ return (patterns,)
48
+ return tuple(patterns)
49
+
50
+
51
+ def resolve_entity_names(
52
+ patterns: str | Sequence[str],
53
+ available_names: Sequence[str],
54
+ *,
55
+ preserve_order: bool,
56
+ ) -> tuple[str, ...]:
57
+ """Resolve selector patterns with the semantics shared by both engines.
58
+
59
+ ``preserve_order=False`` follows ``available_names``. ``True`` groups matches by
60
+ pattern order while retaining ``available_names`` order inside each pattern. A name may
61
+ match only one pattern and every pattern must match, matching both native selector helpers.
62
+ Keeping this tiny resolver in the declaration layer lets validation and compilation reason
63
+ about the selected tensor axis before either engine SDK is imported.
64
+ """
65
+ expressions = (patterns,) if isinstance(patterns, str) else tuple(patterns)
66
+ if not expressions:
67
+ raise ValueError("A selector was given no patterns.")
68
+
69
+ matches_by_pattern: list[list[str]] = [[] for _ in expressions]
70
+ source_order: list[str] = []
71
+ for name in available_names:
72
+ matching = [
73
+ index
74
+ for index, expression in enumerate(expressions)
75
+ if re.fullmatch(expression, name)
76
+ ]
77
+ if len(matching) > 1:
78
+ duplicate_patterns = tuple(expressions[index] for index in matching)
79
+ raise ValueError(
80
+ f"Entity name {name!r} matches multiple selector patterns: {duplicate_patterns!r}."
81
+ )
82
+ if matching:
83
+ matches_by_pattern[matching[0]].append(name)
84
+ source_order.append(name)
85
+
86
+ unmatched = [
87
+ expression
88
+ for expression, matches in zip(expressions, matches_by_pattern)
89
+ if not matches
90
+ ]
91
+ if unmatched:
92
+ raise ValueError(
93
+ f"Selector patterns match no entity names: {unmatched!r}. Available names: {tuple(available_names)!r}."
94
+ )
95
+ if not preserve_order:
96
+ return tuple(source_order)
97
+ return tuple(name for matches in matches_by_pattern for name in matches)
98
+
99
+
100
+ @dataclass(frozen=True)
101
+ class EntityRef:
102
+ """A subset of one scene entity, named by pattern rather than by index.
103
+
104
+ Patterns are regular expressions matched against the entity's own names, which is what both
105
+ engines do; the matching helper is byte-identical between them, so a pattern selects the same
106
+ thing either way.
107
+
108
+ Args:
109
+ entity: Key of the entity in the scene.
110
+ joints: Joint name patterns, or a single pattern.
111
+ bodies: Body name patterns, or a single pattern.
112
+ other: Patterns for selector kinds outside :data:`UNIVERSAL_KINDS`, keyed by kind. A
113
+ backend that cannot express a kind must reject the reference rather than drop it.
114
+ preserve_order: When true the selection follows the order of the patterns; when false it
115
+ follows the entity's own order. Note what the entity's own order is not: it is whatever
116
+ the engine built, which is a breadth-first walk under PhysX and model-file order under
117
+ MuJoCo, and those two disagree. D1's canonical depth-first order is therefore reached
118
+ only by passing the catalog's joint names explicitly with this flag set -- a bare ``.*``
119
+ preserves the order of a one-element pattern list and so changes nothing.
120
+ """
121
+
122
+ entity: str = "robot"
123
+ joints: str | Sequence[str] | None = None
124
+ bodies: str | Sequence[str] | None = None
125
+ other: Mapping[str, str | Sequence[str]] = field(default_factory=dict)
126
+ preserve_order: bool = False
127
+
128
+ def __post_init__(self) -> None:
129
+ object.__setattr__(self, "joints", _normalise(self.joints))
130
+ object.__setattr__(self, "bodies", _normalise(self.bodies))
131
+ normalised = {
132
+ kind: _normalise(patterns) for kind, patterns in dict(self.other).items()
133
+ }
134
+ for kind, patterns in normalised.items():
135
+ if kind in UNIVERSAL_KINDS:
136
+ raise ValueError(
137
+ f"'{kind}' has its own field on EntityRef; do not pass it through 'other'."
138
+ )
139
+ if not patterns:
140
+ raise ValueError(f"Selector '{kind}' was given no patterns.")
141
+ # Sorted so that two references built from equivalent mappings compare equal.
142
+ object.__setattr__(self, "other", dict(sorted(normalised.items())))
143
+
144
+ def selectors(self) -> dict[str, tuple[str, ...]]:
145
+ """Every selector on this reference, keyed by kind, in a single mapping."""
146
+ out: dict[str, tuple[str, ...]] = {}
147
+ if self.joints is not None:
148
+ out["joint"] = self.joints
149
+ if self.bodies is not None:
150
+ out["body"] = self.bodies
151
+ out.update(self.other) # type: ignore[arg-type]
152
+ return out
153
+
154
+ def kinds(self) -> frozenset[str]:
155
+ """The selector kinds this reference uses. A backend checks it against what it supports."""
156
+ return frozenset(self.selectors())