diff --git a/ltc/run.py b/ltc/run.py index e3ec9a8..c500554 100644 --- a/ltc/run.py +++ b/ltc/run.py @@ -1,4 +1,5 @@ import os + os.environ['XLA_PYTHON_CLIENT_PREALLOCATE'] = 'false' from dataclasses import replace @@ -13,7 +14,7 @@ from ltc.agents import BayesianDDQN, DCF, QNetwork, QNetworkDropout, StochasticVariationalNetwork from ltc.sim import InitialStateConf, cox_traffic, process_output, simulate -from ltc.sim.constants import INITIAL_CAPACITY, Actions +from ltc.sim.constants import INITIAL_CAPACITY, Actions, ObservationIndex from ltc.utils.scan_states import Carry, Output from ltc.utils.plots import plot_all, plot_first @@ -51,7 +52,7 @@ def init_traffic(traffic, key, n): def rl_step(drl_step, dcf_step, traffic_step, n, n_drl, n_bins=50): def rl_step_fn(c, _): - key, drl_keys, dcf_keys, traffic_key = jax.random.split(c.key, 4) + key, drl_keys, dcf_keys, traffic_key, simulation_key = jax.random.split(c.key, 5) drl_keys = jax.random.split(drl_keys, n_drl) dcf_keys = jax.random.split(dcf_keys, n - n_drl) traffic_keys = jax.random.split(traffic_key, n) @@ -61,8 +62,10 @@ def rl_step_fn(c, _): actions = jnp.concatenate([drl_actions, dcf_actions]) traffic_states, new_frames = traffic_step(c.traffic_states, traffic_keys) - buffer_states, channel_state = simulate(c.buffer_states, new_frames, actions) - obs, rewards, powers = process_output(c.buffer_states, buffer_states, c.power_states, channel_state, c.obs, actions, c.terminals) + buffer_states, channel_state, transmission_history = simulate(simulation_key, c.buffer_states, new_frames, + actions, c.transmission_history, n) + obs, rewards, powers = process_output(c.buffer_states, buffer_states, c.power_states, channel_state, c.obs, + actions, c.terminals, transmission_history) terminals = jnp.logical_or(c.terminals, powers < 0) params = c.drl_states.params['model'] if 'model' in c.drl_states.params else c.drl_states.params @@ -72,7 +75,7 @@ def rl_step_fn(c, _): hist, bin_edges = jax.vmap(jnp.histogram, in_axes=(0, None))(flat_params, n_bins) c = Carry( - drl_states, dcf_states, traffic_states, buffer_states, powers, + drl_states, dcf_states, traffic_states, buffer_states, powers, transmission_history, channel_state, key, obs, actions, rewards, terminals ) o = Output( @@ -95,7 +98,10 @@ def rl_step_fn(c, _): buffer_states = jnp.zeros(n, dtype=int) power_states = jnp.full(n, INITIAL_CAPACITY, dtype=int) channel_state = 0 - obs = jnp.zeros((n, window_size, 5), dtype=int).at[:, -1].set(INITIAL_CAPACITY) + obs = jnp.zeros((n, window_size, 6), dtype=int).at[:, ObservationIndex.POWER.value].set(INITIAL_CAPACITY) + id_values = jnp.arange(0, n).reshape(n, 1) + obs = obs.at[:, :, ObservationIndex.ID.value].set(id_values) + transmission_history = jnp.zeros((window_size + 1, 2)) rewards = jnp.zeros(n) terminals = jnp.full(n, False, dtype=bool) @@ -127,7 +133,7 @@ def rl_step_fn(c, _): rl_step_fn = jax.jit(rl_step(drl_step, dcf_step, traffic_step, n, n_drl)) init_carry = Carry( - drl_states, dcf_states, traffic_states, buffer_states, power_states, + drl_states, dcf_states, traffic_states, buffer_states, power_states, transmission_history, channel_state, key, obs, actions, rewards, terminals ) all_outputs = [] diff --git a/ltc/sim/constants.py b/ltc/sim/constants.py index 864db46..ff21baa 100644 --- a/ltc/sim/constants.py +++ b/ltc/sim/constants.py @@ -2,12 +2,10 @@ class Actions(Enum): - """ - Actions for the agent - """ TX = 0 CS = 1 IDLE = 2 + ACK = 3 """ @@ -21,10 +19,12 @@ class Actions(Enum): Rewards and penalties """ TX_REWARD = 1.0 +ACK_REWARD = 1.0 EMPTY_BUFFER_REWARD = 0.5 NO_TX_REWARD = 0.0 NO_TX_PENALTY = -1.0 EMPTY_TX_PENALTY = -0.5 +EMPTY_ACK_PENALTY = -0.5 COLLISION_PENALTY = -1.0 MAX_RETRANSMISSION_PENALTY = -1.0 @@ -33,6 +33,7 @@ class Actions(Enum): """ INITIAL_CAPACITY = 10000 TX_CONSUMPTION = 5 +ACK_CONSUMPTION = 2 CS_CONSUMPTION = 1 IDLE_CONSUMPTION = 0 @@ -41,3 +42,22 @@ class Actions(Enum): https://ieeexplore.ieee.org/document/8930559 """ TAU = 5.484 * 1e-3 + + +class AckState(Enum): + SENT = 1 + NOT_SENT = 0 + + +class ObservationIndex(Enum): + BUFFER = 0 + CHANNEL = 1 + RET_C = 2 + NO_TX = 3 + POWER = 4 + ID = 5 + + +class TransmissionIndex(Enum): + DESTINATION = 0 + ACK = 1 diff --git a/ltc/sim/process_output.py b/ltc/sim/process_output.py index 9b4d8da..b5bc4a8 100644 --- a/ltc/sim/process_output.py +++ b/ltc/sim/process_output.py @@ -5,7 +5,7 @@ def no_transmission(args): - action, buffer_state, _, _, no_tx = args + action, buffer_state, _, _, no_tx, _ = args return jax.lax.cond( action == Actions.IDLE.value, lambda: jax.lax.cond(buffer_state == 0, idle_empty_buffer, idle_full_buffer, args), @@ -21,21 +21,21 @@ def idle_empty_buffer(_): def idle_full_buffer(args): - _, _, ret_c, _, no_tx = args + _, _, ret_c, _, no_tx, _ = args reward = NO_TX_REWARD no_tx = no_tx + 1 return reward, ret_c, no_tx def no_transmission_short(args): - _, buffer_state, ret_c, _, no_tx = args + _, buffer_state, ret_c, _, no_tx, _ = args reward = NO_TX_REWARD no_tx = jnp.where(buffer_state == 0, 0, no_tx + 1) return reward, ret_c, no_tx def no_transmission_long(args): - _, _, ret_c, _, no_tx = args + _, _, ret_c, _, no_tx, _ = args scale = jax.lax.min(1., (no_tx - SAFE_IDLE_PERIOD + 1) / PENALIZED_IDLE_PERIOD) reward = scale * NO_TX_PENALTY no_tx = no_tx + 1 @@ -43,12 +43,12 @@ def no_transmission_long(args): def transmission(args): - _, _, _, channel_state, _ = args + _, _, _, channel_state, _, _ = args return jax.lax.cond(channel_state == 1, transmission_without_collision, transmission_with_collision, args) def transmission_with_collision(args): - _, _, ret_c, _, _ = args + _, _, ret_c, _, _, _ = args return jax.lax.cond(ret_c < MAX_RETRANSMISSION, retransmission, max_retransmission_collision, args) @@ -60,7 +60,7 @@ def max_retransmission_collision(_): def retransmission(args): - _, _, ret_c, _, _ = args + _, _, ret_c, _, _, _ = args reward = COLLISION_PENALTY ret_c = ret_c + 1 no_tx = 0 @@ -68,7 +68,7 @@ def retransmission(args): def transmission_without_collision(args): - _, buffer_state, _, _, _ = args + _, buffer_state, _, _, _, _ = args return jax.lax.cond(buffer_state > 0, successful_transmission, empty_buffer_transmission, args) @@ -80,18 +80,61 @@ def empty_buffer_transmission(_): def successful_transmission(args): - _, _, ret_c, _, _ = args + _, _, ret_c, _, _, _ = args reward = TX_REWARD / (ret_c + 1) ret_c = 0 no_tx = 0 return reward, ret_c, no_tx -def process_output_i(buffer_state, new_buffer_state, power_state, channel_state, obs, action, terminal): - _, _, ret_c, no_tx, _ = obs[-1] - args = (action, buffer_state, ret_c, channel_state, no_tx) +def transmission_ack(args): + _, _, _, channel_state, _, _ = args + return jax.lax.cond(channel_state == 1, transmission_ack_without_collision, transmission_with_collision, args) + + +def transmission_ack_without_collision(args): + _, _, _, _, _, ack_flag = args + return jax.lax.cond(ack_flag == 1, successful_ack_transmission, empty_ack_transmission, args) + + +def empty_ack_transmission(args): + _, _, ret_c, _, no_tx, _ = args + reward = EMPTY_ACK_PENALTY + ret_c = ret_c + no_tx = no_tx + 1 + return reward, ret_c, no_tx + + +def successful_ack_transmission(args): + _, _, ret_c, _, _, _ = args + reward = ACK_REWARD / (ret_c + 1) + ret_c = 0 + no_tx = 0 + return reward, ret_c, no_tx + + +def process_output_i(buffer_state, new_buffer_state, power_state, channel_state, obs, action, terminal, + transmission_history): + _, _, ret_c, no_tx, _, my_id = obs[-1] + mask = (transmission_history[:, TransmissionIndex.DESTINATION.value] == my_id) & ( + transmission_history[:, TransmissionIndex.ACK.value] == AckState.NOT_SENT.value) + matching_rows = jnp.nonzero(mask, size=mask.shape[0], fill_value=-1)[0] + print(matching_rows) + ack_flag = jnp.where(jnp.all(matching_rows == -1), -1, 1) + args = (action, buffer_state, ret_c, channel_state, no_tx, ack_flag) + + reward, ret_c, no_tx = jax.lax.cond( + action == Actions.ACK.value, + transmission_ack, + lambda a: jax.lax.cond( + action == Actions.TX.value, + transmission, + no_transmission, + operand=a + ), + operand=args + ) - reward, ret_c, no_tx = jax.lax.cond(action == Actions.TX.value, transmission, no_transmission, args) reward = jnp.where(terminal, 0., reward) channel_state = jnp.where(action == Actions.CS.value, channel_state, -1) @@ -101,18 +144,24 @@ def process_output_i(buffer_state, new_buffer_state, power_state, channel_state, action == Actions.CS.value, power_state - CS_CONSUMPTION, jnp.where( action == Actions.IDLE.value, power_state - IDLE_CONSUMPTION, - power_state + jnp.where( + action == Actions.ACK.value, power_state - ACK_CONSUMPTION, + power_state + ) ) ) ) - obs_t = jnp.array([new_buffer_state, channel_state, ret_c, no_tx, power]) + obs_t = jnp.array([new_buffer_state, channel_state, ret_c, no_tx, power, my_id]) obs = jnp.roll(obs, -1, axis=0) obs = obs.at[-1].set(obs_t) return obs, reward, power -def process_output(buffer_states, new_buffer_states, power_states, channel_state, obs, actions, terminals): +def process_output(buffer_states, new_buffer_states, power_states, channel_state, obs, actions, terminals, transmission_history): channel_states = jnp.full(buffer_states.shape[0], channel_state) - return jax.vmap(process_output_i)(buffer_states, new_buffer_states, power_states, channel_states, obs, actions, terminals) + return jax.vmap(process_output_i, in_axes=(0, 0, 0, 0, 0, 0, 0, None))(buffer_states, new_buffer_states, + power_states, channel_states, obs, actions, + terminals, transmission_history + ) diff --git a/ltc/sim/sim.py b/ltc/sim/sim.py index 369c854..7e514a2 100644 --- a/ltc/sim/sim.py +++ b/ltc/sim/sim.py @@ -1,6 +1,7 @@ import jax.numpy as jnp +import jax -from ltc.sim.constants import Actions +from ltc.sim.constants import * def channel_state_selector(actions): @@ -11,10 +12,12 @@ def channel_state_selector(actions): int: -1 if more than one STA transmit, 1 if exactly one STA transmit, 0 if noone transmit at the moment. """ - ones_count = jnp.sum(actions == Actions.TX.value) + transmitting = (actions == Actions.TX.value) | (actions == Actions.ACK.value) + active_count = jnp.sum(transmitting) + return jnp.where( - ones_count > 1, -1, - jnp.where(ones_count == 1, 1, 0) + active_count > 1, -1, + jnp.where(active_count == 1, 1, 0) ) @@ -46,8 +49,52 @@ def add_new_frames(buffer_states, new_frames): return jnp.bitwise_or(buffer_states, new_frames) -def simulate(buffer_states, new_frames, actions): +def modify_transmission_history(actions, transmission_history, address): + args = (actions, transmission_history, address) + return jax.lax.cond( + jnp.any(actions == Actions.ACK.value), + ack_transmission, + lambda a: jax.lax.cond( + jnp.any(actions == Actions.TX.value), + data_transmission, + no_transmission, + operand=a + ), + operand=args + ) + + +def data_transmission(args): + actions, transmission_history, address = args + transmission_history_t = jnp.array([address, AckState.NOT_SENT.value]) + transmission_history = jnp.roll(transmission_history, -1, axis=0) + transmission_history = transmission_history.at[-1].set(transmission_history_t) + return transmission_history + + +def ack_transmission(args): + actions, transmission_history, address = args + sta_num = jnp.argmax(actions == Actions.ACK.value) + mask = (transmission_history[:, TransmissionIndex.DESTINATION.value] == sta_num) & ( + transmission_history[:, TransmissionIndex.ACK.value] == AckState.NOT_SENT.value) + matching_rows = jnp.nonzero(mask, size=mask.shape[0], fill_value=-1)[0] + transmission_history = transmission_history.at[matching_rows[0], TransmissionIndex.ACK.value].set( + AckState.SENT.value) + return transmission_history + + +def no_transmission(args): + actions, transmission_history, address = args + return transmission_history + + +def simulate(key, buffer_states, new_frames, actions, transmission_history, nWifi): channel_state = channel_state_selector(actions) + address = jnp.where(channel_state == 1, jax.random.randint(key, shape=(), minval=0, maxval=nWifi), -1) + transmission_history = jnp.where(channel_state == 1, + modify_transmission_history(actions, transmission_history, address), + transmission_history) buffer_states = jnp.where(channel_state == 1, buffer_clearing(buffer_states, actions), buffer_states) buffer_states = add_new_frames(buffer_states, new_frames) - return buffer_states, channel_state + + return buffer_states, channel_state, transmission_history diff --git a/ltc/utils/scan_states.py b/ltc/utils/scan_states.py index 6342ae8..3ae115a 100644 --- a/ltc/utils/scan_states.py +++ b/ltc/utils/scan_states.py @@ -14,6 +14,7 @@ class Carry: traffic_states: ModelState buffer_states: jax.Array power_states: jax.Array + transmission_history: jax.Array channel_state: int key: jax.random.PRNGKey obs: jax.Array diff --git a/test/sim/test_process_output.py b/test/sim/test_process_output.py index a72e231..ca8be43 100644 --- a/test/sim/test_process_output.py +++ b/test/sim/test_process_output.py @@ -8,81 +8,93 @@ class ProcessOutputTestCase(unittest.TestCase): def test_no_transmission(self): - args = (1, 5, 0, 1) + args = (1, 1, 5, 1, 0, 0) reward, r, no_tx = no_transmission(args) self.assertEqual(reward, 0.0) self.assertEqual(r, 5) def test_transmission_without_collision_empty_buffer(self): - args = (0, 3, 1, 1) + args = (1, 0, 0, 1, 0, 0) reward, r, no_tx = transmission_without_collision(args) self.assertEqual(reward, EMPTY_TX_PENALTY) self.assertEqual(r, 0) def test_transmission_without_collision_successful(self): - args = (1, 2, 1, 1) + args = (1, 1, 2, 1, 0, 0) reward, r, no_tx = transmission_without_collision(args) self.assertAlmostEqual(reward, TX_REWARD / (2 + 1)) self.assertEqual(r, 0) def test_transmission_with_collision_retransmission(self): - args = (1, 2, -1, 1) + args = (1, 1, 2, -1, 0, 0) reward, r, no_tx = transmission_with_collision(args) self.assertEqual(reward, COLLISION_PENALTY) self.assertEqual(r, 3) def test_transmission_with_collision_max_retransmission(self): - args = (1, MAX_RETRANSMISSION, -1, 1) + args = (1, 1, MAX_RETRANSMISSION, -1, 0, 0) reward, r, no_tx = transmission_with_collision(args) self.assertEqual(reward, MAX_RETRANSMISSION_PENALTY) self.assertEqual(r, 0) def test_transmission_collision_path(self): - args = (1, 7, -1, 1) + args = (1, 1, 7, -1, 0, 0) reward, r, no_tx = transmission(args) self.assertEqual(reward, COLLISION_PENALTY) self.assertEqual(r, 8) def test_transmission_successful_path(self): - args = (1, 1, 1, 1) + args = (1, 1, 1, 1, 0, 0) reward, r, no_tx = transmission(args) self.assertAlmostEqual(reward, TX_REWARD / (1 + 1)) self.assertEqual(r, 0) def test_transmission_empty_buffer(self): - args = (0, 1, 1, 1) + args = (1, 0, 1, 1, 0, 0) reward, r, no_tx = transmission(args) self.assertEqual(reward, EMPTY_TX_PENALTY) self.assertEqual(r, 0) def test_process_rl_output(self): - buffer_states = jnp.array([0, 0, 1, 0]) - new_buffer_states = jnp.array([0, 0, 1, 0]) + buffer_states = jnp.array([1, 1, 1, 1]) + new_buffer_states = jnp.array([0, 1, 1, 0]) power_states = jnp.array([100, 100, 100, 100]) - actions = jnp.array([Actions.TX.value, Actions.CS.value, Actions.TX.value, Actions.TX.value]) + actions = jnp.array([Actions.ACK.value, Actions.IDLE.value, Actions.TX.value, Actions.TX.value]) terminals = jnp.array([False, False, False, False]) + n = 4 + window_size = 4 channel_state = -1 obs = jnp.array([[ - [1, 1, 6, 0, 100], - [1, 1, 7, 0, 100], - [1, 1, 8, 0, 100] - ]] * 4) + [1, 1, 6, 0, 100, 0], + [1, 1, 7, 0, 100, 0], + [1, 1, 8, 0, 100, 0] + ]] * n) + id_values = jnp.arange(0, n).reshape(n, 1) + obs = obs.at[:, :, ObservationIndex.ID.value].set(id_values) expected_obs_0 = jnp.array([ - [1, 1, 7, 0, 100], - [1, 1, 8, 0, 100], - [0, -1, 0, 0, 100 - TX_CONSUMPTION], + [1, 1, 7, 0, 100, 0], + [1, 1, 8, 0, 100, 0], + [0, -1, 0, 0, 100 - ACK_CONSUMPTION, 0], ]) expected_R_0 = -1.0 - expected_power = power_states + jnp.array([-TX_CONSUMPTION, -CS_CONSUMPTION, -TX_CONSUMPTION, -TX_CONSUMPTION]) + expected_power = power_states + jnp.array([-ACK_CONSUMPTION, -IDLE_CONSUMPTION, -TX_CONSUMPTION, -TX_CONSUMPTION]) + + transmission_history = jnp.array([[0, 1], + [0, 0], + [0, 1], + [3, 0], + [0, 1]]) result_obs, result_R, power = process_output( - buffer_states, new_buffer_states, power_states, channel_state, obs, actions, terminals + buffer_states, new_buffer_states, power_states, channel_state, obs, actions, terminals, transmission_history ) + self.assertTrue(jnp.array_equal(result_obs[0], expected_obs_0)) self.assertEqual(result_R[0], expected_R_0) self.assertTrue(jnp.array_equal(power, expected_power)) + if __name__ == '__main__': unittest.main() diff --git a/test/sim/test_sim.py b/test/sim/test_sim.py index 065caef..120c645 100644 --- a/test/sim/test_sim.py +++ b/test/sim/test_sim.py @@ -58,34 +58,25 @@ def test_simulate(self): expected_buffer_states = jnp.array([0, 1, 1, 1]) expected_channel_state = 1 - - result_buffer_states, result_channel_state = simulate(buffer_states, new_frames, actions) - self.assertTrue(jnp.array_equal(result_buffer_states, expected_buffer_states)) - self.assertEqual(result_channel_state, expected_channel_state) - - # Test simulation with no actions - buffer_states = jnp.array([1, 0, 1, 0]) - new_frames = jnp.array([0, 1, 0, 1]) - actions = jnp.array([Actions.CS.value, Actions.CS.value, Actions.CS.value, Actions.CS.value]) - - expected_buffer_states = jnp.array([1, 1, 1, 1]) - expected_channel_state = 0 - - result_buffer_states, result_channel_state = simulate(buffer_states, new_frames, actions) - self.assertTrue(jnp.array_equal(result_buffer_states, expected_buffer_states)) - self.assertEqual(result_channel_state, expected_channel_state) - - # Test simulation with multiple actions - buffer_states = jnp.array([0, 0, 0, 0]) - new_frames = jnp.array([0, 1, 0, 1]) - actions = jnp.array([Actions.TX.value, Actions.CS.value, Actions.CS.value, Actions.TX.value]) - - expected_buffer_states = jnp.array([0, 1, 0, 1]) - expected_channel_state = -1 - - result_buffer_states, result_channel_state = simulate(buffer_states, new_frames, actions) + nWifi = 4 + seed_number = 1 + + transmission_history = jnp.array([[0, 1], + [0, 1], + [0, 1], + [3, 0], + [0, 1]]) + + transmission_history_expected = jnp.array([[0, 1], + [0, 1], + [3, 0], + [0, 1], + [2, 0]]) + + result_buffer_states, result_channel_state, transmission_history_result = simulate(seed_number, buffer_states, new_frames, actions, transmission_history, nWifi) self.assertTrue(jnp.array_equal(result_buffer_states, expected_buffer_states)) self.assertEqual(result_channel_state, expected_channel_state) + self.assertTrue(jnp.array_equal(transmission_history_result, transmission_history_expected)) if __name__ == "__main__":