YUS200619's picture
Upload environment/env.py with huggingface_hub
11b949a verified
from __future__ import annotations
import random
from typing import Dict, Optional, Tuple
try:
from openenv import Environment
except ImportError: # local fallback for dev without openenv installed
class Environment: # type: ignore
pass
from .consequences import apply_action
from .events import CANONICAL_EVENTS, Event
from .reward import annotate_decision, component_scores, weighted_total
from .state import EpisodeState
class WorkLifeFirewallEnv(Environment):
metadata = {"name": "work-life-firewall", "version": "1.0.0", "max_steps": 14}
def __init__(self, randomize_order: bool = True, seed: Optional[int] = None):
super().__init__()
self.randomize_order = randomize_order
self.rng = random.Random(seed)
self._state = EpisodeState()
self._events = []
self._idx = 0
def reset(self) -> Dict:
self._state = EpisodeState()
self._events = list(CANONICAL_EVENTS)
if self.randomize_order:
opener = self._events[0]
tail = self._events[1:]
self.rng.shuffle(tail)
self._events = [opener] + tail
self._idx = 0
return self._build_obs(self._events[self._idx])
def step(self, action: str) -> Tuple[Dict, float, bool, Dict]:
event = self._events[self._idx]
action_id, action_data = self._decode_action(action, event)
self._state = apply_action(self._state, event.id, action_id, action_data)
# Add comm quality annotations into latest decision.
latest = self._state.decisions[-1]
annotate_decision(latest, action)
step_reward = self._step_reward(latest, action_data)
self._idx += 1
self._state.day = min(4, self._idx // 2)
done = self._idx >= len(self._events) or self._state.energy <= 0.05
if done:
comps = component_scores(self._state, self._state.decisions)
terminal = weighted_total(comps)
info = {"episode_complete": True, "components": comps}
return self._state.to_observation(), round(step_reward + terminal, 4), True, info
return self._build_obs(self._events[self._idx]), round(step_reward, 4), False, {}
def state(self) -> Dict:
return self._state.to_observation()
def _build_obs(self, event: Event) -> Dict:
return {
"state": self._state.to_observation(),
"event": {
"id": event.id,
"title": event.title,
"source": event.source,
"sender": event.sender,
"body": event.body,
"urgency": event.urgency,
"relationship_stake": event.relationship_stake,
"actions": [a["id"] for a in event.actions],
},
"instruction": "Write the response or action you would take. End with a specific next step.",
}
def _decode_action(self, free_text: str, event: Event) -> Tuple[str, Dict]:
text = free_text.lower()
decode_map = {
"E1_staging": {"ssh": "fix_directly", "restart": "fix_directly", "delegate": "delegate_oncall", "infra": "escalate_infra"},
"E2_slack": {"async": "async_boundary", "later": "async_boundary", "call now": "both_calls_now", "ignore": "ignore_slack"},
"E3_client_email": {"thanks": "acknowledge_timeline", "thursday": "acknowledge_timeline", "sorry": "full_apology", "cc": "cc_manager"},
"E4_leave": {"approve": "ping_sundar", "today": "ping_sundar", "wait": "wait_hope", "cancel": "cancel_leave"},
"E5_appraisal": {"block": "block_wednesday", "calendar": "block_wednesday", "tonight": "rush_tonight", "extension": "ask_extension"},
"E6_oncall": {"can't": "no_clearly_kindly", "cannot": "no_clearly_kindly", "swap": "counter_offer", "yes": "yes_people_pleaser", "sure": "yes_people_pleaser"},
"E7_standup": {"async": "decline_async", "skip": "decline_async", "attend": "attend_sacrifice", "delegate": "delegate_teammate"},
}
mapping = decode_map.get(event.id, {})
selected = event.actions[0]
for keyword, action_id in mapping.items():
if keyword in text:
selected = next((a for a in event.actions if a["id"] == action_id), event.actions[0])
break
return selected["id"], selected
def _step_reward(self, decision: Dict, action_data: Dict) -> float:
energy = float(action_data.get("energy_delta", 0.0))
sprint = float(action_data.get("sprint_delta", 0.0))
rel_penalty = sum(abs(v) for v in action_data.get("relationship_delta", {}).values()) * 0.08
comm_bonus = decision.get("message_quality", 0.0) * 0.12
return energy * 0.5 + sprint * 0.4 - rel_penalty + comm_bonus