Hugging Face's logo

Spaces:
aakarsh2007
/
Oceanus-AI
Running

App Files Community
Oceanus-AI / oceanus /models.py
here2rule
initial deploy
ed79989
Raw History Blame Contribute Delete
19.7 kB
"""
OpenEnv Wrapper β€” reset(), step(), ObservationSpace, ActionSpace.
"""
import json
import re
from typing import Dict, Any, Tuple, List, Optional
from oceanus.physics import GameState, GRID_SIZE
def parse_asv_action(raw_output: str, schema_version: str = "v1") -> Tuple[Optional[Dict], bool]:
"""Extract JSON action from LLM output. Returns (action_dict, is_valid)."""
if not raw_output or not isinstance(raw_output, str):
return None, False
try:
# Greedy match to handle nested braces
match = re.search(r'\{[^{}]*\}', raw_output, re.DOTALL)
if not match:
return None, False
action = json.loads(match.group())
intent_key = "action_type" if schema_version == "v2" else "intent"
intent = action.get(intent_key, "")
if not intent:
return None, False
valid_intents = {"move", "scan", "clean", "broadcast"}
if intent not in valid_intents:
return None, False
if intent == "move":
direction = action.get("direction", "")
if direction not in {"north", "south", "east", "west", "stay"}:
return None, False
if intent == "broadcast" and not action.get("message", "").strip():
return None, False
# Normalize to v1 for internal engine use
if schema_version == "v2":
action["intent"] = action.pop("action_type")
return action, True
except (json.JSONDecodeError, AttributeError, TypeError):
return None, False
def parse_policy_action(raw_output: str, schema_version: str = "v1") -> Tuple[Optional[Dict], bool]:
"""Extract JSON action from policy LLM output."""
if not raw_output or not isinstance(raw_output, str):
return None, False
try:
match = re.search(r'\{[^{}]*\}', raw_output, re.DOTALL)
if not match:
return None, False
action = json.loads(match.group())
intent_key = "action_type" if schema_version == "v2" else "intent"
target_key = "recipient" if schema_version == "v2" else "target"
intent = action.get(intent_key, "")
if not intent:
return None, False
valid_intents = {"reply_email", "propose_treaty", "accept_treaty", "reject_treaty"}
if intent not in valid_intents:
return None, False
# Relaxed validation β€” only require target for reply/propose/accept
if intent in ("reply_email", "propose_treaty", "accept_treaty"):
if not action.get(target_key, "").strip():
return None, False
# Normalize to v1 for internal engine use
if schema_version == "v2":
action["intent"] = action.pop("action_type")
if target_key in action:
action["target"] = action.pop("recipient")
return action, True
except (json.JSONDecodeError, AttributeError, TypeError):
return None, False
# ─────────────────────────────────────────────
# Observation Builders
# ─────────────────────────────────────────────
def build_asv_observation(state: GameState, asv_id: str) -> Dict:
asv = state.asvs[asv_id]
sonar = state.get_sonar_reading(asv)
sector = state.get_sector(asv.row, asv.col)
on_net = state.grid[asv.row, asv.col] > 0
return {
"agent_id": asv_id,
"position": {"row": asv.row, "col": asv.col, "sector": sector},
"battery": asv.battery,
"sonar_3x3": sonar,
"on_net": bool(on_net),
"comms_inbox": list(asv.comms_inbox[-3:]),
"step": state.step_count,
"biodiversity_index": round(state.biodiversity_index, 1),
"wind_hint": f"Wind direction: {state.wind_vector}",
"schema_version": state.schema_version,
}
def build_policy_observation(state: GameState, agent_id: str) -> Dict:
if agent_id == "Port_Authority":
inbox_msgs = list(state.port_authority_inbox.messages)[:4]
else:
inbox_msgs = list(state.fleet_manager_inbox.messages)[:4]
return {
"agent_id": agent_id,
"current_policy_status": state.treaty_status,
"net_spawn_rate": state.net_spawn_rate,
"biodiversity_index": round(state.biodiversity_index, 1),
"inbox": inbox_msgs,
"step": state.step_count,
"unanswered_emails": (
state.port_authority_inbox.unanswered_count
if agent_id == "Port_Authority"
else state.fleet_manager_inbox.unanswered_count
),
"schema_version": state.schema_version,
}
def build_asv_prompt(obs: Dict) -> str:
sonar_str = "\n".join([" " + str(row) for row in obs["sonar_3x3"]])
inbox_str = "\n".join([f" - {m}" for m in obs["comms_inbox"]]) or " (empty)"
if obs.get("schema_version") == "v2":
schema_instructions = """
Respond with a single JSON action using schema v2 ('action_type' instead of 'intent'):
- To move: {"action_type": "move", "direction": "north|south|east|west|stay"}
- To clean a net at your position: {"action_type": "clean"}
- To broadcast to other ASVs: {"action_type": "broadcast", "message": "your message here"}
- To scan (reveal sonar): {"action_type": "scan"}
"""
else:
schema_instructions = """
Respond with a single JSON action using schema v1:
- To move: {"intent": "move", "direction": "north|south|east|west|stay"}
- To clean a net at your position: {"intent": "clean"}
- To broadcast to other ASVs: {"intent": "broadcast", "message": "your message here"}
- To scan (reveal sonar): {"intent": "scan"}
"""
return f"""You are {obs['agent_id']}, an Autonomous Surface Vehicle on a ghost-net recovery mission.
CURRENT STATE:
- Position: Row {obs['position']['row']}, Col {obs['position']['col']} (Sector {obs['position']['sector']})
- Battery: {obs['battery']}%
- Biodiversity Index: {obs['biodiversity_index']}%
- On Net: {"YES" if obs.get("on_net") else "NO"}
- Step: {obs['step']}
- {obs['wind_hint']}
SONAR (3x3 grid around you, 'net' = ghost net detected):
{sonar_str}
COMMS INBOX:
{inbox_str}
TASK: Recover ghost nets to restore ocean biodiversity. Coordinate with other ASVs.
{schema_instructions}
Your JSON action:"""
def build_policy_prompt(obs: Dict) -> str:
inbox_str = ""
for msg in obs["inbox"]:
inbox_str += f"\n FROM: {msg.get('from','?')} [Mood: {msg.get('mood','Neutral')}]\n MSG: {msg.get('message','')}\n"
if not inbox_str:
inbox_str = " (no messages)"
if obs.get("schema_version") == "v2":
schema_instructions = """
Respond with a single JSON action using schema v2 ('action_type' instead of 'intent', 'recipient' instead of 'target'):
- To reply to an email: {"action_type": "reply_email", "recipient": "sender_name", "content": "your reply"}
- To propose a treaty: {"action_type": "propose_treaty", "recipient": "Fleet_Manager", "content": "treaty terms"}
- To accept a treaty: {"action_type": "accept_treaty", "recipient": "Port_Authority", "content": "acceptance message"}
"""
else:
schema_instructions = """
Respond with a single JSON action using schema v1:
- To reply to an email: {"intent": "reply_email", "target": "sender_name", "content": "your reply"}
- To propose a treaty: {"intent": "propose_treaty", "target": "Fleet_Manager", "content": "treaty terms"}
- To accept a treaty: {"intent": "accept_treaty", "target": "Port_Authority", "content": "acceptance message"}
"""
return f"""You are the {obs['agent_id']} in the Oceanus Ocean Recovery Program.
CURRENT STATE:
- Policy Status: {obs['current_policy_status']}
- Net Spawn Rate: {obs['net_spawn_rate']} (lower is better)
- Biodiversity Index: {obs['biodiversity_index']}%
- Unanswered Emails: {obs['unanswered_emails']}
- Step: {obs['step']}
YOUR INBOX:
{inbox_str}
TASK: Negotiate a Tagging Mandate treaty to reduce ghost net spawning. Reply to stakeholders diplomatically.
A successful treaty reduces net spawn rate from High to Low, saving the ecosystem.
{schema_instructions}
Your JSON action:"""
# ─────────────────────────────────────────────
# Reward Engine
# ─────────────────────────────────────────────
class RewardEngine:
"""CTDE reward: R_total = Ξ±*Ξ£ R_local_asv + Ξ²*Ξ£ R_local_policy + Ξ³*R_global_ecosystem"""
ALPHA = 1.0
BETA = 1.0
GAMMA = 2.0
def __init__(self):
self.episode_rewards: Dict[str, List[float]] = {}
def compute_asv_reward(
self, action: Optional[Dict], valid: bool, cleaned: bool, broadcast_useful: bool
) -> float:
r = 0.0
if not valid:
r -= 1.0
if cleaned:
r += 5.0
if action and action.get("intent") == "move":
r -= 0.1
if broadcast_useful:
r += 0.5
return r
def compute_policy_reward(
self,
action: Optional[Dict],
valid: bool,
treaty_advanced: bool,
email_replied: bool,
email_ignored_penalty: bool,
) -> float:
r = 0.0
if not valid:
r -= 1.0
if treaty_advanced:
r += 10.0
if email_replied:
r += 2.0
if email_ignored_penalty:
r -= 5.0
return r
def compute_global_reward(self, state: GameState) -> float:
if not state.done:
return 0.0
if state.biodiversity_index >= 75.0:
return 100.0
elif state.biodiversity_index >= 50.0:
return 40.0
elif state.biodiversity_index >= 25.0:
return 10.0
return -20.0
def blend(self, asv_rewards: List[float], policy_rewards: List[float], global_r: float) -> float:
return self.ALPHA * sum(asv_rewards) + self.BETA * sum(policy_rewards) + self.GAMMA * global_r
# ─────────────────────────────────────────────
# Main OpenEnv Environment
# ─────────────────────────────────────────────
class OceanusEnv:
"""
OpenEnv-compatible environment for Oceanus.
Observation space: JSON strings (text)
Action space: Text output (parsed to JSON)
"""
metadata = {
"name": "Oceanus-v1",
"description": "Multi-layer ghost-gear recovery and treaty negotiation arena",
"themes": ["multi-agent", "long-horizon", "world-modeling", "self-improvement"],
"num_agents": 6,
}
def __init__(self, seed: int = 42, max_steps: int = 200):
self.seed = seed
self.max_steps = max_steps
self.state: Optional[GameState] = None
self.reward_engine = RewardEngine()
self._episode_asv_rewards: List[float] = []
self._episode_policy_rewards: List[float] = []
self._step_log: List[Dict] = []
def reset(self) -> Dict[str, Any]:
"""Reset environment, return initial observations for all agents."""
self.state = GameState(seed=self.seed)
self.state.max_steps = self.max_steps
self._episode_asv_rewards = []
self._episode_policy_rewards = []
self._step_log = []
obs = {}
for asv_id in self.state.asvs:
raw_obs = build_asv_observation(self.state, asv_id)
obs[asv_id] = {"observation": raw_obs, "prompt": build_asv_prompt(raw_obs)}
for agent_id in ["Port_Authority", "Fleet_Manager"]:
raw_obs = build_policy_observation(self.state, agent_id)
obs[agent_id] = {"observation": raw_obs, "prompt": build_policy_prompt(raw_obs)}
return obs
def step(self, actions: Dict[str, str]) -> Tuple[Dict, Dict, bool, Dict]:
"""
Process one step for all agents.
actions: {agent_id: raw_llm_output_string}
Returns: (observations, rewards, done, info)
"""
assert self.state is not None, "Call reset() first"
rewards: Dict[str, float] = {}
info: Dict[str, Any] = {"events": [], "cleaned": [], "treaty_events": [], "broadcasts": []}
# ── ASV actions ──
asv_rewards = []
for asv_id, raw_output in actions.items():
if not asv_id.startswith("ASV"):
continue
asv = self.state.asvs.get(asv_id)
if asv is None:
continue
action, valid = parse_asv_action(raw_output, schema_version=self.state.schema_version)
cleaned = False
broadcast_useful = False
if valid and action:
intent = action["intent"]
if intent == "move":
asv.move(action.get("direction", "stay"))
elif intent == "clean":
cleaned = self.state.clean_net_at(asv.row, asv.col)
if cleaned:
asv.nets_cleaned += 1
info["cleaned"].append(f"{asv_id} cleaned net at ({asv.row},{asv.col})")
elif intent == "broadcast":
msg = action.get("message", "").strip()
if msg and len(msg) >= 10 and ("net" in msg.lower() or "sector" in msg.lower() or "at" in msg.lower()):
broadcast_useful = True
for other_id, other_asv in self.state.asvs.items():
if other_id != asv_id:
other_asv.comms_inbox.append(f"{asv_id}: {msg}")
if len(other_asv.comms_inbox) > 10:
other_asv.comms_inbox.pop(0)
info["broadcasts"].append(f"{asv_id}: {msg}")
# "scan" β€” sonar is always active, no extra action needed
r = self.reward_engine.compute_asv_reward(action, valid, cleaned, broadcast_useful)
rewards[asv_id] = r
asv_rewards.append(r)
# ── Policy actions ──
policy_rewards = []
for agent_id, raw_output in actions.items():
if agent_id.startswith("ASV"):
continue
if agent_id not in ("Port_Authority", "Fleet_Manager"):
continue
action, valid = parse_policy_action(raw_output, schema_version=self.state.schema_version)
treaty_advanced = False
email_replied = False
email_ignored_penalty = False
if valid and action:
intent = action["intent"]
target = action.get("target", "")
content = action.get("content", "")
if intent == "reply_email":
inbox = (
self.state.port_authority_inbox
if agent_id == "Port_Authority"
else self.state.fleet_manager_inbox
)
msg = inbox.pop_message(target)
if msg and len(content) > 10:
email_replied = True
info["treaty_events"].append(f"{agent_id} replied to {target}")
elif intent == "propose_treaty":
if self.state.treaty_status == "No Tagging Mandate":
self.state.treaty_status = "Treaty Proposed"
treaty_advanced = True
other_inbox = (
self.state.fleet_manager_inbox
if agent_id == "Port_Authority"
else self.state.port_authority_inbox
)
other_inbox.add_message(agent_id, "Formal", f"Treaty proposal: {content}")
info["treaty_events"].append(
f"{agent_id} proposed treaty: {content[:60]}..."
)
elif intent == "accept_treaty":
if self.state.treaty_status == "Treaty Proposed":
self.state.treaty_status = "Tagging Mandate Active"
treaty_advanced = True
info["treaty_events"].append(
f"TREATY SIGNED by {agent_id}! Spawn rate dropping."
)
elif intent == "reject_treaty":
self.state.treaty_status = "No Tagging Mandate"
info["treaty_events"].append(f"{agent_id} rejected treaty.")
# Penalty for ignored emails > 5 steps old
inbox = (
self.state.port_authority_inbox
if agent_id == "Port_Authority"
else self.state.fleet_manager_inbox
)
if any(m["age"] > 5 for m in inbox.messages):
email_ignored_penalty = True
r = self.reward_engine.compute_policy_reward(
action, valid, treaty_advanced, email_replied, email_ignored_penalty
)
rewards[agent_id] = r
policy_rewards.append(r)
# ── Tick physics ──
self.state.tick()
# ── Global reward at episode end ──
global_r = self.reward_engine.compute_global_reward(self.state)
total_r = self.reward_engine.blend(asv_rewards, policy_rewards, global_r)
rewards["__total__"] = total_r
rewards["__global__"] = global_r
self._episode_asv_rewards.extend(asv_rewards)
self._episode_policy_rewards.extend(policy_rewards)
# ── Build next observations ──
obs: Dict[str, Any] = {}
for asv_id in self.state.asvs:
raw_obs = build_asv_observation(self.state, asv_id)
obs[asv_id] = {"observation": raw_obs, "prompt": build_asv_prompt(raw_obs)}
for agent_id in ["Port_Authority", "Fleet_Manager"]:
raw_obs = build_policy_observation(self.state, agent_id)
obs[agent_id] = {"observation": raw_obs, "prompt": build_policy_prompt(raw_obs)}
info["step"] = self.state.step_count
info["biodiversity"] = self.state.biodiversity_index
info["active_nets"] = len(self.state.ghost_nets)
info["treaty_status"] = self.state.treaty_status
info["total_cleaned"] = self.state.total_cleaned
self._step_log.append({
"step": self.state.step_count,
"rewards": dict(rewards),
"biodiversity": self.state.biodiversity_index,
"active_nets": len(self.state.ghost_nets),
"treaty": self.state.treaty_status,
})
return obs, rewards, self.state.done, info
def get_episode_summary(self) -> Dict:
return {
"total_steps": self.state.step_count if self.state else 0,
"total_cleaned": self.state.total_cleaned if self.state else 0,
"biodiversity_final": self.state.biodiversity_index if self.state else 0.0,
"treaty_status": self.state.treaty_status if self.state else "N/A",
"avg_asv_reward": (
sum(self._episode_asv_rewards) / max(1, len(self._episode_asv_rewards))
),
"avg_policy_reward": (
sum(self._episode_policy_rewards) / max(1, len(self._episode_policy_rewards))
),
"step_log": self._step_log,
}