Spaces:
Sleeping
Sleeping
File size: 4,723 Bytes
aab74bd | 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 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 | # Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.
"""Store Manager Environment Client."""
from typing import Any, Dict
from openenv.core import EnvClient
from openenv.core.client_types import StepResult
from .models import (
ProductState,
StoremanagerAction,
StoremanagerObservation,
StoremanagerState,
)
class StoremanagerEnv(
EnvClient[StoremanagerAction, StoremanagerObservation, StoremanagerState]
):
"""
Client for the Store Manager Environment.
Maintains a persistent WebSocket connection to the environment server,
enabling efficient multi-step interactions with lower latency.
Example:
>>> with StoremanagerEnv(base_url="http://localhost:8000") as env:
... result = env.reset()
... print([p.name for p in result.observation.inventory])
...
... action = StoremanagerAction(product_id=0, discount_pct=20)
... result = env.step(action)
... print(f"Reward: {result.reward:.2f}")
Example with Docker:
>>> client = StoremanagerEnv.from_docker_image("StoreManager-env:latest")
>>> try:
... result = client.reset()
... result = client.step(StoremanagerAction(product_id=0, discount_pct=10))
... finally:
... client.close()
"""
def _step_payload(self, action: StoremanagerAction) -> Dict:
"""Convert StoremanagerAction to JSON payload for the step message."""
payload: Dict = {
"action_type": action.action_type,
"product_id": action.product_id,
"discount_pct": action.discount_pct,
}
if action.restock_quantity is not None:
payload["restock_quantity"] = action.restock_quantity
if action.target_zone is not None:
payload["target_zone"] = action.target_zone
return payload
def _parse_result(self, payload: Dict) -> StepResult[StoremanagerObservation]:
"""Parse server response into StepResult[StoremanagerObservation]."""
obs_data = payload.get("observation", {})
raw_inventory = obs_data.get("inventory", [])
inventory = [ProductState(**p) for p in raw_inventory if isinstance(p, dict)]
observation = StoremanagerObservation(
inventory=inventory,
current_step=obs_data.get("current_step", 0),
max_steps=obs_data.get("max_steps", 30),
steps_remaining=obs_data.get("steps_remaining", 0),
cumulative_profit=obs_data.get("cumulative_profit", 0.0),
last_step_profit=obs_data.get("last_step_profit", 0.0),
last_step_units_sold=obs_data.get("last_step_units_sold", 0),
last_expired_products=obs_data.get("last_expired_products", []),
num_customers_per_step=obs_data.get("num_customers_per_step", 10),
error=obs_data.get("error"),
zone_capacity=obs_data.get("zone_capacity", {1: 2, 2: 4, 3: 99}),
zone_occupancy=obs_data.get("zone_occupancy", {1: 0, 2: 0, 3: 0}),
zone_multipliers=obs_data.get("zone_multipliers", {1: 1.5, 2: 1.0, 3: 0.6}),
holding_cost_this_step=obs_data.get("holding_cost_this_step", 0.0),
placement_cost_this_step=obs_data.get("placement_cost_this_step", 0.0),
restock_cost_this_step=obs_data.get("restock_cost_this_step", 0.0),
unjustified_discount_penalty=obs_data.get("unjustified_discount_penalty", 0.0),
done=payload.get("done", False),
reward=payload.get("reward"),
metadata=obs_data.get("metadata", {}),
)
return StepResult(
observation=observation,
reward=payload.get("reward"),
done=payload.get("done", False),
)
def _parse_state(self, payload: Dict[str, Any]) -> StoremanagerState:
"""Parse state endpoint response into StoremanagerState."""
return StoremanagerState(
episode_id=payload.get("episode_id"),
step_count=payload.get("step_count", 0),
num_products=payload.get("num_products", 8),
max_steps=payload.get("max_steps", 30),
cumulative_profit=payload.get("cumulative_profit", 0.0),
last_reward=payload.get("last_reward", 0.0),
active_product_count=payload.get("active_product_count", 0),
inventory_snapshot=payload.get("inventory_snapshot", []),
)
|