server/StoreManager_environment.py

# 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 Implementation.



A retail store simulator where an LLM agent acts as the store manager.

The agent can apply discounts, restock inventory, or move products between

shelf zones each step to maximise cumulative profit over a fixed episode horizon.



Episode dynamics:

  - N products, each with a cost price, selling price, stock quantity,

    expiry step, a base pick-probability weight, a shelf zone, and restock params.

  - Each step the agent chooses one action:

      - "discount": apply a discount (0/10/20/50%) to one product

      - "restock":  place a reorder for a product (delivered after lead_time steps)

      - "place":    move a product to a different shelf zone (1/2/3)

  - Discounts are ephemeral — they apply only for the current step.

  - Zone multiplier scales pick probability: zone 1 = 1.5x, 2 = 1.0x, 3 = 0.6x.

  - 10 customers sample products per step; probabilities incorporate zone + discount.

  - Holding cost is charged every step based on quantity on hand.

  - Restock cost is charged upfront; stock arrives after lead_time steps.

  - Placement cost is a flat labor fee per move.

  - Reward = step_revenue - step_cost - expiry_penalty - holding_cost

             - placement_cost - restock_cost.

  - Expiry penalty = remaining_quantity * cost_price for each expired product.

  - Episode ends after max_steps steps, or when all products are inactive.

"""

from __future__ import annotations

import sys
import os

# Ensure the project root (parent of server/) is on sys.path so that
# models.py and tasks.py are importable whether this file is run directly
# or imported as part of the package.
_ROOT = os.path.join(os.path.dirname(__file__), "..")
if _ROOT not in sys.path:
    sys.path.insert(0, _ROOT)

import math
from typing import Any, Dict, List, Optional
from uuid import uuid4

import numpy as np
from openenv.core.env_server.interfaces import Environment

from models import (
    ProductState,
    StoremanagerAction,
    StoremanagerObservation,
    StoremanagerState,
)

try:
    from tasks import TaskConfig
except ImportError:
    TaskConfig = None  # type: ignore[assignment,misc]

# ── Curated product catalogue ──────────────────────────────────────────────────
_PRODUCT_NAMES: List[str] = [
    "Milk", "Bread", "Cheese", "Yogurt", "Juice", "Butter",
    "Eggs", "Cream", "Honey", "Jam", "Soda", "Water",
    "Tea", "Coffee", "Flour", "Rice",
]

# ── Zone constants ─────────────────────────────────────────────────────────────
ZONE_MULTIPLIERS: Dict[int, float] = {1: 1.5, 2: 1.0, 3: 0.6}
ZONE_CAPACITY: Dict[int, int] = {1: 2, 2: 4, 3: 99}
PLACEMENT_COST: float = 3.0  # flat labor fee per zone move


class StoremanagerEnvironment(Environment):
    """

    Retail store manager RL environment.



    The agent observes full inventory state and chooses one action per step:

    apply a discount, place a restock order, or move a product to a better zone.



    Args:

        num_products: Number of distinct products in the store (default 8).

        max_steps: Episode length in steps (default 30).

        num_customers: Customers simulated per step (default 10).

        seed: Optional base random seed for reproducible episodes.

    """

    SUPPORTS_CONCURRENT_SESSIONS: bool = True

    def __init__(

        self,

        task_config: Optional["TaskConfig"] = None,  # type: ignore[name-defined]

        num_products: int = 8,

        max_steps: int = 30,

        num_customers: int = 10,

        seed: Optional[int] = None,

    ) -> None:
        super().__init__()
        if task_config is not None:
            num_products = task_config.num_products
            max_steps = task_config.max_steps
            num_customers = task_config.num_customers
            self._expiry_offset_min = task_config.expiry_offset_min
            self._expiry_offset_max = task_config.expiry_offset_max
            self._task_name = task_config.name
            self._profit_target = task_config.profit_target
        else:
            self._expiry_offset_min = 8
            self._expiry_offset_max = 25
            self._task_name = "custom"
            # Default: assume ~$8 average profit per step for custom configs
            self._profit_target = 8.0 * max_steps

        self._num_products = min(num_products, len(_PRODUCT_NAMES))
        self._max_steps = max_steps
        self._num_customers = num_customers
        self._base_seed = seed

        self._rng: np.random.Generator = np.random.default_rng(seed)
        self._inventory: List[ProductState] = []
        self._pending_orders: List[Dict[str, Any]] = []
        self._current_step: int = 0
        self._cumulative_profit: float = 0.0

        self._state = StoremanagerState(
            num_products=self._num_products,
            max_steps=self._max_steps,
        )

    # ── OpenEnv API ───────────────────────────────────────────────────────────

    def reset(

        self,

        seed: Optional[int] = None,

        episode_id: Optional[str] = None,

        **kwargs: Any,

    ) -> StoremanagerObservation:
        """Reset the store and generate a fresh inventory.



        Accepts an optional ``task`` keyword argument to dynamically switch the

        task configuration without restarting the server.  This allows a single

        server instance to serve all difficulty levels sequentially.

        """
        # ── Dynamic task reconfiguration ────────────────────────────────────
        task_name = kwargs.get("task")
        if task_name:
            try:
                from tasks import TASKS  # local import to avoid circular deps
                task_config = TASKS.get(task_name)
                if task_config is not None:
                    self._num_products = min(task_config.num_products, len(_PRODUCT_NAMES))
                    self._max_steps = task_config.max_steps
                    self._num_customers = task_config.num_customers
                    self._expiry_offset_min = task_config.expiry_offset_min
                    self._expiry_offset_max = task_config.expiry_offset_max
                    self._task_name = task_config.name
                    self._profit_target = task_config.profit_target
            except ImportError:
                pass  # tasks module unavailable; keep current config

        self._reset_rubric()

        effective_seed = seed if seed is not None else self._base_seed
        self._rng = np.random.default_rng(effective_seed)

        self._current_step = 0
        self._cumulative_profit = 0.0
        self._pending_orders = []
        self._inventory = self._generate_inventory()
        self._recompute_pick_probs()

        self._update_state(episode_id=episode_id or str(uuid4()), last_reward=self._normalize_reward(0.0))
        self._refresh_product_derived_fields()

        return StoremanagerObservation(
            inventory=list(self._inventory),
            current_step=0,
            max_steps=self._max_steps,
            steps_remaining=self._max_steps,
            cumulative_profit=0.0,
            last_step_profit=0.0,
            last_step_units_sold=0,
            last_expired_products=[],
            num_customers_per_step=self._num_customers,
            task_name=self._task_name,
            zone_capacity=dict(ZONE_CAPACITY),
            zone_occupancy=self._zone_occupancy(),
            zone_multipliers=dict(ZONE_MULTIPLIERS),
            holding_cost_this_step=0.0,
            placement_cost_this_step=0.0,
            restock_cost_this_step=0.0,
            unjustified_discount_penalty=0.0,
            done=False,
            reward=self._normalize_reward(0.0),
        )

    def step(

        self,

        action: StoremanagerAction,

        timeout_s: Optional[float] = None,

        **kwargs: Any,

    ) -> StoremanagerObservation:
        """Execute one step: process deliveries, apply action, simulate customers, handle expiry."""
        # ── 1. Validate product_id ──────────────────────────────────────────
        valid_ids = {p.product_id for p in self._inventory}
        if action.product_id not in valid_ids:
            return self._error_observation(
                f"Invalid product_id {action.product_id}. "
                f"Valid IDs: {sorted(valid_ids)}"
            )

        # ── 2. Process pending deliveries ───────────────────────────────────
        delivered = [o for o in self._pending_orders if o["arrive_at_step"] <= self._current_step]
        for order in delivered:
            product = next((p for p in self._inventory if p.product_id == order["product_id"]), None)
            if product is not None:
                product.quantity += order["quantity"]
                product.is_active = True
                product.pending_restock_quantity = 0
                product.pending_restock_arrives_at_step = None
        self._pending_orders = [o for o in self._pending_orders if o["arrive_at_step"] > self._current_step]

        # ── 3. Compute holding cost (charged before action) ─────────────────
        holding_cost = sum(
            p.quantity * p.cost_price * p.holding_cost_rate
            for p in self._inventory
            if p.is_active and p.quantity > 0
        )

        # ── 4. Reset per-step discount tracker (selling_price persists across steps) ──
        # current_discount_pct tracks only this step's discount for the pick-prob formula.
        # selling_price is NOT reset — discounts compound: each action cuts the current price.
        for product in self._inventory:
            product.current_discount_pct = 0.0

        # ── 5. Dispatch action ──────────────────────────────────────────────
        placement_cost = 0.0
        restock_cost = 0.0
        target = next(p for p in self._inventory if p.product_id == action.product_id)

        if action.action_type == "discount":
            if not target.is_active:
                return self._error_observation(
                    f"Product '{target.name}' (id={target.product_id}) is no longer active "
                    f"(sold out or expired). Choose an active product to discount."
                )
            if target.quantity <= 0:
                return self._error_observation(
                    f"Product '{target.name}' (id={target.product_id}) has no stock remaining."
                )
            target.current_discount_pct = float(action.discount_pct)
            # Phase 1: dynamic expiry-aware price floor.
            # Far from expiry: protect margin — floor at cost_price.
            # Near expiry: allow progressive below-cost markdown to clear stock
            # before the far-worse full expiry penalty hits.
            steps_left = target.expiry_step - self._current_step
            if steps_left >= 4:
                price_floor = target.cost_price           # full protection
            elif steps_left == 3:
                price_floor = round(target.cost_price * 0.70, 2)  # 30% loss/unit max
            elif steps_left == 2:
                price_floor = round(target.cost_price * 0.40, 2)  # 60% loss/unit max
            else:                                                    # 1 step left
                price_floor = round(target.cost_price * 0.10, 2)  # 90% loss/unit max
            target.selling_price = round(
                max(price_floor, target.selling_price * (1.0 - action.discount_pct / 100.0)),
                2,
            )
            # popularity_multiplier compounds up: repeated discounting makes the
            # product progressively more visible/attractive to customers
            if action.discount_pct > 0:
                target.popularity_multiplier = round(
                    target.popularity_multiplier * (1.0 + action.discount_pct / 100.0), 6
                )

        elif action.action_type == "restock":
            qty = action.restock_quantity
            if qty is None or qty <= 0:
                return self._error_observation(
                    f"restock_quantity must be a positive integer, got {qty!r}."
                )
            if target.expiry_step <= self._current_step:
                return self._error_observation(
                    f"Product '{target.name}' (id={target.product_id}) has expired "
                    f"(expiry_step={target.expiry_step}, current_step={self._current_step}). "
                    f"Restocking an expired product wastes money — the stock would "
                    f"immediately re-expire and incur an expiry penalty."
                )
            if target.pending_restock_quantity > 0:
                return self._error_observation(
                    f"Product '{target.name}' already has a pending restock order "
                    f"for {target.pending_restock_quantity} units arriving at step "
                    f"{target.pending_restock_arrives_at_step}. Wait for delivery first."
                )
            arrive_at = self._current_step + target.restock_lead_time
            if arrive_at >= target.expiry_step:
                return self._error_observation(
                    f"Restock order for '{target.name}' would arrive at step {arrive_at} "
                    f"but the product expires at step {target.expiry_step}. "
                    f"The delivered units would immediately expire and be penalized. "
                    f"Order cancelled to prevent wasted spend."
                )
            self._pending_orders.append({
                "product_id": target.product_id,
                "quantity": qty,
                "arrive_at_step": arrive_at,
            })
            target.pending_restock_quantity = qty
            target.pending_restock_arrives_at_step = arrive_at
            restock_cost = qty * target.restock_cost_per_unit

        elif action.action_type == "place":
            zone = action.target_zone
            if zone is None:
                return self._error_observation("target_zone must be specified for 'place' actions.")
            if zone == target.zone:
                return self._error_observation(
                    f"Product '{target.name}' is already in zone {zone}."
                )
            # Count active products currently in the target zone (excluding this product)
            occupants = sum(
                1 for p in self._inventory
                if p.is_active and p.zone == zone and p.product_id != target.product_id
            )
            if occupants >= ZONE_CAPACITY[zone]:
                return self._error_observation(
                    f"Zone {zone} is at capacity ({ZONE_CAPACITY[zone]} products). "
                    f"Move another product out first."
                )
            target.zone = zone
            placement_cost = PLACEMENT_COST

        # ── 6. Gather sellable products ─────────────────────────────────────
        active = [p for p in self._inventory if p.is_active and p.quantity > 0]
        if not active:
            done = True
            self._update_state(last_reward=self._normalize_reward(0.0))
            self._refresh_product_derived_fields()
            return StoremanagerObservation(
                inventory=list(self._inventory),
                current_step=self._current_step,
                max_steps=self._max_steps,
                steps_remaining=max(0, self._max_steps - self._current_step),
                cumulative_profit=self._cumulative_profit,
                last_step_profit=0.0,
                last_step_units_sold=0,
                last_expired_products=[],
                num_customers_per_step=self._num_customers,
                task_name=self._task_name,
                zone_capacity=dict(ZONE_CAPACITY),
                zone_occupancy=self._zone_occupancy(),
                zone_multipliers=dict(ZONE_MULTIPLIERS),
                holding_cost_this_step=round(holding_cost, 4),
                placement_cost_this_step=placement_cost,
                restock_cost_this_step=round(restock_cost, 4),
                unjustified_discount_penalty=0.0,
                done=done,
                reward=self._normalize_reward(0.0),
            )

        # ── 7. Compute pick probabilities (zone + discount + popularity) ────
        effective_weights = np.array([
            p.base_probability
            * ZONE_MULTIPLIERS[p.zone]
            * (1.0 + p.current_discount_pct / 100.0)
            * p.popularity_multiplier
            for p in active
        ])
        pick_probs = effective_weights / effective_weights.sum()

        for product, prob in zip(active, pick_probs):
            product.effective_pick_prob = round(float(prob), 6)

        # ── 8. Simulate customers ───────────────────────────────────────────
        chosen_ids = self._rng.choice(
            a=[p.product_id for p in active],
            size=self._num_customers,
            p=pick_probs,
            replace=True,
        )
        sales_counts: dict[int, int] = {}
        for pid in chosen_ids:
            sales_counts[pid] = sales_counts.get(int(pid), 0) + 1

        # ── 9. Apply sales (clamp by available quantity) ────────────────────
        step_revenue = 0.0
        step_cost = 0.0
        units_sold_total = 0
        unjustified_discount_penalty = 0.0

        for product in active:
            wanted = sales_counts.get(product.product_id, 0)
            sold = min(wanted, product.quantity)
            if sold > 0:
                step_revenue += sold * product.selling_price
                step_cost += sold * product.cost_price
                product.quantity -= sold
                units_sold_total += sold
                product.sales_history.append(
                    {
                        "step": self._current_step,
                        "discount_pct": product.current_discount_pct,
                        "units_sold": sold,
                    }
                )
                # Phase 3: if selling below cost with no urgency, double the per-unit
                # loss as an extra penalty — gives a sharp "this is bad" signal.
                if product.selling_price < product.cost_price:
                    steps_left = product.expiry_step - self._current_step
                    if steps_left >= 4:
                        unjustified_discount_penalty += sold * (product.cost_price - product.selling_price)
            if product.quantity == 0:
                product.is_active = False
                product.effective_pick_prob = 0.0

        # ── 10. Advance step counter ────────────────────────────────────────
        self._current_step += 1

        # ── 11. Check expiry ────────────────────────────────────────────────
        expiry_penalty = 0.0
        expired_names: List[str] = []
        for product in self._inventory:
            if product.is_active and product.expiry_step <= self._current_step:
                expiry_penalty += product.quantity * product.cost_price
                expired_names.append(product.name)
                product.is_active = False
                product.quantity = 0
                product.effective_pick_prob = 0.0

        # ── 12. Reward and cumulative profit ────────────────────────────────
        step_profit = (
            step_revenue - step_cost - expiry_penalty
            - holding_cost - placement_cost - restock_cost
            - unjustified_discount_penalty
        )
        self._cumulative_profit += step_profit

        # ── 13. Terminal condition ───────────────────────────────────────────
        done = self._current_step >= self._max_steps or all(
            not p.is_active for p in self._inventory
        )

        # ── 14. State update and return ─────────────────────────────────────
        normalized_reward = self._normalize_reward(step_profit)
        self._update_state(last_reward=normalized_reward)
        self._refresh_product_derived_fields()

        return StoremanagerObservation(
            inventory=list(self._inventory),
            current_step=self._current_step,
            max_steps=self._max_steps,
            steps_remaining=max(0, self._max_steps - self._current_step),
            cumulative_profit=self._cumulative_profit,
            last_step_profit=step_profit,
            last_step_units_sold=units_sold_total,
            last_expired_products=expired_names,
            num_customers_per_step=self._num_customers,
            task_name=self._task_name,
            zone_capacity=dict(ZONE_CAPACITY),
            zone_occupancy=self._zone_occupancy(),
            zone_multipliers=dict(ZONE_MULTIPLIERS),
            holding_cost_this_step=round(holding_cost, 4),
            placement_cost_this_step=placement_cost,
            restock_cost_this_step=round(restock_cost, 4),
            unjustified_discount_penalty=round(unjustified_discount_penalty, 4),
            done=done,
            reward=normalized_reward,
            metadata={
                "step_revenue": round(step_revenue, 4),
                "step_cost": round(step_cost, 4),
                "expiry_penalty": round(expiry_penalty, 4),
                "holding_cost": round(holding_cost, 4),
                "placement_cost": placement_cost,
                "restock_cost": round(restock_cost, 4),
                "unjustified_discount_penalty": round(unjustified_discount_penalty, 4),
                "step_profit": round(step_profit, 4),
            },
        )

    @property
    def state(self) -> StoremanagerState:
        return self._state

    # ── Private helpers ───────────────────────────────────────────────────────

    def _normalize_reward(self, step_profit: float) -> float:
        """

        Map raw step profit to a meaningful reward strictly in (0.001, 0.999) using a

        sigmoid centered at break-even (step_profit = 0).



        Formula:

            x = step_profit / per_step_target   (normalise by average needed per step)

            reward = clamp(sigmoid(k * x), 0.001, 0.999)



        Where k=2 gives a useful spread:

            step_profit =  0            -> 0.50  (neutral, broke even)

            step_profit =  per_step_tgt -> 0.88  (on-track for target)

            step_profit = 2*per_step_tgt -> 0.98 (excellent step)

            step_profit = -per_step_tgt -> 0.12  (below break-even)

            step_profit = -2*per_step_tgt -> 0.02 (large penalty step)



        x is clamped to [-7, 7] to prevent math.exp() OverflowError on extreme

        profits, and the result is clamped to (0.001, 0.999) to ensure the reward

        is never exactly 0.0 or 1.0.

        """
        per_step_target = self._profit_target / self._max_steps
        x = step_profit / per_step_target if per_step_target != 0 else 0.0
        x = max(-7.0, min(7.0, x))
        raw = 1.0 / (1.0 + math.exp(-2.0 * x))
        return round(max(0.001, min(0.999, raw)), 6)

    def _zone_occupancy(self) -> Dict[int, int]:
        """Count active products in each zone."""
        occupancy: Dict[int, int] = {1: 0, 2: 0, 3: 0}
        for p in self._inventory:
            if p.is_active and p.zone in occupancy:
                occupancy[p.zone] += 1
        return occupancy

    def _refresh_product_derived_fields(self) -> None:
        """Recompute ephemeral derived fields on every product before returning an observation."""
        for p in self._inventory:
            p.steps_until_expiry = max(0, p.expiry_step - self._current_step)
            p.margin_per_unit = round(p.selling_price - p.cost_price, 4)

    def _recompute_pick_probs(self) -> None:
        """Recompute normalised pick probabilities for all active products (zone + popularity)."""
        active = [p for p in self._inventory if p.is_active and p.quantity > 0]
        if not active:
            return
        effs = np.array([
            p.base_probability
            * ZONE_MULTIPLIERS[p.zone]
            * (1.0 + p.current_discount_pct / 100.0)
            * p.popularity_multiplier
            for p in active
        ])
        probs = effs / effs.sum()
        for product, prob in zip(active, probs):
            product.effective_pick_prob = round(float(prob), 6)

    def _generate_inventory(self) -> List[ProductState]:
        """Generate a randomised product inventory for a new episode."""
        chosen_names = self._rng.choice(
            _PRODUCT_NAMES, size=self._num_products, replace=False
        )

        # Draw base probability weights via Dirichlet, then rescale to [0.05, 0.4]
        raw_probs = self._rng.dirichlet(alpha=[1.0] * self._num_products)
        max_raw = raw_probs.max()
        base_probs = 0.05 + raw_probs * (0.40 - 0.05) / max_raw

        # Assign initial zones: fill zone 1 first (up to capacity), then zone 2, rest to zone 3
        zone_slots = {z: ZONE_CAPACITY[z] for z in [1, 2, 3]}
        product_zones: List[int] = []
        for _ in range(self._num_products):
            for zone in [1, 2, 3]:
                if zone_slots[zone] > 0:
                    product_zones.append(zone)
                    zone_slots[zone] -= 1
                    break

        inventory: List[ProductState] = []
        for i in range(self._num_products):
            cost = round(float(self._rng.uniform(1.0, 8.0)), 2)
            margin = float(self._rng.uniform(1.3, 2.5))
            selling = round(cost * margin, 2)
            qty = int(self._rng.integers(15, 41))
            expiry_offset = int(self._rng.integers(self._expiry_offset_min, self._expiry_offset_max + 1))
            holding_rate = round(float(self._rng.uniform(0.01, 0.04)), 4)
            restock_cost_per_unit = round(cost * float(self._rng.uniform(1.0, 1.25)), 2)
            lead_time = int(self._rng.integers(1, 4))  # 1–3 steps

            inventory.append(
                ProductState(
                    product_id=i,
                    name=str(chosen_names[i]),
                    cost_price=cost,
                    selling_price=selling,
                    base_selling_price=selling,
                    quantity=qty,
                    expiry_step=self._current_step + expiry_offset,
                    base_probability=float(base_probs[i]),
                    current_discount_pct=0.0,
                    is_active=True,
                    sales_history=[],
                    zone=product_zones[i],
                    holding_cost_rate=holding_rate,
                    restock_cost_per_unit=restock_cost_per_unit,
                    restock_lead_time=lead_time,
                    pending_restock_quantity=0,
                    pending_restock_arrives_at_step=None,
                    popularity_multiplier=1.0,
                )
            )
        return inventory

    def _update_state(

        self,

        *,

        episode_id: Optional[str] = None,

        last_reward: float = 0.0,

    ) -> None:
        """Sync the internal StoremanagerState with current environment data."""
        if episode_id is not None:
            self._state.episode_id = episode_id
        self._state.step_count = self._current_step
        self._state.num_products = self._num_products
        self._state.max_steps = self._max_steps
        self._state.cumulative_profit = self._cumulative_profit
        self._state.last_reward = last_reward
        self._state.active_product_count = sum(
            1 for p in self._inventory if p.is_active
        )
        self._state.inventory_snapshot = [p.model_dump() for p in self._inventory]

    def _error_observation(self, message: str) -> StoremanagerObservation:
        """Return an observation flagging an invalid action (step does NOT advance)."""
        self._refresh_product_derived_fields()
        return StoremanagerObservation(
            inventory=list(self._inventory),
            current_step=self._current_step,
            max_steps=self._max_steps,
            steps_remaining=max(0, self._max_steps - self._current_step),
            cumulative_profit=self._cumulative_profit,
            last_step_profit=0.0,
            last_step_units_sold=0,
            last_expired_products=[],
            num_customers_per_step=self._num_customers,
            task_name=self._task_name,
            zone_capacity=dict(ZONE_CAPACITY),
            zone_occupancy=self._zone_occupancy(),
            zone_multipliers=dict(ZONE_MULTIPLIERS),
            holding_cost_this_step=0.0,
            placement_cost_this_step=0.0,
            restock_cost_this_step=0.0,
            unjustified_discount_penalty=0.0,
            done=False,
            reward=self._normalize_reward(0.0),
            error=message,
        )


# ── Standalone smoke test ─────────────────────────────────────────────────────

if __name__ == "__main__":
    env = StoremanagerEnvironment(num_products=4, max_steps=6, seed=42)
    obs = env.reset()
    print("=== RESET ===")
    for p in obs.inventory:
        print(f"  [{p.product_id}] {p.name:10s}  zone={p.zone}  qty={p.quantity}  "
              f"lead={p.restock_lead_time}  hold_rate={p.holding_cost_rate:.3f}")
    print(f"  Zone occupancy: {obs.zone_occupancy}")

    actions = [
        StoremanagerAction(action_type="discount", product_id=0, discount_pct=20),
        StoremanagerAction(action_type="restock", product_id=1, restock_quantity=10),
        StoremanagerAction(action_type="place", product_id=2, target_zone=1),
        StoremanagerAction(action_type="place", product_id=3, target_zone=1),  # should fail — zone 1 at capacity
        StoremanagerAction(action_type="discount", product_id=0, discount_pct=10),
        StoremanagerAction(action_type="discount", product_id=1, discount_pct=0),
    ]

    for step_num, action in enumerate(actions):
        obs = env.step(action)
        print(f"\n=== STEP {step_num + 1} ({action.action_type.upper()}) ===")
        if obs.error:
            print(f"  ERROR: {obs.error}")
        else:
            print(f"  Reward:          {obs.reward:.2f}")
            print(f"  Holding cost:    {obs.holding_cost_this_step:.4f}")
            print(f"  Placement cost:  {obs.placement_cost_this_step:.2f}")
            print(f"  Restock cost:    {obs.restock_cost_this_step:.4f}")
            print(f"  Cumulative:      {obs.cumulative_profit:.2f}")
            print(f"  Zone occupancy:  {obs.zone_occupancy}")
            if obs.last_expired_products:
                print(f"  Expired:         {obs.last_expired_products}")
        if obs.done:
            print("  Episode done.")
            break