Add shared reward test
This commit is contained in:
@@ -1,10 +1,15 @@
|
||||
import yaml
|
||||
|
||||
from primaite.game.agent.rewards import GreenAdminDatabaseUnreachablePenalty, WebpageUnavailablePenalty
|
||||
from primaite.game.game import PrimaiteGame
|
||||
from primaite.session.environment import PrimaiteGymEnv
|
||||
from primaite.simulator.network.hardware.nodes.host.server import Server
|
||||
from primaite.simulator.network.hardware.nodes.network.router import ACLAction, Router
|
||||
from primaite.simulator.network.transmission.network_layer import IPProtocol
|
||||
from primaite.simulator.network.transmission.transport_layer import Port
|
||||
from primaite.simulator.system.applications.database_client import DatabaseClient
|
||||
from primaite.simulator.system.services.database.database_service import DatabaseService
|
||||
from tests import TEST_ASSETS_ROOT
|
||||
from tests.conftest import ControlledAgent
|
||||
|
||||
|
||||
@@ -80,3 +85,25 @@ def test_uc2_rewards(game_and_agent):
|
||||
state = game.get_sim_state()
|
||||
reward_value = comp.calculate(state)
|
||||
assert reward_value == -1.0
|
||||
|
||||
|
||||
def test_shared_reward():
|
||||
CFG_PATH = TEST_ASSETS_ROOT / "configs/shared_rewards.yaml"
|
||||
with open(CFG_PATH, "r") as f:
|
||||
cfg = yaml.safe_load(f)
|
||||
|
||||
env = PrimaiteGymEnv(game_config=cfg)
|
||||
|
||||
env.reset()
|
||||
|
||||
order = env.game._reward_calculation_order
|
||||
assert order.index("defender") > order.index("client_1_green_user")
|
||||
assert order.index("defender") > order.index("client_2_green_user")
|
||||
|
||||
for step in range(256):
|
||||
act = env.action_space.sample()
|
||||
env.step(act)
|
||||
g1_reward = env.game.agents["client_1_green_user"].reward_function.current_reward
|
||||
g2_reward = env.game.agents["client_2_green_user"].reward_function.current_reward
|
||||
blue_reward = env.game.agents["defender"].reward_function.current_reward
|
||||
assert blue_reward == g1_reward + g2_reward
|
||||
|
||||
Reference in New Issue
Block a user