CrisisWorldCortex / manual_testing.py
Angshuman28's picture
Upload folder using huggingface_hub
d2cb56e verified
Raw
History Blame Contribute Delete
1.72 kB
from .models import NoOp, DeployResource, RestrictMovement, CrisisworldcortexAction, CrisisworldcortexObservation, PublicCommunication
from .server import CrisisWorld as CrisisWorldEnvironment
def logger(count, obs: CrisisworldcortexObservation, action: CrisisworldcortexAction | None) -> None:
if action:
print(f"[Action Taken]: {action.model_dump(exclude_none=True)}")
else:
print("Initial Observation:")
print(f"turn={count}, done={obs.done}, reward={obs.reward}")
for reg in obs.regions:
print(
f"{reg.region_id} pop={reg.population} "
f"inf={reg.infected} rec={reg.recovered} "
f"dec={reg.deceased}"
)
print("-" * 40)
print()
env = CrisisWorldEnvironment(config=EnvConfig(max_turns=6, num_regions=4))
obs = env.reset(seed=42, episode_id="manual-test")
logger(0, obs, None)
# Turn 1 (NoOp)
action = DeployResource(resource="medical", region_id="r0", amount=10)
# action = NoOp()
obs = env.step(action)
logger(1, obs, action)
# Turn 2 (NoOp)
action = NoOp()
obs = env.step(action)
logger(2, obs, action)
# Turn 3 (DeployResource to r0)
action = RestrictMovement(region_id="r0", level=2)
obs = env.step(action)
logger(3, obs, action)
# Turn 4 (PublicCommunication)
action = PublicCommunication(audience="public", message="Stay safe! Wear masks and maintain social distancing. Don't move between regions, stay home if you can, and seek medical care if you have symptoms.")
obs = env.step(action)
logger(4, obs, action)
# Turn 5 (NoOp)
action = DeployResource(resource="medical", region_id="r0", amount=40)
# action = NoOp()
obs = env.step(action)
logger(5, obs, action)