MiniGridEnv / env /models.py
yashu2000's picture
Upload folder using huggingface_hub
a03a89b verified
"""OpenEnv Pydantic models for MiniGridEnv."""
from typing import Any, Optional
from pydantic import BaseModel, ConfigDict, Field
try:
from openenv.core.env_server.types import Action as _ActionBase
from openenv.core.env_server.types import Observation as _ObservationBase
from openenv.core.env_server.types import State as _StateBase
except ImportError:
_ActionBase = BaseModel
_ObservationBase = BaseModel
_StateBase = BaseModel
class MiniGridAction(_ActionBase):
"""Agent action represented as a natural language command."""
if _ActionBase is BaseModel:
model_config = ConfigDict(extra="forbid")
metadata: dict[str, Any] = Field(default_factory=dict)
command: str = Field(
...,
description="Natural language command (for example: 'go forward', 'turn left', 'pickup').",
)
thought: Optional[str] = Field(
default=None,
description="Optional reasoning trace for logging/analysis.",
)
class MiniGridObservation(_ObservationBase):
"""Text observation derived from MiniGrid's egocentric 7x7 view."""
if _ObservationBase is BaseModel:
model_config = ConfigDict(extra="forbid")
done: bool = Field(default=False)
reward: Optional[float] = Field(default=None)
metadata: dict[str, Any] = Field(default_factory=dict)
text: str = Field(..., description="Natural-language environment description.")
mission: str = Field(..., description="Current mission instruction.")
step_idx: int = Field(..., ge=0, description="Current 0-indexed step number.")
steps_remaining: int = Field(..., ge=0, description="Steps left before truncation.")
max_steps: int = Field(..., ge=1, description="Maximum number of steps in this episode.")
history: list[dict[str, Any]] = Field(
default_factory=list,
description="Recent step history entries for prompt reconstruction.",
)
level_name: str = Field(default="", description="Selected BabyAI level short name.")
last_action: Optional[str] = Field(default=None, description="Canonical last action string.")
action_success: Optional[bool] = Field(
default=None,
description="Whether the last action had an observable effect.",
)
class MiniGridState(_StateBase):
"""Episode-level state metrics for logging and debugging."""
if _StateBase is BaseModel:
model_config = ConfigDict(extra="allow")
episode_id: Optional[str] = Field(default=None)
step_count: int = Field(default=0, ge=0)
level_name: str = Field(default="", description="Selected level name.")
level_difficulty: int = Field(default=0, ge=0, description="Difficulty stage index.")
completed: bool = Field(default=False, description="True when mission succeeds.")
truncated: bool = Field(default=False, description="True when max step budget is exhausted.")
total_reward: float = Field(default=0.0, description="Cumulative reward this episode.")
steps_taken: int = Field(default=0, ge=0, description="Total steps executed.")
optimal_steps: Optional[int] = Field(
default=None,
ge=0,
description="Optional best-path length from the BabyAI bot.",
)
efficiency_ratio: Optional[float] = Field(
default=None,
description="optimal_steps / steps_taken when available.",
)
valid_actions: int = Field(default=0, ge=0, description="Count of valid parsed actions.")
invalid_actions: int = Field(default=0, ge=0, description="Count of invalid parsed actions.")
action_distribution: dict[str, int] = Field(
default_factory=dict,
description="Histogram of canonical actions used so far.",
)