Spaces:
Running
Running
File size: 2,860 Bytes
e11e4fe |
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 |
from urllib.parse import urlparse, parse_qs
def _behavior_to_agent_id(behavior_name: str, unique_id: int) -> str:
return f"{behavior_name}?agent_id={unique_id}"
def _agent_id_to_behavior(agent_id: str) -> str:
return agent_id.split("?agent_id=")[0]
def _unwrap_batch_steps(batch_steps, behavior_name):
decision_batch, termination_batch = batch_steps
decision_id = [
_behavior_to_agent_id(behavior_name, i) for i in decision_batch.agent_id
]
termination_id = [
_behavior_to_agent_id(behavior_name, i) for i in termination_batch.agent_id
]
agents = decision_id + termination_id
obs = {
agent_id: [batch_obs[i] for batch_obs in termination_batch.obs]
for i, agent_id in enumerate(termination_id)
}
if decision_batch.action_mask is not None:
obs.update(
{
agent_id: {
"observation": [batch_obs[i] for batch_obs in decision_batch.obs],
"action_mask": [mask[i] for mask in decision_batch.action_mask],
}
for i, agent_id in enumerate(decision_id)
}
)
else:
obs.update(
{
agent_id: [batch_obs[i] for batch_obs in decision_batch.obs]
for i, agent_id in enumerate(decision_id)
}
)
obs = {k: v if len(v) > 1 else v[0] for k, v in obs.items()}
dones = {agent_id: True for agent_id in termination_id}
dones.update({agent_id: False for agent_id in decision_id})
rewards = {
agent_id: termination_batch.reward[i]
for i, agent_id in enumerate(termination_id)
}
rewards.update(
{agent_id: decision_batch.reward[i] for i, agent_id in enumerate(decision_id)}
)
cumulative_rewards = {k: v for k, v in rewards.items()}
infos = {}
for i, agent_id in enumerate(decision_id):
infos[agent_id] = {}
infos[agent_id]["behavior_name"] = behavior_name
infos[agent_id]["group_id"] = decision_batch.group_id[i]
infos[agent_id]["group_reward"] = decision_batch.group_reward[i]
for i, agent_id in enumerate(termination_id):
infos[agent_id] = {}
infos[agent_id]["behavior_name"] = behavior_name
infos[agent_id]["group_id"] = termination_batch.group_id[i]
infos[agent_id]["group_reward"] = termination_batch.group_reward[i]
infos[agent_id]["interrupted"] = termination_batch.interrupted[i]
id_map = {agent_id: i for i, agent_id in enumerate(decision_id)}
return agents, obs, dones, rewards, cumulative_rewards, infos, id_map
def _parse_behavior(full_behavior):
parsed = urlparse(full_behavior)
name = parsed.path
ids = parse_qs(parsed.query)
team_id: int = 0
if "team" in ids:
team_id = int(ids["team"][0])
return name, team_id
|