Spaces:
Sleeping
Sleeping
File size: 1,693 Bytes
1b72fa2 | 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 | import numpy as np
from traffic_rl.env.traffic_env import TrafficEnv
def make_env(**overrides):
config = {
"max_steps": 20,
"arrival_mode": "deterministic",
"arrival_sequence": [
[2, 1, 0, 0],
[1, 0, 2, 0],
[0, 1, 1, 0],
[1, 1, 0, 2],
],
"service_rate": 2,
"seed": 123,
}
config.update(overrides)
return TrafficEnv(config=config)
def test_reset_returns_valid_state():
env = make_env()
state = env.reset()
assert isinstance(state, np.ndarray)
assert state.shape == (10,)
assert np.all(state[:8] >= 0)
assert state[8] in (0, 1)
def test_step_updates_queues_and_non_negative():
env = make_env(arrival_sequence=[[1, 1, 0, 0]], service_rate=1)
env.reset()
next_state, reward, done, info = env.step(1)
assert next_state.shape == (10,)
assert np.all(next_state[:8] >= 0)
assert isinstance(reward, float)
assert isinstance(done, bool)
assert "throughput" in info
def test_environment_deterministic_transitions():
env1 = make_env()
env2 = make_env()
s1 = env1.reset()
s2 = env2.reset()
assert np.allclose(s1, s2)
actions = [1, 0, 2, 0, 1]
for a in actions:
ns1, r1, d1, i1 = env1.step(a)
ns2, r2, d2, i2 = env2.step(a)
assert np.allclose(ns1, ns2)
assert r1 == r2
assert d1 == d2
assert i1["throughput"] == i2["throughput"]
def test_no_negative_values_over_rollout():
env = make_env()
env.reset()
for _ in range(10):
state, _, done, _ = env.step(0)
assert np.all(state[:8] >= 0)
if done:
break
|