Spaces:
Sleeping
Sleeping
File size: 3,212 Bytes
dfea88e | 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 | from pydantic import BaseModel, Field
from typing import List, Dict, Any, Tuple
from .simulator import CacheSimulator
from .workloads import generate_easy_task, generate_medium_task, generate_hard_task
class Observation(BaseModel):
incoming_request: int = Field(description="The ID of the data item being requested.")
cache_state: List[int] = Field(description="Current items in the cache. -1 means empty.")
idle_times: List[int] = Field(description="Time steps since each cache slot was last accessed.")
class Action(BaseModel):
evict_index: int = Field(description="The index (0 to capacity-1) of the cache slot to evict.")
class AdaptiveCacheEnv:
def __init__(self, task_level: str = "easy", capacity: int = 10):
self.capacity = capacity
self.task_level = task_level
self.sim = CacheSimulator(capacity)
if task_level == "easy":
self.workload = generate_easy_task()
elif task_level == "medium":
self.workload = generate_medium_task(cache_size=capacity)
else:
self.workload = generate_hard_task()
self.step_count = 0
self.hits = 0
def reset(self) -> Observation:
self.sim = CacheSimulator(self.capacity)
self.step_count = 0
self.hits = 0
return self.state()
def state(self) -> Observation:
# Safe check for the terminal state to prevent IndexError
if self.step_count >= len(self.workload):
current_item = -1 # Simulation is over, no more incoming requests
else:
current_item = self.workload[self.step_count]
idle_times = [(self.sim.current_time - t) if t > 0 else 0 for t in self.sim.last_access_time]
return Observation(
incoming_request=current_item,
cache_state=self.sim.cache.tolist(),
idle_times=idle_times
)
def step(self, action: Action) -> Tuple[Observation, float, bool, Dict[str, Any]]:
# 1. Apply Action (Evict and Insert)
current_item = self.workload[self.step_count]
self.sim.evict_and_insert(action.evict_index, current_item)
# 2. Advance time strictly by 1 step
self.step_count += 1
# 3. Check Episode Boundary
done = self.step_count >= len(self.workload)
reward = 0.0
if done:
final_score = self.hits / max(1, len(self.workload))
return self.state(), reward, True, {"score": final_score}
# 4. Evaluate the *next* state strictly without fast-forwarding
next_item = self.workload[self.step_count]
is_hit = self.sim.request_item(next_item)
if is_hit:
reward = 1.0
self.hits += 1
# If it's a hit, the agent will see this in the next observation
# and can essentially choose a "safe" eviction slot that doesn't hurt.
else:
reward = -1.0
current_score = self.hits / max(1, self.step_count)
info = {"score": current_score, "hits": self.hits, "steps": self.step_count}
return self.state(), reward, done, info |