File size: 1,183 Bytes
05b0a5c 29edd33 05b0a5c 29edd33 05b0a5c 29edd33 05b0a5c 29edd33 05b0a5c 29edd33 05b0a5c 29edd33 05b0a5c 094509b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 | from dataclasses import dataclass
from typing import Any
from openenv.env import SupportEnv
from openenv.models import Action as SupportAction
@dataclass
class ResetResult:
observation: Any
done: bool
@dataclass
class StepResult:
observation: Any
reward: float
done: bool
phase: int # phase AFTER the transition (what the env is now waiting for)
class SupportEnvWrapper:
def __init__(self, task: str = "easy"):
self.env = SupportEnv(task_name=task)
async def reset(self) -> ResetResult:
result = self.env.reset()
return ResetResult(observation=result.observation, done=result.done)
async def step(self, action: SupportAction) -> StepResult:
result = self.env.step(action)
return StepResult(
observation=result.observation,
reward=result.reward,
done=result.done,
phase=self.env.phase, # actual phase from env, not a counter
)
async def close(self) -> None:
self.env.close()
@classmethod
async def from_docker_image(cls, image_name: str = None, task: str = "easy") -> "SupportEnvWrapper":
return cls(task=task)
|