File size: 918 Bytes
d88694d
 
 
 
 
 
 
97ac6b2
d88694d
 
3da7d69
 
d88694d
 
 
 
 
 
 
 
 
 
 
 
 
 
97ac6b2
d88694d
 
 
 
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
"""
Client wrapper for the warehouse fulfillment environment.
"""

from __future__ import annotations

from .env import WarehouseFulfillmentEnv
from .models import WarehouseAction, WarehouseObservation, WarehouseState

try:
    from openenv.core.env_client import EnvClient  
except ImportError:  
    EnvClient = object


class WarehouseEnvClient(EnvClient):
    """
    Local client facade used when OpenEnv transport is not available.
    """

    def __init__(self, task_id: str = "easy_single_pick", seed: int = 7) -> None:
        self.env = WarehouseFulfillmentEnv(task_id=task_id, seed=seed)

    def reset(self, task_id: str | None = None, seed: int | None = None) -> WarehouseObservation:
        return self.env.reset(task_id=task_id, seed=seed)

    def step(self, action: WarehouseAction | str):
        return self.env.step(action)

    def state(self) -> WarehouseState:
        return self.env.state()