Upload environment/consequences.py with huggingface_hub

environment/consequences.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+from dataclasses import replace
+import random
+from typing import Dict
+
+from .state import EpisodeState, RelationshipState, clamp01
+
+
+def _copy_state(state: EpisodeState) -> EpisodeState:
+    return replace(
+        state,
+        relationships=replace(state.relationships),
+        decisions=[*state.decisions],
+        history=[*state.history],
+    )
+
+
+def apply_action(state: EpisodeState, event_id: str, action_id: str, action_data: Dict) -> EpisodeState:
+    new_state = _copy_state(state)
+
+    new_state.energy = clamp01(state.energy + float(action_data.get("energy_delta", 0.0)))
+    new_state.sprint_health = clamp01(state.sprint_health + float(action_data.get("sprint_delta", 0.0)))
+
+    rel = new_state.relationships
+    for person, delta in action_data.get("relationship_delta", {}).items():
+        if hasattr(rel, person):
+            setattr(rel, person, clamp01(getattr(rel, person) + float(delta)))
+    new_state.relationships = RelationshipState(**vars(rel)).clamped()
+
+    if event_id == "E1_staging" and action_id == "fix_directly":
+        new_state.staging_fixed = True
+    if event_id == "E4_leave":
+        if action_id == "ping_sundar":
+            new_state.leave_status = "approved" if random.random() > 0.20 else "pending"
+        elif action_id == "cancel_leave":
+            new_state.leave_status = "cancelled"
+    if event_id == "E5_appraisal" and action_id == "block_wednesday":
+        new_state.appraisal_done = True
+    if event_id == "E6_oncall" and action_id == "yes_people_pleaser":
+        new_state.oncall_accepted = True
+
+    new_state.decisions.append(
+        {
+            "day": state.day,
+            "event_id": event_id,
+            "action_id": action_id,
+            "energy_before": state.energy,
+            "energy_after": new_state.energy,
+        }
+    )
+
+    new_state.history.append(f"{event_id}:{action_id}")
+    return new_state.normalized()