| from __future__ import annotations |
|
|
| from dataclasses import dataclass, field |
| from typing import Any, Protocol |
|
|
| from dovla_cil.tasks.schema import SceneSpec, TaskSpec |
|
|
| Observation = dict[str, Any] |
|
|
|
|
| @dataclass(frozen=True) |
| class SimState: |
| task_id: str |
| scene_id: str | None |
| seed: int | None |
| symbolic_state: dict[str, Any] |
| metadata: dict[str, Any] = field(default_factory=dict) |
|
|
|
|
| class ActionChunk: |
| """Semantic simulator action chunk. |
| |
| Examples: |
| - `ActionChunk.single("move_to", object="red_mug")` |
| - `ActionChunk([{"command": "grasp", "object": "red_mug"}])` |
| """ |
|
|
| def __init__( |
| self, |
| actions: dict[str, Any] | list[dict[str, Any]] | tuple[dict[str, Any], ...] | None = None, |
| **single_action: Any, |
| ) -> None: |
| if actions is None: |
| if not single_action: |
| raise ValueError("ActionChunk requires at least one action") |
| actions = single_action |
| if isinstance(actions, dict): |
| action_list = [actions] |
| else: |
| action_list = list(actions) |
| if not action_list: |
| raise ValueError("ActionChunk requires at least one action") |
| normalized: list[dict[str, Any]] = [] |
| for action in action_list: |
| if not isinstance(action, dict): |
| raise TypeError("Each semantic action must be a dict") |
| if not _action_command(action): |
| raise ValueError("Each semantic action requires 'command', 'type', 'op', or 'name'") |
| normalized.append(dict(action)) |
| self.actions = tuple(normalized) |
|
|
| @classmethod |
| def single(cls, command: str, **params: Any) -> ActionChunk: |
| return cls({"command": command, **params}) |
|
|
| def to_dict(self) -> dict[str, Any]: |
| return {"actions": [dict(action) for action in self.actions]} |
|
|
| def __iter__(self): |
| return iter(self.actions) |
|
|
| def __eq__(self, other: Any) -> bool: |
| return isinstance(other, ActionChunk) and self.actions == other.actions |
|
|
| def __repr__(self) -> str: |
| return f"ActionChunk(actions={self.actions!r})" |
|
|
|
|
| @dataclass(frozen=True) |
| class RolloutResult: |
| observation: Observation |
| reward: float |
| done: bool |
| info: dict[str, Any] = field(default_factory=dict) |
| trajectory: list[dict[str, Any]] = field(default_factory=list) |
| contacts: list[dict[str, Any]] = field(default_factory=list) |
| before_state: dict[str, Any] = field(default_factory=dict) |
| after_state: dict[str, Any] = field(default_factory=dict) |
|
|
| @property |
| def next_observation(self) -> Observation: |
| """Compatibility alias for the earlier scaffold.""" |
|
|
| return self.observation |
|
|
|
|
| SimulatorTransition = RolloutResult |
|
|
|
|
| class SimulatorBackend(Protocol): |
| name: str |
|
|
| def seed(self, seed: int) -> None: |
| """Set deterministic simulator seed.""" |
|
|
| def reset_task(self, task: TaskSpec, scene: SceneSpec | None = None) -> SimState: |
| """Reset simulator to a task/scene and return the initial symbolic state.""" |
|
|
| def serialize_state(self) -> bytes: |
| """Serialize the exact simulator state.""" |
|
|
| def restore_state(self, state_blob: bytes) -> None: |
| """Restore an exact serialized simulator state.""" |
|
|
| def render_observation(self) -> Observation: |
| """Render or expose the current observation.""" |
|
|
| def get_symbolic_state(self) -> dict[str, Any]: |
| """Return a symbolic state dictionary.""" |
|
|
| def execute_action_chunk(self, action: ActionChunk) -> RolloutResult: |
| """Execute an action chunk and return rollout metadata.""" |
|
|
| def close(self) -> None: |
| """Release resources.""" |
|
|
|
|
| def _action_command(action: dict[str, Any]) -> str: |
| return str( |
| action.get("command") or action.get("type") or action.get("op") or action.get("name") or "" |
| ) |
|
|
