# Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # BSD-style license """ Organization Simulation Environment. A multi-agent environment simulating hierarchical organization with teams, escalation, and resource management. """ import uuid from typing import Any from openenv.core.env_server.interfaces import Environment from ..models import OrgAction, OrgObservation, OrgState TASK_SCENARIOS = { "solo_bug_fix": { "description": "Fix a high-priority bug in the engineering team", "tasks": [ { "id": "task_bug", "team": "engineering", "type": "fix_bug", "priority": "high", "deadline": 10, "difficulty": 1, } ], }, "cross_team_launch": { "description": "Engineering and Sales coordinate a product launch", "tasks": [ { "id": "task_eng", "team": "engineering", "type": "implement_feature", "priority": "high", "deadline": 15, "difficulty": 2, }, { "id": "task_sales", "team": "sales", "type": "prepare_proposal", "priority": "medium", "deadline": 12, "difficulty": 2, }, ], }, "startup_crisis": { "description": "Simultaneous incident + resource-gated feature + cross-team client meeting", "tasks": [ { "id": "task_incident", "team": "engineering", "type": "fix_bug", "priority": "critical", "deadline": 5, "difficulty": 3, }, { "id": "task_feature", "team": "engineering", "type": "implement_feature", "priority": "high", "deadline": 20, "difficulty": 2, "requires_resource": "senior_engineer", }, { "id": "task_client", "team": "sales", "type": "client_meeting", "priority": "critical", "deadline": 8, "difficulty": 2, }, ], }, } TASK_IDS = list(TASK_SCENARIOS.keys()) class OrgSimEnvironment(Environment): """Organization simulation environment.""" def __init__(self, max_steps: int = 50): super().__init__() self.max_steps = max_steps self._current_task_id = "cross_team_launch" self._state = OrgState(episode_id="", step_count=0) self._current_agent = "eng_swe1" self._tasks: dict[str, dict] = {} self._messages: list[dict] = [] self._episode_resources: dict[str, bool] = {} self._init_organization() def _init_organization(self): """Initialize static organization structure.""" self.teams = { "engineering": { "lead": "eng_lead", "members": ["eng_swe1", "eng_swe2", "eng_swe3"], }, "sales": { "lead": "sales_mgr", "members": ["sales_rep1", "sales_rep2", "sales_rep3"], }, "operations": { "lead": "ops_lead", "members": ["ops_analyst1", "ops_analyst2", "ops_analyst3"], }, } def reset(self, agent_id: str = "eng_swe1", task_id: str = "cross_team_launch", **kwargs: Any) -> OrgObservation: """Reset environment for new episode. Args: agent_id: The agent playing this episode. task_id: One of 'solo_bug_fix', 'cross_team_launch', 'startup_crisis'. """ self._current_task_id = task_id if task_id in TASK_SCENARIOS else "cross_team_launch" self._state.episode_id = str(uuid.uuid4()) self._state.step_count = 0 self._current_agent = agent_id self._messages = [] scenario = TASK_SCENARIOS[self._current_task_id] self._tasks = { t["id"]: { **t, "status": "pending", "assignee": None, "progress": 0.0, } for t in scenario["tasks"] } # Reset shared resources per episode self._episode_resources = { "senior_engineer": True, "cloud_instances": True, "budget_approval": True, } return self._get_observation() def step(self, action: OrgAction) -> OrgObservation: """Execute action and return observation.""" self._state.step_count += 1 reward = 0.0 if action.action_type == "REQUEST_TASK": reward = self._handle_request_task(action) elif action.action_type == "ACCEPT_TASK": reward = self._handle_accept_task(action) elif action.action_type == "COMPLETE_TASK": reward = self._handle_complete_task(action) elif action.action_type == "REQUEST_HELP": reward = self._handle_request_help(action) elif action.action_type == "PROVIDE_HELP": reward = self._handle_provide_help(action) elif action.action_type == "ESCALATE": reward = self._handle_escalate(action) elif action.action_type == "REQUEST_RESOURCE": reward = self._handle_request_resource(action) elif action.action_type == "REPORT_STATUS": reward = 0.0 done = self._check_done() return self._get_observation(reward=reward, done=done) def _handle_request_task(self, action: OrgAction) -> float: my_team = self._get_team_for_agent(self._current_agent) available = [ t for t in self._tasks.values() if t["team"] == my_team and t["status"] == "pending" and t.get("assignee") is None ] if available: task = available[0] task["status"] = "in_progress" task["assignee"] = self._current_agent return 0.1 return -0.05 # No task available — slight penalty to discourage looping def _handle_accept_task(self, action: OrgAction) -> float: task_id = action.payload.get("task_id") if not task_id or task_id not in self._tasks: return -0.1 task = self._tasks[task_id] if task["status"] != "pending": return -0.1 task["status"] = "in_progress" task["assignee"] = self._current_agent return 0.1 def _handle_complete_task(self, action: OrgAction) -> float: task_id = action.payload.get("task_id") if not task_id or task_id not in self._tasks: return -0.2 task = self._tasks[task_id] if task.get("assignee") != self._current_agent: return -0.1 if task["status"] != "in_progress": return -0.1 # Block completion of resource-gated tasks if resource was never locked required = task.get("requires_resource") if required and self._episode_resources.get(required, True): # Resource still available (never locked) — penalize and don't complete return -0.2 deadline = task.get("deadline", 20) time_taken = self._state.step_count if time_taken <= deadline: base_reward = 1.0 time_bonus = max(0.0, (deadline - time_taken) / deadline) * 0.3 priority_scores = {"critical": 0.3, "high": 0.2, "medium": 0.1, "low": 0.0} priority_bonus = priority_scores.get(task.get("priority", "medium"), 0.1) task["status"] = "completed" task["completed_at"] = self._state.step_count return base_reward + time_bonus + priority_bonus else: task["status"] = "failed" return -0.3 def _handle_request_help(self, action: OrgAction) -> float: task_id = action.payload.get("task_id") if not task_id or task_id not in self._tasks: return -0.1 task = self._tasks[task_id] if task.get("assignee") != self._current_agent: return -0.1 task["progress"] = min(1.0, task.get("progress", 0.0) + 0.2) lead = self._get_team_lead(self._get_team_for_agent(self._current_agent)) self._messages.append({ "from": self._current_agent, "to": lead, "type": "help_request", "task_id": task_id, }) # Scripted reply gives agent real information to act on self._messages.append({ "from": lead, "to": self._current_agent, "type": "help_response", "task_id": task_id, "content": f"Acknowledged. Assigned eng_swe2 to assist with {task_id}. Check internal docs for context.", }) return 0.1 def _handle_provide_help(self, action: OrgAction) -> float: task_id = action.payload.get("task_id") if not task_id or task_id not in self._tasks: return -0.1 task = self._tasks[task_id] task["progress"] = min(1.0, task.get("progress", 0.0) + 0.3) return 0.2 def _handle_escalate(self, action: OrgAction) -> float: task_id = action.payload.get("task_id") if not task_id or task_id not in self._tasks: return -0.2 task = self._tasks[task_id] if task["status"] in ("completed", "failed", "escalated"): return -0.1 # Already terminal # Appropriate: cross-team task OR difficulty >= 2 when stuck (progress < 0.5) my_team = self._get_team_for_agent(self._current_agent) is_cross_team = task["team"] != my_team is_hard_stuck = task.get("difficulty", 1) >= 2 and task.get("progress", 0.0) < 0.5 if is_cross_team or is_hard_stuck: task["status"] = "escalated" task["escalated_at"] = self._state.step_count return 0.3 else: return -0.2 # Premature / unnecessary escalation def _handle_request_resource(self, action: OrgAction) -> float: resource_id = action.payload.get("resource_id") if not resource_id or resource_id not in self._episode_resources: return -0.2 if self._episode_resources[resource_id]: self._episode_resources[resource_id] = False # Lock it return 0.2 return -0.1 # Already locked def _check_done(self) -> bool: """Episode ends when all tasks are in a terminal state OR max_steps reached. Checks for terminal states (completed/failed/escalated), NOT just non-pending. A task in 'in_progress' is NOT done — episode must continue. """ if self._state.step_count >= self.max_steps: return True active = sum( 1 for t in self._tasks.values() if t["status"] in ("pending", "in_progress") ) return active == 0 def grade_episode(self) -> float: """Grade the completed episode. Returns 0.0–1.0. Called after done=True. Score varies with timing, escalation quality, and resource management — never identical across different trajectories. """ if self._current_task_id == "solo_bug_fix": return self._grade_solo_bug_fix() elif self._current_task_id == "cross_team_launch": return self._grade_cross_team_launch() elif self._current_task_id == "startup_crisis": return self._grade_startup_crisis() return 0.0 def _grade_solo_bug_fix(self) -> float: task = self._tasks.get("task_bug", {}) if task.get("status") == "completed": completed_at = task.get("completed_at", self._state.step_count) deadline = task.get("deadline", 10) time_bonus = max(0.0, (deadline - completed_at) / deadline) * 0.30 return min(1.0, 0.50 + time_bonus + 0.15) # base + time_bonus + priority elif task.get("status") == "failed": return 0.15 # partial credit for attempting return 0.0 def _grade_cross_team_launch(self) -> float: score = 0.0 eng = self._tasks.get("task_eng", {}) sales = self._tasks.get("task_sales", {}) if eng.get("status") == "completed": completed_at = eng.get("completed_at", self._state.step_count) deadline = eng.get("deadline", 15) time_bonus = max(0.0, (deadline - completed_at) / deadline) * 0.20 score += 0.40 + time_bonus # up to 0.60 if sales.get("status") == "escalated": score += 0.25 # correct action for cross-team task elif sales.get("status") == "completed": score += 0.20 # completed somehow — less expected but fine return min(1.0, score) def _grade_startup_crisis(self) -> float: score = 0.0 incident = self._tasks.get("task_incident", {}) feature = self._tasks.get("task_feature", {}) client = self._tasks.get("task_client", {}) # Incident: must complete by step 5 if incident.get("status") == "completed": completed_at = incident.get("completed_at", self._state.step_count) if completed_at <= 5: score += 0.40 else: score += 0.10 # completed but late # Feature: resource must have been locked (False = locked) resource_locked = not self._episode_resources.get("senior_engineer", True) if feature.get("status") == "completed" and resource_locked: score += 0.30 elif feature.get("status") == "completed": score += 0.10 # completed without resource # Client: correct cross-team escalation if client.get("status") == "escalated": score += 0.20 elif client.get("status") == "completed": score += 0.10 # Resource management bonus if resource_locked: score += 0.10 return min(1.0, score) def _get_observation(self, reward: float = 0.0, done: bool = False) -> OrgObservation: my_team = self._get_team_for_agent(self._current_agent) available = [ t for t in self._tasks.values() if t["team"] == my_team and t["status"] == "pending" and t.get("assignee") is None ] active = next( (t for t in self._tasks.values() if t.get("assignee") == self._current_agent and t["status"] == "in_progress"), None, ) team_status = { team_name: { "busy": sum(1 for t in self._tasks.values() if t["team"] == team_name and t.get("assignee")), "total": len(team_data["members"]), } for team_name, team_data in self.teams.items() } return OrgObservation( my_agent_id=self._current_agent, my_team=my_team, my_role="lead" if self._current_agent == self.teams[my_team]["lead"] else "member", available_tasks=available, active_task=active, inbox=[m for m in self._messages if m.get("to") == self._current_agent], team_status=team_status, resources=self._episode_resources.copy(), metrics={ "tasks_completed": sum(1 for t in self._tasks.values() if t["status"] == "completed"), "tasks_failed": sum(1 for t in self._tasks.values() if t["status"] == "failed"), "tasks_escalated": sum(1 for t in self._tasks.values() if t["status"] == "escalated"), "current_task_id": self._current_task_id, "step_count": self._state.step_count, }, done=done, reward=reward, metadata={"step_count": self._state.step_count, "task_id": self._current_task_id}, ) def _get_team_for_agent(self, agent_id: str) -> str: """Get team for agent — checks both lead and members.""" for team_name, team_data in self.teams.items(): if agent_id == team_data["lead"] or agent_id in team_data["members"]: return team_name return "engineering" def _get_team_lead(self, team_name: str) -> str: return self.teams.get(team_name, {}).get("lead", "") @property def state(self) -> OrgState: self._state.tasks = self._tasks.copy() self._state.team_members = {k: v["members"] for k, v in self.teams.items()} self._state.resource_status = self._episode_resources.copy() return self._state