env/generator.py

import random
from env.tasks import get_task


def generate_env(task_id, rng=None):
    """Build initial items and hidden dynamics. Use *rng* for reproducible episodes."""
    r = rng if rng is not None else random
    config = get_task(task_id)

    state_items = []
    hidden_items = []

    current_time = 0

    for i in range(config["num_items"]):
        base_ttl = r.randint(3, 8)
        update_freq = r.uniform(0.1, config["volatility"])
        last_update = r.randint(0, 3)

        age = current_time - last_update

        is_stale = age > base_ttl or r.random() < update_freq

        last_result = "stale" if is_stale and r.random() < 0.7 else "hit"

        state_items.append(
            {
                "key": f"item_{i}",
                "age": max(age, 0),
                "access_count": r.randint(1, 20),
                "last_result": last_result,
            }
        )

        hidden_items.append(
            {
                "base_ttl": base_ttl,
                "update_freq": update_freq,
                "last_update": last_update,
            }
        )

    return state_items, hidden_items, current_time