From a2f9a2ea347a5ec624b4d0fd8d3c0082f34e209d Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Tue, 21 Apr 2026 12:28:31 +0200 Subject: [PATCH 01/24] Make buffer non-binary --- .gitignore | 2 + ltc/agents/dcf.py | 8 +-- ltc/run.py | 17 +++-- ltc/sim/constants.py | 3 + ltc/sim/process_output.py | 8 ++- ltc/sim/sim.py | 91 ++++++++++++++++++-------- ltc/utils/plots.py | 22 ++++--- ltc/utils/scan_states.py | 1 + test/sim/test_process_output.py | 31 ++++++--- test/sim/test_sim.py | 110 +++++++++++++++----------------- 10 files changed, 175 insertions(+), 118 deletions(-) diff --git a/.gitignore b/.gitignore index d86e4ff..8269ccd 100644 --- a/.gitignore +++ b/.gitignore @@ -169,3 +169,5 @@ cython_debug/ # PyPI configuration file .pypirc + +scripts/ \ No newline at end of file diff --git a/ltc/agents/dcf.py b/ltc/agents/dcf.py index 7231f2d..a76fd7a 100644 --- a/ltc/agents/dcf.py +++ b/ltc/agents/dcf.py @@ -44,10 +44,10 @@ def double_cw(): backoff = jax.random.randint(key, (), 0, state.cw) return DCFState(cw=cw, backoff=backoff) - buffer, channel, ret_c, _, _ = env_state[-1] + buffer_count, channel, ret_c, _, _ = env_state[-1] return jax.lax.cond( - buffer == 0, + buffer_count == 0, reset, lambda: jax.lax.cond( jax.lax.bitwise_and(state.backoff > 0, channel == 0), @@ -70,10 +70,10 @@ def double_cw(): @staticmethod def sample(state, key, env_state): - buffer, channel, _, _, _ = env_state[-1] + buffer_count, channel, _, _, _ = env_state[-1] return jnp.where( - buffer == 0, + buffer_count == 0, Actions.IDLE.value, jnp.where( jax.lax.bitwise_and(state.backoff == 0, channel == 0), diff --git a/ltc/run.py b/ltc/run.py index e16da42..d4bac97 100644 --- a/ltc/run.py +++ b/ltc/run.py @@ -14,7 +14,7 @@ from ltc.agents import BayesianDDQN, DCF, QNetwork, 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, EMPTY_PACKET_ID from ltc.utils.history import build_history_filename, ensure_clean_git_worktree, get_short_commit_hash from ltc.utils.scan_states import Carry, Output from ltc.utils.plots import plot_all, plot_first @@ -128,7 +128,7 @@ def rl_step_coroutine(c, step): actions = jnp.concatenate([drl_actions, legacy_actions]) traffic_states, new_frames = traffic_step(c.traffic_states, traffic_keys) - buffer_states, channel_state = simulate(c.buffer_states, new_frames, actions) + buffer_states, channel_state, packet_seqs = simulate(c.buffer_states, new_frames, actions, c.packet_seqs) obs, rewards, powers = process_output( c.buffer_states, buffer_states, c.power_states, channel_state, c.obs, actions, c.terminals, reward_key ) @@ -144,12 +144,12 @@ def rl_step_coroutine(c, step): hist, bin_edges = None, None c = Carry( - drl_states, legacy_states, traffic_states, buffer_states, powers, + drl_states, legacy_states, traffic_states, packet_seqs, buffer_states, powers, channel_state, key, obs, actions, rewards, terminals, active ) o = Output( legacy_states, obs, actions, rewards, terminals, buffer_states, powers, - (new_frames > 0).astype(int), channel_state, active, hist, bin_edges + new_frames, channel_state, active, hist, bin_edges ) yield c, o @@ -179,6 +179,7 @@ def setup_args(): parser.add_argument('--n_epochs', type=int, default=50, help='Number of training epochs to run.') parser.add_argument('--n_steps', type=int, default=2000, help='Number of steps per epoch.') parser.add_argument('--window_size', type=int, default=1, help='Size of the observation window for each agent.') + parser.add_argument('--queue_size', type=int, default=16, help='Fixed packet-queue length for each station.') parser.add_argument('--seed', type=int, default=42, help='Random seed for reproducibility.') parser.add_argument('--save_plots', action='store_true', default=False, help='Whether to save the generated plots.') parser.add_argument('--loc', type=float, default=5.0, help='loc traffic generator parameter.') @@ -207,6 +208,7 @@ def setup_args(): n_steps = args.n_steps total_steps = n_epochs * n_steps window_size = args.window_size + queue_size = args.queue_size seed = args.seed traffic_type = args.traffic_type @@ -239,10 +241,11 @@ def setup_args(): key = jax.random.key(seed) num_actions = 2 actions = jnp.zeros(n, dtype=int) - buffer_states = jnp.zeros(n, dtype=int) + packet_seqs = jnp.zeros(n, dtype=jnp.int32) + buffer_states = jnp.full((n, queue_size), EMPTY_PACKET_ID, dtype=jnp.int32) power_states = jnp.full(n, INITIAL_CAPACITY, dtype=int) channel_state = 0 - obs = jnp.zeros((n, window_size, 5), dtype=int) + obs = jnp.zeros((n, window_size, 5), dtype=jnp.int32) rewards = jnp.zeros(n) terminals = jnp.full(n, False, dtype=bool) active = jnp.ones(n, dtype=bool).at[n_init:].set(False) @@ -306,7 +309,7 @@ def setup_args(): ) rl_step_fn = jax.jit(rl_step_fn) carry = Carry( - drl_states, legacy_states, traffic_states, buffer_states, power_states, + drl_states, legacy_states, traffic_states, packet_seqs, buffer_states, power_states, channel_state, key, obs, actions, rewards, terminals, active ) all_outputs = [] diff --git a/ltc/sim/constants.py b/ltc/sim/constants.py index 203828e..cd87fe5 100644 --- a/ltc/sim/constants.py +++ b/ltc/sim/constants.py @@ -42,3 +42,6 @@ class Actions(Enum): https://ieeexplore.ieee.org/document/8930559 """ TAU = 5.484 * 1e-3 + +# Queue representation +EMPTY_PACKET_ID = -1 diff --git a/ltc/sim/process_output.py b/ltc/sim/process_output.py index ee55f3e..43d0780 100644 --- a/ltc/sim/process_output.py +++ b/ltc/sim/process_output.py @@ -2,11 +2,12 @@ import jax.numpy as jnp from ltc.sim.constants import * +from ltc.sim.sim import is_buffer_empty def no_transmission(args): _, buffer_state, _, _, _, _ = args - return jax.lax.cond(buffer_state == 0, idle_empty_buffer, idle_full_buffer, args) + return jax.lax.cond(is_buffer_empty(buffer_state), idle_empty_buffer, idle_full_buffer, args) def idle_empty_buffer(_): @@ -67,7 +68,7 @@ def retransmission(args): def transmission_without_collision(args): _, buffer_state, _, _, _, _ = args - return jax.lax.cond(buffer_state > 0, successful_transmission, empty_buffer_transmission, args) + return jax.lax.cond(~is_buffer_empty(buffer_state), successful_transmission, empty_buffer_transmission, args) def empty_buffer_transmission(_): @@ -104,7 +105,8 @@ def process_output_i(buffer_state, new_buffer_state, power_state, channel_state, ) ) - obs_t = jnp.array([new_buffer_state, channel_state, ret_c, no_tx, action]) + buffer_count = jnp.sum((new_buffer_state != EMPTY_PACKET_ID).astype(jnp.int32)) + obs_t = jnp.array([buffer_count, channel_state, ret_c, no_tx, action], dtype=jnp.int32) obs = jnp.roll(obs, -1, axis=0) obs = obs.at[-1].set(obs_t) diff --git a/ltc/sim/sim.py b/ltc/sim/sim.py index 369c854..e9919a0 100644 --- a/ltc/sim/sim.py +++ b/ltc/sim/sim.py @@ -1,6 +1,7 @@ +import jax import jax.numpy as jnp -from ltc.sim.constants import Actions +from ltc.sim.constants import Actions, EMPTY_PACKET_ID def channel_state_selector(actions): @@ -18,36 +19,76 @@ def channel_state_selector(actions): ) -def buffer_clearing(buffer_states, actions): +def _cantor_pairing(a, b): + s = a + b + return (s * (s + 1)) // 2 + b + + +def is_buffer_empty(buffer_state): + return jnp.all(buffer_state == EMPTY_PACKET_ID) + + +def remove_packets_from_queue(queue_state, remove_mask): """ - Updates the buffer_states vector based on the action vector. - Sets the value to 0 in buffer_states where there is a 1 in both vectors. - :param buffer_states: vector with binary value of STAs' buffer occupation. - 0 - STA's buffer is empty - 1 - STA's buffer is full - :param actions: binary vector describes STAs' actions - 0 - channel sensing - 1 - transmission - :return: - jnp.ndarray: updated buffer_states. + Removes packets selected by remove_mask and compacts remaining packets to queue head. """ - return jnp.where((buffer_states == 1) & (actions == Actions.TX.value), 0, buffer_states) + n = queue_state.shape[0] + valid = queue_state != EMPTY_PACKET_ID + remove_mask = jnp.asarray(remove_mask, dtype=bool) + keep = valid & ~remove_mask + # Stable compaction using argsort over static-size keys. + idx = jnp.arange(n) + sort_key = jnp.where(keep, idx, n + idx) + order = jnp.argsort(sort_key) + compacted = queue_state[order] + kept_count = jnp.sum(keep.astype(jnp.int32)) + tail_mask = jnp.arange(n) >= kept_count + return jnp.where(tail_mask, EMPTY_PACKET_ID, compacted) -def add_new_frames(buffer_states, new_frames): + +def remove_transmitted_packets(buffer_states, tx_masks): + return jax.vmap(remove_packets_from_queue)(buffer_states, tx_masks) + + +def enqueue_generated_packets(buffer_state, new_frames, station_id, packet_seq): """ - Updates the buffer_states by adding new frames from generator. - :param buffer_states: vector with binary value of STAs' buffer occupation. - :param new_frames: vector with new frames generated for each STA. - :return: - jnp.ndarray: updated buffer_states. + Appends generated packets at queue tail up to free capacity. """ - new_frames = (new_frames > 0).astype(int) - return jnp.bitwise_or(buffer_states, new_frames) + q = buffer_state.shape[0] + new_frames = jnp.maximum(new_frames.astype(jnp.int32), 0) + occupied = jnp.sum((buffer_state != EMPTY_PACKET_ID).astype(jnp.int32)) + capacity = q - occupied + to_add = jnp.minimum(new_frames.astype(jnp.int32), capacity) + idx = jnp.arange(q, dtype=jnp.int32) + rel = idx - occupied + add_mask = (idx >= occupied) & (rel < to_add) -def simulate(buffer_states, new_frames, actions): + local_ids = packet_seq + jnp.maximum(rel.astype(jnp.int32), 0) + packet_ids = _cantor_pairing(station_id.astype(jnp.int32), local_ids) + next_state = jnp.where(add_mask, packet_ids, buffer_state) + next_packet_seq = packet_seq + new_frames + + return next_state, next_packet_seq + + +def add_new_frames(buffer_states, new_frames, packet_seqs): + station_ids = jnp.arange(buffer_states.shape[0], dtype=jnp.int32) + return jax.vmap(enqueue_generated_packets)(buffer_states, new_frames, station_ids, packet_seqs) + + +def simulate(buffer_states, new_frames, actions, packet_seqs): + """ + One simulation step with fixed-size packet queues. + On successful TX, one packet from queue head is removed for transmitting station. + """ channel_state = channel_state_selector(actions) - 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 + success_tx = (actions == Actions.TX.value) & (channel_state == 1) + queue_size = buffer_states.shape[1] + head_mask = jnp.arange(queue_size) == 0 + tx_masks = success_tx[:, None] & head_mask[None, :] + + buffer_states = remove_transmitted_packets(buffer_states, tx_masks) + buffer_states, packet_seqs = add_new_frames(buffer_states, new_frames, packet_seqs) + return buffer_states, channel_state, packet_seqs diff --git a/ltc/utils/plots.py b/ltc/utils/plots.py index 913c689..9731885 100644 --- a/ltc/utils/plots.py +++ b/ltc/utils/plots.py @@ -10,7 +10,7 @@ from ltc.utils.scan_states import Output from ltc.utils.history import parse_history_filename, unpack_history -from ltc.sim.constants import NO_TX_REWARD, Actions, INITIAL_CAPACITY +from ltc.sim.constants import NO_TX_REWARD, Actions, INITIAL_CAPACITY, EMPTY_PACKET_ID SAVE_FORMAT = 'normal' @@ -304,14 +304,15 @@ def plot_buffer_states(buffer_states, terminals, n, n_drl, seed, name): else: xs = np.linspace(0, n_epochs * n_steps, n_epochs) + 1 - buffer_states = np.where(1 - terminals, buffer_states, 0).sum(axis=1) - buffer_states = buffer_states / (1 - terminals).sum(axis=1) + queue_fill = np.count_nonzero(buffer_states != EMPTY_PACKET_ID, axis=-1) / buffer_states.shape[-1] + queue_fill = np.where(1 - terminals, queue_fill, 0).sum(axis=1) + queue_fill = queue_fill / (1 - terminals).sum(axis=1) for i in range(n_drl): - plt.plot(xs, buffer_states[:, i], color='red') + plt.plot(xs, queue_fill[:, i], color='red') for i in range(n_drl, n): - plt.plot(xs, buffer_states[:, i], color='blue', linestyle='--') + plt.plot(xs, queue_fill[:, i], color='blue', linestyle='--') plt.plot([], color='blue', linestyle='--', label='Legacy') plt.plot([], color='red', label='DRL') @@ -507,8 +508,9 @@ def plot_channel_access_delay(buffer_states, new_frames, terminals, n, n_drl, se else: xs = np.linspace(0, n_epochs * n_steps, n_epochs) + 1 - buffer_states, new_frames = (1 - terminals) * buffer_states, (1 - terminals) * new_frames - cad = buffer_states.sum(axis=1) / (new_frames.sum(axis=1) + 1e-6) + queue_fill = np.count_nonzero(buffer_states != EMPTY_PACKET_ID, axis=-1) + queue_fill, new_frames = (1 - terminals) * queue_fill, (1 - terminals) * new_frames + cad = queue_fill.sum(axis=1) / (new_frames.sum(axis=1) + 1e-6) for i in range(n_drl): plt.plot(xs, cad[:, i], color='red') @@ -539,13 +541,15 @@ def plot_xnor(actions, channel_states, buffer_states, terminals, n, n_drl, seed, else: xs = np.linspace(0, n_epochs * n_steps, n_epochs) + 1 + empty = np.all(buffer_states == EMPTY_PACKET_ID, axis=-1) + non_empty = ~empty idle_and_empty = jax.lax.bitwise_and( jax.lax.bitwise_not(terminals), - jax.lax.bitwise_and(actions == Actions.IDLE.value, buffer_states == 0) + jax.lax.bitwise_and(actions == Actions.IDLE.value, empty) ).sum(axis=1) tx_and_full = jax.lax.bitwise_and( jax.lax.bitwise_and(jax.lax.bitwise_not(terminals), np.expand_dims(channel_states, axis=-1) == 1), - jax.lax.bitwise_and(actions == Actions.TX.value, buffer_states != 0) + jax.lax.bitwise_and(actions == Actions.TX.value, non_empty) ).sum(axis=1) xnor = idle_and_empty + tx_and_full diff --git a/ltc/utils/scan_states.py b/ltc/utils/scan_states.py index 51bdf93..2009f92 100644 --- a/ltc/utils/scan_states.py +++ b/ltc/utils/scan_states.py @@ -12,6 +12,7 @@ class Carry: drl_states: AgentState legacy_states: AgentState traffic_states: ModelState + packet_seqs: jax.Array buffer_states: jax.Array power_states: jax.Array channel_state: int diff --git a/test/sim/test_process_output.py b/test/sim/test_process_output.py index 963cc4f..039e25c 100644 --- a/test/sim/test_process_output.py +++ b/test/sim/test_process_output.py @@ -7,60 +7,71 @@ from ltc.sim.process_output import * KEY = jax.random.PRNGKey(0) +EMPTY = EMPTY_PACKET_ID class ProcessOutputTestCase(unittest.TestCase): def test_no_transmission(self): - args = (Actions.IDLE.value, 1, 5, 0, 1, KEY) + args = (Actions.IDLE.value, jnp.array([1, EMPTY, EMPTY]), 5, 0, 1, KEY) 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 = (Actions.TX.value, 0, 3, 1, 1, KEY) + args = (Actions.TX.value, jnp.array([EMPTY, EMPTY, EMPTY]), 3, 1, 1, KEY) 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 = (Actions.TX.value, 1, 2, 1, 1, KEY) + args = (Actions.TX.value, jnp.array([5, EMPTY, EMPTY]), 2, 1, 1, KEY) reward, r, no_tx = transmission_without_collision(args) self.assertAlmostEqual(reward, TX_REWARD) self.assertEqual(r, 0) def test_transmission_with_collision_retransmission(self): - args = (Actions.TX.value, 1, 2, -1, 1, KEY) + args = (Actions.TX.value, jnp.array([5, EMPTY, EMPTY]), 2, -1, 1, KEY) 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 = (Actions.TX.value, 1, MAX_RETRANSMISSION, -1, 1, KEY) + args = (Actions.TX.value, jnp.array([5, EMPTY, EMPTY]), MAX_RETRANSMISSION, -1, 1, KEY) 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 = (Actions.TX.value, 1, 7, -1, 1, KEY) + args = (Actions.TX.value, jnp.array([5, EMPTY, EMPTY]), 7, -1, 1, KEY) reward, r, no_tx = transmission(args) self.assertEqual(reward, COLLISION_PENALTY) self.assertEqual(r, 8) def test_transmission_successful_path(self): - args = (Actions.TX.value, 1, 1, 1, 1, KEY) + args = (Actions.TX.value, jnp.array([5, EMPTY, EMPTY]), 1, 1, 1, KEY) reward, r, no_tx = transmission(args) self.assertAlmostEqual(reward, TX_REWARD) self.assertEqual(r, 0) def test_transmission_empty_buffer(self): - args = (Actions.TX.value, 0, 1, 1, 1, KEY) + args = (Actions.TX.value, jnp.array([EMPTY, EMPTY, EMPTY]), 1, 1, 1, KEY) 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([ + [EMPTY, EMPTY, EMPTY], + [EMPTY, EMPTY, EMPTY], + [9, EMPTY, EMPTY], + [EMPTY, EMPTY, EMPTY], + ]) + new_buffer_states = jnp.array([ + [EMPTY, EMPTY, EMPTY], + [EMPTY, EMPTY, EMPTY], + [9, EMPTY, EMPTY], + [EMPTY, EMPTY, EMPTY], + ]) power_states = jnp.array([100, 100, 100, 100]) actions = jnp.array([Actions.TX.value, Actions.CS.value, Actions.TX.value, Actions.TX.value]) terminals = jnp.array([False, False, False, False]) diff --git a/test/sim/test_sim.py b/test/sim/test_sim.py index 065caef..829bcc1 100644 --- a/test/sim/test_sim.py +++ b/test/sim/test_sim.py @@ -2,11 +2,13 @@ import jax.numpy as jnp -from ltc.sim.constants import Actions +from ltc.sim.constants import Actions, EMPTY_PACKET_ID from ltc.sim.sim import * class SimulateTestCase(unittest.TestCase): + QUEUE = 4 + def test_channel_state_selector(self): # Test no actions actions = jnp.array([Actions.CS.value, Actions.CS.value, Actions.CS.value, Actions.CS.value]) @@ -20,72 +22,60 @@ def test_channel_state_selector(self): actions = jnp.array([Actions.TX.value, Actions.TX.value, Actions.CS.value, Actions.CS.value]) self.assertEqual(channel_state_selector(actions), -1) - def test_buffer_clearing(self): - # Test clearing buffers when actions match - buffer_states = jnp.array([1, 0, 1, 0]) - actions = jnp.array([Actions.TX.value, Actions.CS.value, Actions.CS.value, Actions.CS.value]) - expected = jnp.array([0, 0, 1, 0]) - result = buffer_clearing(buffer_states, actions) - self.assertTrue(jnp.array_equal(result, expected)) - - # Test no clearing when no matching actions - buffer_states = jnp.array([1, 0, 1, 0]) - actions = jnp.array([Actions.CS.value, Actions.CS.value, Actions.CS.value, Actions.CS.value]) - expected = jnp.array([1, 0, 1, 0]) - result = buffer_clearing(buffer_states, actions) + def test_remove_transmitted_packets(self): + buffer_states = jnp.array([ + [10, 11, 12, EMPTY_PACKET_ID], + [20, 21, EMPTY_PACKET_ID, EMPTY_PACKET_ID], + ]) + tx_masks = jnp.array([ + [1, 0, 1, 0], # remove 1st and 3rd + [0, 1, 0, 0], # remove 2nd + ]) + expected = jnp.array([ + [11, EMPTY_PACKET_ID, EMPTY_PACKET_ID, EMPTY_PACKET_ID], + [20, EMPTY_PACKET_ID, EMPTY_PACKET_ID, EMPTY_PACKET_ID], + ]) + result = remove_transmitted_packets(buffer_states, tx_masks) self.assertTrue(jnp.array_equal(result, expected)) def test_add_new_frames(self): - # Test adding new frames - buffer_states = jnp.array([1, 0, 1, 0]) - new_frames = jnp.array([0, 1, 0, 0]) - expected = jnp.array([1, 1, 1, 0]) - result = add_new_frames(buffer_states, new_frames) - self.assertTrue(jnp.array_equal(result, expected)) - - # Test no new frames added - buffer_states = jnp.array([1, 0, 1, 0]) - new_frames = jnp.array([0, 0, 0, 0]) - expected = jnp.array([1, 0, 1, 0]) - result = add_new_frames(buffer_states, new_frames) - self.assertTrue(jnp.array_equal(result, expected)) + buffer_states = jnp.array([ + [EMPTY_PACKET_ID, EMPTY_PACKET_ID, EMPTY_PACKET_ID, EMPTY_PACKET_ID], + [100, 101, EMPTY_PACKET_ID, EMPTY_PACKET_ID], + ], dtype=jnp.int32) + packet_seqs = jnp.array([0, 2], dtype=jnp.int32) + new_frames = jnp.array([3, 5], dtype=jnp.int32) + + result, new_packet_seqs = add_new_frames(buffer_states, new_frames, packet_seqs) + self.assertEqual(new_packet_seqs[0], 3) + self.assertEqual(new_packet_seqs[1], 7) + self.assertEqual(jnp.sum(result[0] != EMPTY_PACKET_ID), 3) + # Overflow should cap to queue capacity. + self.assertEqual(jnp.sum(result[1] != EMPTY_PACKET_ID), self.QUEUE) def test_simulate(self): - # Test full simulation - buffer_states = jnp.array([1, 1, 1, 0]) - new_frames = jnp.array([0, 1, 1, 1]) + buffer_states = jnp.array([ + [7, 8, EMPTY_PACKET_ID, EMPTY_PACKET_ID], + [EMPTY_PACKET_ID, EMPTY_PACKET_ID, EMPTY_PACKET_ID, EMPTY_PACKET_ID], + [30, EMPTY_PACKET_ID, EMPTY_PACKET_ID, EMPTY_PACKET_ID], + [40, 41, 42, EMPTY_PACKET_ID], + ], dtype=jnp.int32) + packet_seqs = jnp.array([2, 0, 1, 3], dtype=jnp.int32) + new_frames = jnp.array([0, 2, 1, 0], dtype=jnp.int32) actions = jnp.array([Actions.TX.value, Actions.CS.value, Actions.CS.value, Actions.CS.value]) - 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) - self.assertTrue(jnp.array_equal(result_buffer_states, expected_buffer_states)) - self.assertEqual(result_channel_state, expected_channel_state) + result_buffer_states, result_channel_state, result_packet_seqs = simulate( + buffer_states, new_frames, actions, packet_seqs + ) + + self.assertEqual(result_channel_state, 1) + # Successful TX on station 0 removes queue head (7). + self.assertEqual(int(result_buffer_states[0][0]), 8) + self.assertTrue(jnp.array_equal(result_buffer_states[0][1:], jnp.array([EMPTY_PACKET_ID] * 3))) + # New arrivals were appended. + self.assertEqual(jnp.sum(result_buffer_states[1] != EMPTY_PACKET_ID), 2) + self.assertEqual(jnp.sum(result_buffer_states[2] != EMPTY_PACKET_ID), 2) + self.assertTrue(jnp.array_equal(result_packet_seqs, packet_seqs + new_frames)) if __name__ == "__main__": From afab178fa05a44bf7170debbb312cafecc6090a2 Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Tue, 21 Apr 2026 16:05:49 +0200 Subject: [PATCH 02/24] Add new observations --- ltc/agents/dcf.py | 14 ++-- ltc/run.py | 98 +++++++++++++++++++++++++-- ltc/sim/observation.py | 18 +++++ ltc/sim/process_output.py | 113 ++++++++++++++++++++++++++++++-- ltc/sim/sim.py | 60 ++++++++++++----- ltc/utils/scan_states.py | 7 ++ test/sim/test_process_output.py | 47 +++++++++---- test/sim/test_sim.py | 32 +++++++-- 8 files changed, 339 insertions(+), 50 deletions(-) create mode 100644 ltc/sim/observation.py diff --git a/ltc/agents/dcf.py b/ltc/agents/dcf.py index a76fd7a..9b4db66 100644 --- a/ltc/agents/dcf.py +++ b/ltc/agents/dcf.py @@ -3,6 +3,7 @@ from chex import dataclass from reinforced_lib.agents import BaseAgent, AgentState +from ltc.sim.observation import ObsIdx from ltc.sim.constants import Actions @@ -44,16 +45,18 @@ def double_cw(): backoff = jax.random.randint(key, (), 0, state.cw) return DCFState(cw=cw, backoff=backoff) - buffer_count, channel, ret_c, _, _ = env_state[-1] + buffer_count = env_state[-1][ObsIdx.BUFFER_PACKET_COUNT] + channel_busy = env_state[-1][ObsIdx.CHANNEL_LAST_CS_BUSY] + ret_c = env_state[-1][ObsIdx.STATUS_RETRY_COUNTER] return jax.lax.cond( buffer_count == 0, reset, lambda: jax.lax.cond( - jax.lax.bitwise_and(state.backoff > 0, channel == 0), + jax.lax.bitwise_and(state.backoff > 0, channel_busy == 0), countdown, lambda: jax.lax.cond( - jax.lax.bitwise_and(state.backoff > 0, channel != 0), + jax.lax.bitwise_and(state.backoff > 0, channel_busy != 0), freeze, lambda: jax.lax.cond( jax.lax.bitwise_or(reward > 0, ret_c == 0), @@ -70,13 +73,14 @@ def double_cw(): @staticmethod def sample(state, key, env_state): - buffer_count, channel, _, _, _ = env_state[-1] + buffer_count = env_state[-1][ObsIdx.BUFFER_PACKET_COUNT] + channel_busy = env_state[-1][ObsIdx.CHANNEL_LAST_CS_BUSY] return jnp.where( buffer_count == 0, Actions.IDLE.value, jnp.where( - jax.lax.bitwise_and(state.backoff == 0, channel == 0), + jax.lax.bitwise_and(state.backoff == 0, channel_busy == 0), Actions.TX.value, Actions.CS.value ) diff --git a/ltc/run.py b/ltc/run.py index d4bac97..25fb3b1 100644 --- a/ltc/run.py +++ b/ltc/run.py @@ -13,6 +13,7 @@ from tqdm import trange from ltc.agents import BayesianDDQN, DCF, QNetwork, StochasticVariationalNetwork +from ltc.sim.observation import OBS_SIZE from ltc.sim import InitialStateConf, cox_traffic, process_output, simulate from ltc.sim.constants import INITIAL_CAPACITY, Actions, EMPTY_PACKET_ID from ltc.utils.history import build_history_filename, ensure_clean_git_worktree, get_short_commit_hash @@ -100,6 +101,7 @@ def rl_step( drl_step, legacy_step, traffic_step, n, n_drl, n_bins=50, one_shot_step=None, one_shot_target=None, station_change_interval=None, station_change_delta=0, station_change_start_step=0, station_change_stop_step=None, station_change_target=None, + obs_enable_cs_tx_same_type=False, obs_enable_tx_collision_other=False, ): def rl_step_coroutine(c, step): active = schedule_active_stations( @@ -128,9 +130,71 @@ def rl_step_coroutine(c, step): actions = jnp.concatenate([drl_actions, legacy_actions]) traffic_states, new_frames = traffic_step(c.traffic_states, traffic_keys) - buffer_states, channel_state, packet_seqs = simulate(c.buffer_states, new_frames, actions, c.packet_seqs) + ( + buffer_states, + buffer_birth_steps, + channel_state, + packet_seqs, + planned_packets, + successful_packets, + ) = simulate(c.buffer_states, c.buffer_birth_steps, new_frames, actions, c.packet_seqs, step) + + is_ltc = jnp.arange(n) < n_drl + tx_mask = actions == Actions.TX.value + tx_idx = jnp.argmax(tx_mask.astype(jnp.int32)) + tx_is_ltc = is_ltc[tx_idx] + cs_tx_same_type_now = (tx_is_ltc == is_ltc).astype(jnp.float32) + diff_type = is_ltc[:, None] != is_ltc[None, :] + tx_collision_other_now = jnp.any(diff_type & tx_mask[None, :], axis=1).astype(jnp.float32) + + arrival_hist = jnp.roll(c.arrival_hist, -1, axis=1).at[:, -1].set(new_frames.astype(jnp.float32)) + planned_tx_hist = jnp.roll(c.planned_tx_hist, -1, axis=1).at[:, -1].set(planned_packets.astype(jnp.float32)) + success_tx_hist = jnp.roll(c.success_tx_hist, -1, axis=1).at[:, -1].set(successful_packets.astype(jnp.float32)) + channel_busy_hist = jnp.roll(c.channel_busy_hist, -1).at[-1].set((channel_state != 0).astype(jnp.float32)) + collision_hist = jnp.roll(c.collision_hist, -1).at[-1].set((channel_state == -1).astype(jnp.float32)) + tx_hist = jnp.roll(c.tx_hist, -1, axis=0).at[-1].set(tx_mask.astype(jnp.int32)) + + arrival_valid = arrival_hist >= 0.0 + arrival_count = jnp.sum(arrival_valid, axis=1) + traffic_mean_arrival_rate = jnp.where( + arrival_count > 0, + jnp.sum(jnp.where(arrival_valid, arrival_hist, 0.0), axis=1) / arrival_count, + -1.0, + ) + + busy_valid = channel_busy_hist >= 0.0 + busy_count = jnp.sum(busy_valid) + channel_occupancy_pct_window = jnp.where( + busy_count > 0, + jnp.sum(jnp.where(busy_valid, channel_busy_hist, 0.0)) / busy_count, + -1.0, + ) + + coll_valid = collision_hist >= 0.0 + coll_count = jnp.sum(coll_valid) + channel_collisions_pct_window = jnp.where( + coll_count > 0, + jnp.sum(jnp.where(coll_valid, collision_hist, 0.0)) / coll_count, + -1.0, + ) + + planned_valid = planned_tx_hist >= 0.0 + planned_sum = jnp.sum(jnp.where(planned_valid, planned_tx_hist, 0.0), axis=1) + success_sum = jnp.sum(jnp.where(planned_valid, success_tx_hist, 0.0), axis=1) + back_pct = jnp.where(planned_sum > 0, success_sum / planned_sum, -1.0) + + tx_valid_rows = tx_hist[:, 0] >= 0 + unique_ltc_tx_window = jnp.where( + jnp.any(tx_valid_rows), + jnp.sum(jnp.any(jnp.where(tx_valid_rows[:, None], tx_hist[:, :n_drl], 0) > 0, axis=0).astype(jnp.float32)), + -1.0, + ) + obs, rewards, powers = process_output( - c.buffer_states, buffer_states, c.power_states, channel_state, c.obs, actions, c.terminals, reward_key + c.buffer_states, buffer_states, buffer_birth_steps, c.power_states, channel_state, c.obs, actions, + c.terminals, reward_key, step, traffic_mean_arrival_rate, channel_occupancy_pct_window, + channel_collisions_pct_window, back_pct, unique_ltc_tx_window, cs_tx_same_type_now, + tx_collision_other_now, obs_enable_cs_tx_same_type, obs_enable_tx_collision_other ) terminals = jnp.logical_or(c.terminals, powers < 0) @@ -144,7 +208,8 @@ def rl_step_coroutine(c, step): hist, bin_edges = None, None c = Carry( - drl_states, legacy_states, traffic_states, packet_seqs, buffer_states, powers, + drl_states, legacy_states, traffic_states, packet_seqs, buffer_states, buffer_birth_steps, + arrival_hist, planned_tx_hist, success_tx_hist, channel_busy_hist, collision_hist, tx_hist, powers, channel_state, key, obs, actions, rewards, terminals, active ) o = Output( @@ -180,6 +245,9 @@ def setup_args(): parser.add_argument('--n_steps', type=int, default=2000, help='Number of steps per epoch.') parser.add_argument('--window_size', type=int, default=1, help='Size of the observation window for each agent.') parser.add_argument('--queue_size', type=int, default=16, help='Fixed packet-queue length for each station.') + parser.add_argument('--obs_stats_window', type=int, default=1000, help='Window length for rolling observation statistics.') + parser.add_argument('--obs_enable_cs_tx_same_type', action='store_true', default=False, help='Enable CS me-vs-other transmitter feature.') + parser.add_argument('--obs_enable_tx_collision_other', action='store_true', default=False, help='Enable TX collision-involves-other feature.') parser.add_argument('--seed', type=int, default=42, help='Random seed for reproducibility.') parser.add_argument('--save_plots', action='store_true', default=False, help='Whether to save the generated plots.') parser.add_argument('--loc', type=float, default=5.0, help='loc traffic generator parameter.') @@ -209,6 +277,7 @@ def setup_args(): total_steps = n_epochs * n_steps window_size = args.window_size queue_size = args.queue_size + obs_stats_window = args.obs_stats_window seed = args.seed traffic_type = args.traffic_type @@ -219,6 +288,13 @@ def setup_args(): station_change_delta = args.station_change_delta station_change_start_step = args.station_change_start_step station_change_stop_step = args.station_change_stop_step + obs_enable_cs_tx_same_type = args.obs_enable_cs_tx_same_type + obs_enable_tx_collision_other = args.obs_enable_tx_collision_other + + if queue_size <= 0: + raise ValueError('--queue_size must be positive.') + if obs_stats_window <= 0: + raise ValueError('--obs_stats_window must be positive.') if station_change_interval is not None and station_change_interval <= 0: raise ValueError('--station_change_interval must be positive.') @@ -243,9 +319,16 @@ def setup_args(): actions = jnp.zeros(n, dtype=int) packet_seqs = jnp.zeros(n, dtype=jnp.int32) buffer_states = jnp.full((n, queue_size), EMPTY_PACKET_ID, dtype=jnp.int32) + buffer_birth_steps = jnp.full((n, queue_size), -1, dtype=jnp.int32) + arrival_hist = jnp.full((n, obs_stats_window), -1.0, dtype=jnp.float32) + planned_tx_hist = jnp.full((n, obs_stats_window), -1.0, dtype=jnp.float32) + success_tx_hist = jnp.full((n, obs_stats_window), -1.0, dtype=jnp.float32) + channel_busy_hist = jnp.full(obs_stats_window, -1.0, dtype=jnp.float32) + collision_hist = jnp.full(obs_stats_window, -1.0, dtype=jnp.float32) + tx_hist = jnp.full((obs_stats_window, n), -1, dtype=jnp.int32) power_states = jnp.full(n, INITIAL_CAPACITY, dtype=int) channel_state = 0 - obs = jnp.zeros((n, window_size, 5), dtype=jnp.int32) + obs = jnp.zeros((n, window_size, OBS_SIZE), dtype=jnp.float32) rewards = jnp.zeros(n) terminals = jnp.full(n, False, dtype=bool) active = jnp.ones(n, dtype=bool).at[n_init:].set(False) @@ -272,7 +355,7 @@ def setup_args(): ) key, init_key = jax.random.split(key) drl_states, drl_step = init_agents(drl, init_key, n) - drl_states = drl_states.replace(prev_env_state=drl_states.prev_env_state.astype(int)) + drl_states = drl_states.replace(prev_env_state=drl_states.prev_env_state.astype(jnp.float32)) dcf = DCF() key, init_key = jax.random.split(key) @@ -306,10 +389,13 @@ def setup_args(): station_change_start_step=station_change_start_step, station_change_stop_step=station_change_stop_step, station_change_target=station_change_target, + obs_enable_cs_tx_same_type=obs_enable_cs_tx_same_type, + obs_enable_tx_collision_other=obs_enable_tx_collision_other, ) rl_step_fn = jax.jit(rl_step_fn) carry = Carry( - drl_states, legacy_states, traffic_states, packet_seqs, buffer_states, power_states, + drl_states, legacy_states, traffic_states, packet_seqs, buffer_states, buffer_birth_steps, + arrival_hist, planned_tx_hist, success_tx_hist, channel_busy_hist, collision_hist, tx_hist, power_states, channel_state, key, obs, actions, rewards, terminals, active ) all_outputs = [] diff --git a/ltc/sim/observation.py b/ltc/sim/observation.py new file mode 100644 index 0000000..79869a5 --- /dev/null +++ b/ltc/sim/observation.py @@ -0,0 +1,18 @@ +class ObsIdx: + BUFFER_OCCUPANCY_PCT = 0 + BUFFER_PACKET_COUNT = 1 + BUFFER_OLDEST_AGE_NORM = 2 + TRAFFIC_MEAN_ARRIVAL_RATE = 3 + CHANNEL_LAST_CS_BUSY = 4 + CHANNEL_LAST_CS_TX_SAME_TYPE = 5 + CHANNEL_LAST_TX_COLLISION_OTHER = 6 + CHANNEL_OCCUPANCY_PCT_WINDOW = 7 + CHANNEL_COLLISIONS_PCT_WINDOW = 8 + ACTION_LAST = 9 + ACTION_BACK_PCT = 10 + STATUS_RETRY_COUNTER = 11 + STATUS_NO_TX_COUNTER = 12 + STATUS_UNIQUE_LTC_TX_WINDOW = 13 + + +OBS_SIZE = 14 diff --git a/ltc/sim/process_output.py b/ltc/sim/process_output.py index 43d0780..d8ac60c 100644 --- a/ltc/sim/process_output.py +++ b/ltc/sim/process_output.py @@ -1,6 +1,7 @@ import jax import jax.numpy as jnp +from ltc.sim.observation import ObsIdx from ltc.sim.constants import * from ltc.sim.sim import is_buffer_empty @@ -86,8 +87,37 @@ def successful_transmission(args): return reward, ret_c, no_tx -def process_output_i(buffer_state, new_buffer_state, power_state, channel_state, obs, action, terminal, key): - _, _, ret_c, no_tx, _ = obs[-1] +def _oldest_packet_age_norm(buffer_state, buffer_birth_state, current_step): + valid = buffer_state != EMPTY_PACKET_ID + oldest_birth = jnp.min(jnp.where(valid, buffer_birth_state, current_step)) + age = jnp.where(jnp.any(valid), current_step - oldest_birth, -1) + queue_size = jnp.maximum(buffer_state.shape[0], 1) + return jnp.where(age >= 0, age.astype(jnp.float32) / queue_size, -1.0) + + +def process_output_i( + buffer_state, + new_buffer_state, + new_buffer_birth_state, + power_state, + channel_state, + obs, + action, + terminal, + key, + current_step, + traffic_mean_arrival_rate, + channel_occupancy_pct_window, + channel_collisions_pct_window, + back_pct, + unique_ltc_tx_window, + cs_tx_same_type_now, + tx_collision_other_now, + enable_cs_tx_same_type, + enable_tx_collision_other, +): + ret_c = obs[-1, ObsIdx.STATUS_RETRY_COUNTER].astype(jnp.int32) + no_tx = obs[-1, ObsIdx.STATUS_NO_TX_COUNTER].astype(jnp.int32) args = (action, buffer_state, ret_c, channel_state, no_tx, key) reward, ret_c, no_tx = jax.lax.cond(action == Actions.TX.value, transmission, no_transmission, args) @@ -106,14 +136,85 @@ def process_output_i(buffer_state, new_buffer_state, power_state, channel_state, ) buffer_count = jnp.sum((new_buffer_state != EMPTY_PACKET_ID).astype(jnp.int32)) - obs_t = jnp.array([buffer_count, channel_state, ret_c, no_tx, action], dtype=jnp.int32) + buffer_occupancy_pct = buffer_count.astype(jnp.float32) / new_buffer_state.shape[0] + oldest_packet_age_norm = _oldest_packet_age_norm(new_buffer_state, new_buffer_birth_state, current_step) + + cs_busy = jnp.where(action == Actions.CS.value, channel_state != 0, -1).astype(jnp.float32) + cs_tx_same_type = jnp.where( + enable_cs_tx_same_type, + jnp.where(jnp.logical_and(action == Actions.CS.value, channel_state == 1), cs_tx_same_type_now, -1.0), + -1.0, + ) + tx_collision_other = jnp.where( + enable_tx_collision_other, + jnp.where(jnp.logical_and(action == Actions.TX.value, channel_state == -1), tx_collision_other_now, -1.0), + -1.0, + ) + + obs_t = jnp.array( + [ + buffer_occupancy_pct, + buffer_count.astype(jnp.float32), + oldest_packet_age_norm, + traffic_mean_arrival_rate, + cs_busy, + cs_tx_same_type.astype(jnp.float32), + tx_collision_other.astype(jnp.float32), + channel_occupancy_pct_window, + channel_collisions_pct_window, + action.astype(jnp.float32), + back_pct, + ret_c.astype(jnp.float32), + no_tx.astype(jnp.float32), + unique_ltc_tx_window, + ], + dtype=jnp.float32, + ) 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, key): - return jax.vmap(process_output_i, in_axes=(0, 0, 0, None, 0, 0, 0, None))( - buffer_states, new_buffer_states, power_states, channel_state, obs, actions, terminals, key +def process_output( + buffer_states, + new_buffer_states, + new_buffer_birth_states, + power_states, + channel_state, + obs, + actions, + terminals, + key, + current_step, + traffic_mean_arrival_rate, + channel_occupancy_pct_window, + channel_collisions_pct_window, + back_pct, + unique_ltc_tx_window, + cs_tx_same_type_now, + tx_collision_other_now, + enable_cs_tx_same_type, + enable_tx_collision_other, +): + return jax.vmap(process_output_i, in_axes=(0, 0, 0, 0, None, 0, 0, 0, None, None, 0, None, None, 0, None, 0, 0, None, None))( + buffer_states, + new_buffer_states, + new_buffer_birth_states, + power_states, + channel_state, + obs, + actions, + terminals, + key, + current_step, + traffic_mean_arrival_rate, + channel_occupancy_pct_window, + channel_collisions_pct_window, + back_pct, + unique_ltc_tx_window, + cs_tx_same_type_now, + tx_collision_other_now, + enable_cs_tx_same_type, + enable_tx_collision_other, ) diff --git a/ltc/sim/sim.py b/ltc/sim/sim.py index e9919a0..abdacaf 100644 --- a/ltc/sim/sim.py +++ b/ltc/sim/sim.py @@ -51,44 +51,72 @@ def remove_transmitted_packets(buffer_states, tx_masks): return jax.vmap(remove_packets_from_queue)(buffer_states, tx_masks) -def enqueue_generated_packets(buffer_state, new_frames, station_id, packet_seq): +def enqueue_generated_packets(buffer_state, buffer_birth_state, new_frames, station_id, packet_seq, current_step): """ - Appends generated packets at queue tail up to free capacity. + Queue semantics: + - transmitted packets are removed elsewhere, + - arrivals are appended at queue tail, + - if overflow occurs, oldest packets are dropped. """ q = buffer_state.shape[0] new_frames = jnp.maximum(new_frames.astype(jnp.int32), 0) occupied = jnp.sum((buffer_state != EMPTY_PACKET_ID).astype(jnp.int32)) - capacity = q - occupied - to_add = jnp.minimum(new_frames.astype(jnp.int32), capacity) + + kept_new = jnp.minimum(new_frames, q) + kept_old = jnp.minimum(occupied, q - kept_new) + old_start = occupied - kept_old + new_start = packet_seq + (new_frames - kept_new) idx = jnp.arange(q, dtype=jnp.int32) - rel = idx - occupied - add_mask = (idx >= occupied) & (rel < to_add) - local_ids = packet_seq + jnp.maximum(rel.astype(jnp.int32), 0) - packet_ids = _cantor_pairing(station_id.astype(jnp.int32), local_ids) - next_state = jnp.where(add_mask, packet_ids, buffer_state) - next_packet_seq = packet_seq + new_frames + take_old = idx < kept_old + old_idx = jnp.clip(old_start + idx, 0, jnp.maximum(occupied - 1, 0)) + old_packets = jnp.where(take_old, buffer_state[old_idx], EMPTY_PACKET_ID) + old_births = jnp.where(take_old, buffer_birth_state[old_idx], -1) - return next_state, next_packet_seq + new_rel = idx - kept_old + take_new = jnp.logical_and(idx >= kept_old, new_rel < kept_new) + new_local_ids = new_start + jnp.maximum(new_rel, 0) + new_packets = _cantor_pairing(station_id.astype(jnp.int32), new_local_ids) + new_births = jnp.full((q,), current_step, dtype=jnp.int32) + + next_state = jnp.where(take_old, old_packets, jnp.where(take_new, new_packets, EMPTY_PACKET_ID)) + next_birth = jnp.where(take_old, old_births, jnp.where(take_new, new_births, -1)) + next_packet_seq = packet_seq + new_frames + return next_state, next_birth, next_packet_seq -def add_new_frames(buffer_states, new_frames, packet_seqs): +def add_new_frames(buffer_states, buffer_birth_steps, new_frames, packet_seqs, current_step): station_ids = jnp.arange(buffer_states.shape[0], dtype=jnp.int32) - return jax.vmap(enqueue_generated_packets)(buffer_states, new_frames, station_ids, packet_seqs) + return jax.vmap(enqueue_generated_packets, in_axes=(0, 0, 0, 0, 0, None))( + buffer_states, buffer_birth_steps, new_frames, station_ids, packet_seqs, current_step + ) -def simulate(buffer_states, new_frames, actions, packet_seqs): +def simulate(buffer_states, buffer_birth_steps, new_frames, actions, packet_seqs, current_step): """ One simulation step with fixed-size packet queues. On successful TX, one packet from queue head is removed for transmitting station. """ channel_state = channel_state_selector(actions) + pre_nonempty = jnp.any(buffer_states != EMPTY_PACKET_ID, axis=1) success_tx = (actions == Actions.TX.value) & (channel_state == 1) + planned_packets = (actions == Actions.TX.value).astype(jnp.int32) + successful_packets = (success_tx & pre_nonempty).astype(jnp.int32) queue_size = buffer_states.shape[1] head_mask = jnp.arange(queue_size) == 0 tx_masks = success_tx[:, None] & head_mask[None, :] buffer_states = remove_transmitted_packets(buffer_states, tx_masks) - buffer_states, packet_seqs = add_new_frames(buffer_states, new_frames, packet_seqs) - return buffer_states, channel_state, packet_seqs + buffer_birth_steps = remove_transmitted_packets(buffer_birth_steps, tx_masks) + buffer_states, buffer_birth_steps, packet_seqs = add_new_frames( + buffer_states, buffer_birth_steps, new_frames, packet_seqs, current_step + ) + return ( + buffer_states, + buffer_birth_steps, + channel_state, + packet_seqs, + planned_packets, + successful_packets, + ) diff --git a/ltc/utils/scan_states.py b/ltc/utils/scan_states.py index 2009f92..05045f0 100644 --- a/ltc/utils/scan_states.py +++ b/ltc/utils/scan_states.py @@ -14,6 +14,13 @@ class Carry: traffic_states: ModelState packet_seqs: jax.Array buffer_states: jax.Array + buffer_birth_steps: jax.Array + arrival_hist: jax.Array + planned_tx_hist: jax.Array + success_tx_hist: jax.Array + channel_busy_hist: jax.Array + collision_hist: jax.Array + tx_hist: jax.Array power_states: jax.Array channel_state: int key: jax.random.PRNGKey diff --git a/test/sim/test_process_output.py b/test/sim/test_process_output.py index 039e25c..f4177b0 100644 --- a/test/sim/test_process_output.py +++ b/test/sim/test_process_output.py @@ -3,6 +3,7 @@ import jax import jax.numpy as jnp +from ltc.sim.observation import ObsIdx, OBS_SIZE from ltc.sim.constants import * from ltc.sim.process_output import * @@ -66,6 +67,12 @@ def test_process_rl_output(self): [9, EMPTY, EMPTY], [EMPTY, EMPTY, EMPTY], ]) + new_buffer_birth_states = jnp.array([ + [-1, -1, -1], + [-1, -1, -1], + [3, -1, -1], + [-1, -1, -1], + ], dtype=jnp.int32) new_buffer_states = jnp.array([ [EMPTY, EMPTY, EMPTY], [EMPTY, EMPTY, EMPTY], @@ -76,25 +83,39 @@ def test_process_rl_output(self): actions = jnp.array([Actions.TX.value, Actions.CS.value, Actions.TX.value, Actions.TX.value]) terminals = jnp.array([False, False, False, False]) channel_state = -1 - obs = jnp.array([[ - [1, 1, 6, 0, 100], - [1, 1, 7, 0, 100], - [1, 1, 8, 0, 100] - ]] * 4) - - expected_obs_0 = jnp.array([ - [1, 1, 7, 0, 100], - [1, 1, 8, 0, 100], - [0, -1, 0, 0, Actions.TX.value], - ]) + obs = jnp.zeros((4, 3, OBS_SIZE), dtype=jnp.float32) + obs = obs.at[:, -1, ObsIdx.STATUS_RETRY_COUNTER].set(8) expected_R_0 = -1.0 expected_power = power_states + jnp.array([-TX_CONSUMPTION, -CS_CONSUMPTION, -TX_CONSUMPTION, -TX_CONSUMPTION]) result_obs, result_R, power = process_output( - buffer_states, new_buffer_states, power_states, channel_state, obs, actions, terminals, KEY + buffer_states=buffer_states, + new_buffer_states=new_buffer_states, + new_buffer_birth_states=new_buffer_birth_states, + power_states=power_states, + channel_state=channel_state, + obs=obs, + actions=actions, + terminals=terminals, + key=KEY, + current_step=10, + traffic_mean_arrival_rate=jnp.array([1.0, 2.0, 3.0, 4.0], dtype=jnp.float32), + channel_occupancy_pct_window=0.25, + channel_collisions_pct_window=0.5, + back_pct=jnp.array([0.0, 0.6, 1.0, 0.0], dtype=jnp.float32), + unique_ltc_tx_window=2.0, + cs_tx_same_type_now=jnp.array([0.0, 1.0, 0.0, 0.0], dtype=jnp.float32), + tx_collision_other_now=jnp.array([1.0, 0.0, 1.0, 1.0], dtype=jnp.float32), + enable_cs_tx_same_type=True, + enable_tx_collision_other=True, ) - self.assertTrue(jnp.array_equal(result_obs[0], expected_obs_0)) + self.assertEqual(result_obs.shape[-1], OBS_SIZE) + self.assertEqual(result_obs[1, -1, ObsIdx.CHANNEL_LAST_CS_BUSY], 1.0) + self.assertEqual(result_obs[1, -1, ObsIdx.CHANNEL_LAST_CS_TX_SAME_TYPE], -1.0) + self.assertEqual(result_obs[0, -1, ObsIdx.CHANNEL_LAST_TX_COLLISION_OTHER], 1.0) + self.assertEqual(result_obs[2, -1, ObsIdx.BUFFER_PACKET_COUNT], 1.0) + self.assertEqual(result_obs[2, -1, ObsIdx.ACTION_BACK_PCT], 1.0) self.assertEqual(result_R[0], expected_R_0) self.assertTrue(jnp.array_equal(power, expected_power)) diff --git a/test/sim/test_sim.py b/test/sim/test_sim.py index 829bcc1..ff816eb 100644 --- a/test/sim/test_sim.py +++ b/test/sim/test_sim.py @@ -43,15 +43,23 @@ def test_add_new_frames(self): [EMPTY_PACKET_ID, EMPTY_PACKET_ID, EMPTY_PACKET_ID, EMPTY_PACKET_ID], [100, 101, EMPTY_PACKET_ID, EMPTY_PACKET_ID], ], dtype=jnp.int32) + buffer_birth_steps = jnp.array([ + [-1, -1, -1, -1], + [1, 2, -1, -1], + ], dtype=jnp.int32) packet_seqs = jnp.array([0, 2], dtype=jnp.int32) new_frames = jnp.array([3, 5], dtype=jnp.int32) - result, new_packet_seqs = add_new_frames(buffer_states, new_frames, packet_seqs) + result, result_births, new_packet_seqs = add_new_frames( + buffer_states, buffer_birth_steps, new_frames, packet_seqs, current_step=10 + ) self.assertEqual(new_packet_seqs[0], 3) self.assertEqual(new_packet_seqs[1], 7) self.assertEqual(jnp.sum(result[0] != EMPTY_PACKET_ID), 3) - # Overflow should cap to queue capacity. + # Overflow should cap to queue capacity and drop oldest first. self.assertEqual(jnp.sum(result[1] != EMPTY_PACKET_ID), self.QUEUE) + self.assertEqual(int(result[1][0]), 13) + self.assertEqual(int(result_births[1][0]), 10) def test_simulate(self): buffer_states = jnp.array([ @@ -60,22 +68,38 @@ def test_simulate(self): [30, EMPTY_PACKET_ID, EMPTY_PACKET_ID, EMPTY_PACKET_ID], [40, 41, 42, EMPTY_PACKET_ID], ], dtype=jnp.int32) + buffer_birth_steps = jnp.array([ + [1, 2, -1, -1], + [-1, -1, -1, -1], + [5, -1, -1, -1], + [2, 3, 4, -1], + ], dtype=jnp.int32) packet_seqs = jnp.array([2, 0, 1, 3], dtype=jnp.int32) new_frames = jnp.array([0, 2, 1, 0], dtype=jnp.int32) actions = jnp.array([Actions.TX.value, Actions.CS.value, Actions.CS.value, Actions.CS.value]) - result_buffer_states, result_channel_state, result_packet_seqs = simulate( - buffer_states, new_frames, actions, packet_seqs + ( + result_buffer_states, + result_buffer_birth_steps, + result_channel_state, + result_packet_seqs, + planned_packets, + successful_packets, + ) = simulate( + buffer_states, buffer_birth_steps, new_frames, actions, packet_seqs, current_step=10 ) self.assertEqual(result_channel_state, 1) # Successful TX on station 0 removes queue head (7). self.assertEqual(int(result_buffer_states[0][0]), 8) self.assertTrue(jnp.array_equal(result_buffer_states[0][1:], jnp.array([EMPTY_PACKET_ID] * 3))) + self.assertTrue(jnp.array_equal(result_buffer_birth_steps[0][1:], jnp.array([-1, -1, -1], dtype=jnp.int32))) # New arrivals were appended. self.assertEqual(jnp.sum(result_buffer_states[1] != EMPTY_PACKET_ID), 2) self.assertEqual(jnp.sum(result_buffer_states[2] != EMPTY_PACKET_ID), 2) self.assertTrue(jnp.array_equal(result_packet_seqs, packet_seqs + new_frames)) + self.assertTrue(jnp.array_equal(planned_packets, jnp.array([1, 0, 0, 0], dtype=jnp.int32))) + self.assertTrue(jnp.array_equal(successful_packets, jnp.array([1, 0, 0, 0], dtype=jnp.int32))) if __name__ == "__main__": From ef73a36e4c4e07f14ed582e3c6481a1a035f6e7f Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Tue, 21 Apr 2026 16:36:31 +0200 Subject: [PATCH 03/24] Add variable length action space --- ltc/agents/dcf.py | 3 +- ltc/run.py | 351 +++++++++++++++++++++++++------- ltc/sim/constants.py | 26 +++ ltc/sim/observation.py | 19 +- ltc/sim/process_output.py | 14 +- ltc/utils/scan_states.py | 1 + test/sim/test_process_output.py | 1 - test/test_action_space.py | 50 +++++ 8 files changed, 362 insertions(+), 103 deletions(-) create mode 100644 test/test_action_space.py diff --git a/ltc/agents/dcf.py b/ltc/agents/dcf.py index 9b4db66..0bb48d2 100644 --- a/ltc/agents/dcf.py +++ b/ltc/agents/dcf.py @@ -3,8 +3,7 @@ from chex import dataclass from reinforced_lib.agents import BaseAgent, AgentState -from ltc.sim.observation import ObsIdx -from ltc.sim.constants import Actions +from ltc.sim.constants import Actions, ObsIdx @dataclass diff --git a/ltc/run.py b/ltc/run.py index 25fb3b1..7beacef 100644 --- a/ltc/run.py +++ b/ltc/run.py @@ -13,9 +13,15 @@ from tqdm import trange from ltc.agents import BayesianDDQN, DCF, QNetwork, StochasticVariationalNetwork -from ltc.sim.observation import OBS_SIZE from ltc.sim import InitialStateConf, cox_traffic, process_output, simulate -from ltc.sim.constants import INITIAL_CAPACITY, Actions, EMPTY_PACKET_ID +from ltc.sim.constants import ( + INITIAL_CAPACITY, + Actions, + EMPTY_PACKET_ID, + OBS_SIZE, + MAX_MACRO_DURATION, + LEGACY_TX_DURATION, +) from ltc.utils.history import build_history_filename, ensure_clean_git_worktree, get_short_commit_hash from ltc.utils.scan_states import Carry, Output from ltc.utils.plots import plot_all, plot_first @@ -97,11 +103,76 @@ def schedule_active_stations( return next_active +def action_space_size(action_space_variant, max_duration): + if action_space_variant in { + 'txcs_duration_set', + 'txcs_to_discrete_duration', + 'txcs_to_continuous_duration', + 'txcs_to_geometric_duration', + }: + return 2 * max_duration + if action_space_variant == 'tx_stage_commit': + return 3 + raise ValueError(f'Unknown action-space variant: {action_space_variant}') + + +def decode_drl_actions(raw_actions, staged_tx, keys, action_space_variant, max_duration): + raw_actions = raw_actions.astype(jnp.int32) + + if action_space_variant in { + 'txcs_duration_set', + 'txcs_to_discrete_duration', + 'txcs_to_continuous_duration', + 'txcs_to_geometric_duration', + }: + is_tx = raw_actions < max_duration + duration_idx = jnp.where(is_tx, raw_actions, raw_actions - max_duration) + action_type = jnp.where(is_tx, Actions.TX.value, Actions.CS.value) + + if action_space_variant == 'txcs_to_geometric_duration': + p = (duration_idx.astype(jnp.float32) + 1.0) / (max_duration + 1.0) + duration = jax.vmap(lambda k, pp: jax.random.geometric(k, pp))(keys, p).astype(jnp.int32) + duration = jnp.clip(duration, 1, max_duration) + else: + duration = duration_idx + 1 + + return action_type.astype(jnp.int32), duration.astype(jnp.int32), staged_tx + + if action_space_variant == 'tx_stage_commit': + is_stage = raw_actions == 1 + is_commit = raw_actions == 2 + staged_next = jnp.where(is_stage, staged_tx + 1, staged_tx) + commit_duration = jnp.clip(staged_tx, 1, max_duration) + has_stage = staged_tx > 0 + action_type = jnp.where(is_commit & has_stage, Actions.TX.value, Actions.CS.value).astype(jnp.int32) + duration = jnp.where(is_commit & has_stage, commit_duration, 1).astype(jnp.int32) + staged_next = jnp.where(is_commit, 0, staged_next).astype(jnp.int32) + return action_type, duration, staged_next + + raise ValueError(f'Unknown action-space variant: {action_space_variant}') + + +def decode_legacy_actions(raw_actions, legacy_tx_duration): + raw_actions = raw_actions.astype(jnp.int32) + action_type = jnp.where( + raw_actions == Actions.TX.value, + Actions.TX.value, + jnp.where(raw_actions == Actions.CS.value, Actions.CS.value, Actions.IDLE.value), + ).astype(jnp.int32) + duration = jnp.where( + raw_actions == Actions.TX.value, + legacy_tx_duration, + 1, + ).astype(jnp.int32) + return action_type, duration + + def rl_step( drl_step, legacy_step, traffic_step, n, n_drl, n_bins=50, one_shot_step=None, one_shot_target=None, station_change_interval=None, station_change_delta=0, station_change_start_step=0, station_change_stop_step=None, station_change_target=None, obs_enable_cs_tx_same_type=False, obs_enable_tx_collision_other=False, + action_space_variant='txcs_duration_set', max_duration=5, legacy_tx_duration=5, ): def rl_step_coroutine(c, step): active = schedule_active_stations( @@ -119,83 +190,189 @@ def rl_step_coroutine(c, step): key, drl_keys, legacy_keys, traffic_key, reward_key = jax.random.split(c.key, 5) drl_keys = jax.random.split(drl_keys, n_drl) legacy_keys = jax.random.split(legacy_keys, n - n_drl) - traffic_keys = jax.random.split(traffic_key, n) + traffic_slot_keys = jax.random.split(traffic_key, max_duration) + reward_slot_keys = jax.random.split(reward_key, max_duration) - drl_states, drl_actions = yield drl_step( + drl_states, drl_actions_raw = yield drl_step( c.drl_states, drl_keys, c.obs[:n_drl], c.actions[:n_drl], c.rewards[:n_drl], c.terminals[:n_drl], active[:n_drl] ) - legacy_states, legacy_actions = legacy_step( + legacy_states, legacy_actions_raw = legacy_step( c.legacy_states, legacy_keys, c.obs[n_drl:], c.actions[n_drl:], c.rewards[n_drl:], c.terminals[n_drl:], active[n_drl:] ) - actions = jnp.concatenate([drl_actions, legacy_actions]) - - traffic_states, new_frames = traffic_step(c.traffic_states, traffic_keys) - ( - buffer_states, - buffer_birth_steps, - channel_state, - packet_seqs, - planned_packets, - successful_packets, - ) = simulate(c.buffer_states, c.buffer_birth_steps, new_frames, actions, c.packet_seqs, step) - - is_ltc = jnp.arange(n) < n_drl - tx_mask = actions == Actions.TX.value - tx_idx = jnp.argmax(tx_mask.astype(jnp.int32)) - tx_is_ltc = is_ltc[tx_idx] - cs_tx_same_type_now = (tx_is_ltc == is_ltc).astype(jnp.float32) - diff_type = is_ltc[:, None] != is_ltc[None, :] - tx_collision_other_now = jnp.any(diff_type & tx_mask[None, :], axis=1).astype(jnp.float32) - - arrival_hist = jnp.roll(c.arrival_hist, -1, axis=1).at[:, -1].set(new_frames.astype(jnp.float32)) - planned_tx_hist = jnp.roll(c.planned_tx_hist, -1, axis=1).at[:, -1].set(planned_packets.astype(jnp.float32)) - success_tx_hist = jnp.roll(c.success_tx_hist, -1, axis=1).at[:, -1].set(successful_packets.astype(jnp.float32)) - channel_busy_hist = jnp.roll(c.channel_busy_hist, -1).at[-1].set((channel_state != 0).astype(jnp.float32)) - collision_hist = jnp.roll(c.collision_hist, -1).at[-1].set((channel_state == -1).astype(jnp.float32)) - tx_hist = jnp.roll(c.tx_hist, -1, axis=0).at[-1].set(tx_mask.astype(jnp.int32)) - - arrival_valid = arrival_hist >= 0.0 - arrival_count = jnp.sum(arrival_valid, axis=1) - traffic_mean_arrival_rate = jnp.where( - arrival_count > 0, - jnp.sum(jnp.where(arrival_valid, arrival_hist, 0.0), axis=1) / arrival_count, - -1.0, - ) - - busy_valid = channel_busy_hist >= 0.0 - busy_count = jnp.sum(busy_valid) - channel_occupancy_pct_window = jnp.where( - busy_count > 0, - jnp.sum(jnp.where(busy_valid, channel_busy_hist, 0.0)) / busy_count, - -1.0, - ) - - coll_valid = collision_hist >= 0.0 - coll_count = jnp.sum(coll_valid) - channel_collisions_pct_window = jnp.where( - coll_count > 0, - jnp.sum(jnp.where(coll_valid, collision_hist, 0.0)) / coll_count, - -1.0, - ) - - planned_valid = planned_tx_hist >= 0.0 - planned_sum = jnp.sum(jnp.where(planned_valid, planned_tx_hist, 0.0), axis=1) - success_sum = jnp.sum(jnp.where(planned_valid, success_tx_hist, 0.0), axis=1) - back_pct = jnp.where(planned_sum > 0, success_sum / planned_sum, -1.0) - - tx_valid_rows = tx_hist[:, 0] >= 0 - unique_ltc_tx_window = jnp.where( - jnp.any(tx_valid_rows), - jnp.sum(jnp.any(jnp.where(tx_valid_rows[:, None], tx_hist[:, :n_drl], 0) > 0, axis=0).astype(jnp.float32)), - -1.0, + raw_actions = jnp.concatenate([drl_actions_raw, legacy_actions_raw]) + drl_action_types, drl_durations, staged_tx = decode_drl_actions( + drl_actions_raw, c.staged_tx[:n_drl], drl_keys, action_space_variant, max_duration ) + legacy_action_types, legacy_durations = decode_legacy_actions(legacy_actions_raw, legacy_tx_duration) + action_types = jnp.concatenate([drl_action_types, legacy_action_types]) + durations = jnp.concatenate([drl_durations, legacy_durations]) + max_macro = jnp.max(durations) + staged_tx = jnp.concatenate([staged_tx, c.staged_tx[n_drl:]]) + + traffic_states = c.traffic_states + buffer_states = c.buffer_states + buffer_birth_steps = c.buffer_birth_steps + packet_seqs = c.packet_seqs + obs = c.obs + powers = c.power_states + rewards = jnp.zeros_like(c.rewards) + arrival_hist = c.arrival_hist + planned_tx_hist = c.planned_tx_hist + success_tx_hist = c.success_tx_hist + channel_busy_hist = c.channel_busy_hist + collision_hist = c.collision_hist + tx_hist = c.tx_hist + channel_state = c.channel_state + last_new_frames = jnp.zeros((n,), dtype=jnp.int32) + + for inner in range(max_duration): + do_slot = inner < max_macro + + def step_slot(args): + ( + traffic_states, + buffer_states, + buffer_birth_steps, + packet_seqs, + obs, + powers, + rewards, + arrival_hist, + planned_tx_hist, + success_tx_hist, + channel_busy_hist, + collision_hist, + tx_hist, + channel_state, + last_new_frames, + ) = args + slot_traffic_keys = jax.random.split(traffic_slot_keys[inner], n) + traffic_states, new_frames = traffic_step(traffic_states, slot_traffic_keys) + slot_actions = jnp.where(inner < durations, action_types, Actions.IDLE.value) + ( + new_buffer_states, + new_buffer_birth_steps, + channel_state, + packet_seqs, + planned_packets, + successful_packets, + ) = simulate(buffer_states, buffer_birth_steps, new_frames, slot_actions, packet_seqs, step * max_duration + inner) + + is_ltc = jnp.arange(n) < n_drl + tx_mask = slot_actions == Actions.TX.value + tx_idx = jnp.argmax(tx_mask.astype(jnp.int32)) + tx_is_ltc = is_ltc[tx_idx] + cs_tx_same_type_now = (tx_is_ltc == is_ltc).astype(jnp.float32) + diff_type = is_ltc[:, None] != is_ltc[None, :] + tx_collision_other_now = jnp.any(diff_type & tx_mask[None, :], axis=1).astype(jnp.float32) + + arrival_hist = jnp.roll(arrival_hist, -1, axis=1).at[:, -1].set(new_frames.astype(jnp.float32)) + planned_tx_hist = jnp.roll(planned_tx_hist, -1, axis=1).at[:, -1].set(planned_packets.astype(jnp.float32)) + success_tx_hist = jnp.roll(success_tx_hist, -1, axis=1).at[:, -1].set(successful_packets.astype(jnp.float32)) + channel_busy_hist = jnp.roll(channel_busy_hist, -1).at[-1].set((channel_state != 0).astype(jnp.float32)) + collision_hist = jnp.roll(collision_hist, -1).at[-1].set((channel_state == -1).astype(jnp.float32)) + tx_hist = jnp.roll(tx_hist, -1, axis=0).at[-1].set(tx_mask.astype(jnp.int32)) + + arrival_valid = arrival_hist >= 0.0 + arrival_count = jnp.sum(arrival_valid, axis=1) + traffic_mean_arrival_rate = jnp.where( + arrival_count > 0, + jnp.sum(jnp.where(arrival_valid, arrival_hist, 0.0), axis=1) / arrival_count, + -1.0, + ) + + busy_valid = channel_busy_hist >= 0.0 + busy_count = jnp.sum(busy_valid) + channel_occupancy_pct_window = jnp.where( + busy_count > 0, + jnp.sum(jnp.where(busy_valid, channel_busy_hist, 0.0)) / busy_count, + -1.0, + ) + + coll_valid = collision_hist >= 0.0 + coll_count = jnp.sum(coll_valid) + channel_collisions_pct_window = jnp.where( + coll_count > 0, + jnp.sum(jnp.where(coll_valid, collision_hist, 0.0)) / coll_count, + -1.0, + ) + + planned_valid = planned_tx_hist >= 0.0 + planned_sum = jnp.sum(jnp.where(planned_valid, planned_tx_hist, 0.0), axis=1) + success_sum = jnp.sum(jnp.where(planned_valid, success_tx_hist, 0.0), axis=1) + back_pct = jnp.where(planned_sum > 0, success_sum / planned_sum, -1.0) + + tx_valid_rows = tx_hist[:, 0] >= 0 + unique_ltc_tx_window = jnp.where( + jnp.any(tx_valid_rows), + jnp.sum(jnp.any(jnp.where(tx_valid_rows[:, None], tx_hist[:, :n_drl], 0) > 0, axis=0).astype(jnp.float32)), + -1.0, + ) + + obs, slot_rewards, powers = process_output( + buffer_states, new_buffer_states, new_buffer_birth_steps, powers, channel_state, obs, slot_actions, + c.terminals, reward_slot_keys[inner], step * max_duration + inner, traffic_mean_arrival_rate, channel_occupancy_pct_window, + channel_collisions_pct_window, back_pct, unique_ltc_tx_window, cs_tx_same_type_now, + tx_collision_other_now, obs_enable_cs_tx_same_type, obs_enable_tx_collision_other, roll_obs=(inner + 1 == max_macro) + ) + rewards = rewards + slot_rewards + return ( + traffic_states, + new_buffer_states, + new_buffer_birth_steps, + packet_seqs, + obs, + powers, + rewards, + arrival_hist, + planned_tx_hist, + success_tx_hist, + channel_busy_hist, + collision_hist, + tx_hist, + channel_state, + new_frames, + ) + + ( + traffic_states, + buffer_states, + buffer_birth_steps, + packet_seqs, + obs, + powers, + rewards, + arrival_hist, + planned_tx_hist, + success_tx_hist, + channel_busy_hist, + collision_hist, + tx_hist, + channel_state, + last_new_frames, + ) = jax.lax.cond( + do_slot, + step_slot, + lambda x: x, + ( + traffic_states, + buffer_states, + buffer_birth_steps, + packet_seqs, + obs, + powers, + rewards, + arrival_hist, + planned_tx_hist, + success_tx_hist, + channel_busy_hist, + collision_hist, + tx_hist, + channel_state, + last_new_frames, + ), + ) - obs, rewards, powers = process_output( - c.buffer_states, buffer_states, buffer_birth_steps, c.power_states, channel_state, c.obs, actions, - c.terminals, reward_key, step, traffic_mean_arrival_rate, channel_occupancy_pct_window, - channel_collisions_pct_window, back_pct, unique_ltc_tx_window, cs_tx_same_type_now, - tx_collision_other_now, obs_enable_cs_tx_same_type, obs_enable_tx_collision_other - ) terminals = jnp.logical_or(c.terminals, powers < 0) if n_drl > 0: @@ -209,12 +386,12 @@ def rl_step_coroutine(c, step): c = Carry( drl_states, legacy_states, traffic_states, packet_seqs, buffer_states, buffer_birth_steps, - arrival_hist, planned_tx_hist, success_tx_hist, channel_busy_hist, collision_hist, tx_hist, powers, - channel_state, key, obs, actions, rewards, terminals, active + arrival_hist, planned_tx_hist, success_tx_hist, channel_busy_hist, collision_hist, tx_hist, staged_tx, powers, + channel_state, key, obs, raw_actions, rewards, terminals, active ) o = Output( - legacy_states, obs, actions, rewards, terminals, buffer_states, powers, - new_frames, channel_state, active, hist, bin_edges + legacy_states, obs, raw_actions, rewards, terminals, buffer_states, powers, + last_new_frames, channel_state, active, hist, bin_edges ) yield c, o @@ -248,6 +425,19 @@ def setup_args(): parser.add_argument('--obs_stats_window', type=int, default=1000, help='Window length for rolling observation statistics.') parser.add_argument('--obs_enable_cs_tx_same_type', action='store_true', default=False, help='Enable CS me-vs-other transmitter feature.') parser.add_argument('--obs_enable_tx_collision_other', action='store_true', default=False, help='Enable TX collision-involves-other feature.') + parser.add_argument( + '--action_space_variant', + type=str, + default='txcs_duration_set', + choices=[ + 'txcs_duration_set', + 'txcs_to_discrete_duration', + 'txcs_to_continuous_duration', + 'txcs_to_geometric_duration', + 'tx_stage_commit', + ], + help='Macro action-space variant.', + ) parser.add_argument('--seed', type=int, default=42, help='Random seed for reproducibility.') parser.add_argument('--save_plots', action='store_true', default=False, help='Whether to save the generated plots.') parser.add_argument('--loc', type=float, default=5.0, help='loc traffic generator parameter.') @@ -290,6 +480,7 @@ def setup_args(): station_change_stop_step = args.station_change_stop_step obs_enable_cs_tx_same_type = args.obs_enable_cs_tx_same_type obs_enable_tx_collision_other = args.obs_enable_tx_collision_other + action_space_variant = args.action_space_variant if queue_size <= 0: raise ValueError('--queue_size must be positive.') @@ -315,7 +506,7 @@ def setup_args(): one_shot_target = n_final key = jax.random.key(seed) - num_actions = 2 + num_actions = action_space_size(action_space_variant, MAX_MACRO_DURATION) actions = jnp.zeros(n, dtype=int) packet_seqs = jnp.zeros(n, dtype=jnp.int32) buffer_states = jnp.full((n, queue_size), EMPTY_PACKET_ID, dtype=jnp.int32) @@ -326,6 +517,7 @@ def setup_args(): channel_busy_hist = jnp.full(obs_stats_window, -1.0, dtype=jnp.float32) collision_hist = jnp.full(obs_stats_window, -1.0, dtype=jnp.float32) tx_hist = jnp.full((obs_stats_window, n), -1, dtype=jnp.int32) + staged_tx = jnp.zeros(n, dtype=jnp.int32) power_states = jnp.full(n, INITIAL_CAPACITY, dtype=int) channel_state = 0 obs = jnp.zeros((n, window_size, OBS_SIZE), dtype=jnp.float32) @@ -391,11 +583,14 @@ def setup_args(): station_change_target=station_change_target, obs_enable_cs_tx_same_type=obs_enable_cs_tx_same_type, obs_enable_tx_collision_other=obs_enable_tx_collision_other, + action_space_variant=action_space_variant, + max_duration=MAX_MACRO_DURATION, + legacy_tx_duration=LEGACY_TX_DURATION, ) rl_step_fn = jax.jit(rl_step_fn) carry = Carry( drl_states, legacy_states, traffic_states, packet_seqs, buffer_states, buffer_birth_steps, - arrival_hist, planned_tx_hist, success_tx_hist, channel_busy_hist, collision_hist, tx_hist, power_states, + arrival_hist, planned_tx_hist, success_tx_hist, channel_busy_hist, collision_hist, tx_hist, staged_tx, power_states, channel_state, key, obs, actions, rewards, terminals, active ) all_outputs = [] diff --git a/ltc/sim/constants.py b/ltc/sim/constants.py index cd87fe5..42a8f88 100644 --- a/ltc/sim/constants.py +++ b/ltc/sim/constants.py @@ -45,3 +45,29 @@ class Actions(Enum): # Queue representation EMPTY_PACKET_ID = -1 + +""" +Macro-action defaults +""" +MAX_MACRO_DURATION = 5 +LEGACY_TX_DURATION = 5 + + +class ObsIdx: + BUFFER_OCCUPANCY_PCT = 0 + BUFFER_PACKET_COUNT = 1 + BUFFER_OLDEST_AGE_NORM = 2 + TRAFFIC_MEAN_ARRIVAL_RATE = 3 + CHANNEL_LAST_CS_BUSY = 4 + CHANNEL_LAST_CS_TX_SAME_TYPE = 5 + CHANNEL_LAST_TX_COLLISION_OTHER = 6 + CHANNEL_OCCUPANCY_PCT_WINDOW = 7 + CHANNEL_COLLISIONS_PCT_WINDOW = 8 + ACTION_LAST = 9 + ACTION_BACK_PCT = 10 + STATUS_RETRY_COUNTER = 11 + STATUS_NO_TX_COUNTER = 12 + STATUS_UNIQUE_LTC_TX_WINDOW = 13 + + +OBS_SIZE = 14 diff --git a/ltc/sim/observation.py b/ltc/sim/observation.py index 79869a5..00bc499 100644 --- a/ltc/sim/observation.py +++ b/ltc/sim/observation.py @@ -1,18 +1 @@ -class ObsIdx: - BUFFER_OCCUPANCY_PCT = 0 - BUFFER_PACKET_COUNT = 1 - BUFFER_OLDEST_AGE_NORM = 2 - TRAFFIC_MEAN_ARRIVAL_RATE = 3 - CHANNEL_LAST_CS_BUSY = 4 - CHANNEL_LAST_CS_TX_SAME_TYPE = 5 - CHANNEL_LAST_TX_COLLISION_OTHER = 6 - CHANNEL_OCCUPANCY_PCT_WINDOW = 7 - CHANNEL_COLLISIONS_PCT_WINDOW = 8 - ACTION_LAST = 9 - ACTION_BACK_PCT = 10 - STATUS_RETRY_COUNTER = 11 - STATUS_NO_TX_COUNTER = 12 - STATUS_UNIQUE_LTC_TX_WINDOW = 13 - - -OBS_SIZE = 14 +from ltc.sim.constants import ObsIdx, OBS_SIZE diff --git a/ltc/sim/process_output.py b/ltc/sim/process_output.py index d8ac60c..49177e2 100644 --- a/ltc/sim/process_output.py +++ b/ltc/sim/process_output.py @@ -1,7 +1,6 @@ import jax import jax.numpy as jnp -from ltc.sim.observation import ObsIdx from ltc.sim.constants import * from ltc.sim.sim import is_buffer_empty @@ -115,6 +114,7 @@ def process_output_i( tx_collision_other_now, enable_cs_tx_same_type, enable_tx_collision_other, + roll_obs, ): ret_c = obs[-1, ObsIdx.STATUS_RETRY_COUNTER].astype(jnp.int32) no_tx = obs[-1, ObsIdx.STATUS_NO_TX_COUNTER].astype(jnp.int32) @@ -170,8 +170,12 @@ def process_output_i( ], dtype=jnp.float32, ) - obs = jnp.roll(obs, -1, axis=0) - obs = obs.at[-1].set(obs_t) + obs = jax.lax.cond( + roll_obs, + lambda x: jnp.roll(x, -1, axis=0).at[-1].set(obs_t), + lambda x: x.at[-1].set(obs_t), + obs, + ) return obs, reward, power @@ -196,8 +200,9 @@ def process_output( tx_collision_other_now, enable_cs_tx_same_type, enable_tx_collision_other, + roll_obs=True, ): - return jax.vmap(process_output_i, in_axes=(0, 0, 0, 0, None, 0, 0, 0, None, None, 0, None, None, 0, None, 0, 0, None, None))( + return jax.vmap(process_output_i, in_axes=(0, 0, 0, 0, None, 0, 0, 0, None, None, 0, None, None, 0, None, 0, 0, None, None, None))( buffer_states, new_buffer_states, new_buffer_birth_states, @@ -217,4 +222,5 @@ def process_output( tx_collision_other_now, enable_cs_tx_same_type, enable_tx_collision_other, + roll_obs, ) diff --git a/ltc/utils/scan_states.py b/ltc/utils/scan_states.py index 05045f0..c403563 100644 --- a/ltc/utils/scan_states.py +++ b/ltc/utils/scan_states.py @@ -21,6 +21,7 @@ class Carry: channel_busy_hist: jax.Array collision_hist: jax.Array tx_hist: jax.Array + staged_tx: jax.Array power_states: jax.Array channel_state: int key: jax.random.PRNGKey diff --git a/test/sim/test_process_output.py b/test/sim/test_process_output.py index f4177b0..15718ae 100644 --- a/test/sim/test_process_output.py +++ b/test/sim/test_process_output.py @@ -3,7 +3,6 @@ import jax import jax.numpy as jnp -from ltc.sim.observation import ObsIdx, OBS_SIZE from ltc.sim.constants import * from ltc.sim.process_output import * diff --git a/test/test_action_space.py b/test/test_action_space.py new file mode 100644 index 0000000..c2ad5b3 --- /dev/null +++ b/test/test_action_space.py @@ -0,0 +1,50 @@ +import unittest + +import jax +import jax.numpy as jnp + +from ltc.run import decode_drl_actions, decode_legacy_actions +from ltc.sim.constants import Actions + + +class ActionSpaceDecodeTestCase(unittest.TestCase): + def test_duration_set_decode(self): + raw = jnp.array([0, 4, 5, 9], dtype=jnp.int32) + keys = jax.random.split(jax.random.PRNGKey(0), 4) + action_type, duration, staged = decode_drl_actions( + raw, jnp.zeros((4,), dtype=jnp.int32), keys, "txcs_duration_set", 5 + ) + self.assertTrue(jnp.array_equal(action_type, jnp.array([Actions.TX.value, Actions.TX.value, Actions.CS.value, Actions.CS.value]))) + self.assertTrue(jnp.array_equal(duration, jnp.array([1, 5, 1, 5]))) + self.assertTrue(jnp.array_equal(staged, jnp.zeros((4,), dtype=jnp.int32))) + + def test_stage_commit_decode(self): + raw = jnp.array([1, 2, 2], dtype=jnp.int32) # stage, commit, commit + staged = jnp.array([0, 2, 1], dtype=jnp.int32) + keys = jax.random.split(jax.random.PRNGKey(1), 3) + action_type, duration, staged_next = decode_drl_actions( + raw, staged, keys, "tx_stage_commit", 5 + ) + self.assertTrue(jnp.array_equal(action_type, jnp.array([Actions.CS.value, Actions.TX.value, Actions.TX.value]))) + self.assertTrue(jnp.array_equal(duration, jnp.array([1, 2, 1]))) + self.assertTrue(jnp.array_equal(staged_next, jnp.array([1, 0, 0]))) + + def test_geometric_decode_range(self): + raw = jnp.array([0, 9], dtype=jnp.int32) + keys = jax.random.split(jax.random.PRNGKey(2), 2) + action_type, duration, _ = decode_drl_actions( + raw, jnp.zeros((2,), dtype=jnp.int32), keys, "txcs_to_geometric_duration", 5 + ) + self.assertTrue(jnp.array_equal(action_type, jnp.array([Actions.TX.value, Actions.CS.value]))) + self.assertTrue(jnp.all(duration >= 1)) + self.assertTrue(jnp.all(duration <= 5)) + + def test_legacy_decode(self): + raw = jnp.array([Actions.TX.value, Actions.CS.value, Actions.IDLE.value], dtype=jnp.int32) + action_type, duration = decode_legacy_actions(raw, legacy_tx_duration=5) + self.assertTrue(jnp.array_equal(action_type, raw)) + self.assertTrue(jnp.array_equal(duration, jnp.array([5, 1, 1], dtype=jnp.int32))) + + +if __name__ == "__main__": + unittest.main() From 16a46f9a514a95716d3ff912b299647f137a9f01 Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Tue, 21 Apr 2026 17:03:37 +0200 Subject: [PATCH 04/24] Fix agent independence --- ltc/run.py | 299 +++++++++++++++----------------------- ltc/sim/process_output.py | 3 +- ltc/utils/scan_states.py | 3 + 3 files changed, 125 insertions(+), 180 deletions(-) diff --git a/ltc/run.py b/ltc/run.py index 7beacef..0da9710 100644 --- a/ltc/run.py +++ b/ltc/run.py @@ -190,189 +190,125 @@ def rl_step_coroutine(c, step): key, drl_keys, legacy_keys, traffic_key, reward_key = jax.random.split(c.key, 5) drl_keys = jax.random.split(drl_keys, n_drl) legacy_keys = jax.random.split(legacy_keys, n - n_drl) - traffic_slot_keys = jax.random.split(traffic_key, max_duration) - reward_slot_keys = jax.random.split(reward_key, max_duration) + traffic_keys = jax.random.split(traffic_key, n) + + ready = c.macro_remaining == 0 + drl_ready = jnp.logical_and(active[:n_drl], ready[:n_drl]) + legacy_ready = jnp.logical_and(active[n_drl:], ready[n_drl:]) drl_states, drl_actions_raw = yield drl_step( - c.drl_states, drl_keys, c.obs[:n_drl], c.actions[:n_drl], c.rewards[:n_drl], c.terminals[:n_drl], active[:n_drl] + c.drl_states, drl_keys, c.obs[:n_drl], c.actions[:n_drl], c.rewards[:n_drl], c.terminals[:n_drl], drl_ready ) legacy_states, legacy_actions_raw = legacy_step( - c.legacy_states, legacy_keys, c.obs[n_drl:], c.actions[n_drl:], c.rewards[n_drl:], c.terminals[n_drl:], active[n_drl:] + c.legacy_states, legacy_keys, c.obs[n_drl:], c.actions[n_drl:], c.rewards[n_drl:], c.terminals[n_drl:], legacy_ready ) - raw_actions = jnp.concatenate([drl_actions_raw, legacy_actions_raw]) - drl_action_types, drl_durations, staged_tx = decode_drl_actions( - drl_actions_raw, c.staged_tx[:n_drl], drl_keys, action_space_variant, max_duration + + raw_actions = c.actions + raw_actions = raw_actions.at[:n_drl].set(jnp.where(drl_ready, drl_actions_raw, c.actions[:n_drl])) + raw_actions = raw_actions.at[n_drl:].set(jnp.where(legacy_ready, legacy_actions_raw, c.actions[n_drl:])) + + drl_action_types, drl_durations, drl_staged_next = decode_drl_actions( + raw_actions[:n_drl], c.staged_tx[:n_drl], drl_keys, action_space_variant, max_duration + ) + legacy_action_types, legacy_durations = decode_legacy_actions(raw_actions[n_drl:], legacy_tx_duration) + + macro_action_types = c.macro_action_types + macro_remaining = c.macro_remaining + staged_tx = c.staged_tx + + macro_action_types = macro_action_types.at[:n_drl].set( + jnp.where(drl_ready, drl_action_types, c.macro_action_types[:n_drl]) + ) + macro_remaining = macro_remaining.at[:n_drl].set( + jnp.where(drl_ready, drl_durations, c.macro_remaining[:n_drl]) + ) + staged_tx = staged_tx.at[:n_drl].set( + jnp.where(drl_ready, drl_staged_next, c.staged_tx[:n_drl]) + ) + + macro_action_types = macro_action_types.at[n_drl:].set( + jnp.where(legacy_ready, legacy_action_types, c.macro_action_types[n_drl:]) + ) + macro_remaining = macro_remaining.at[n_drl:].set( + jnp.where(legacy_ready, legacy_durations, c.macro_remaining[n_drl:]) + ) + + slot_executing = macro_remaining > 0 + slot_actions = jnp.where(slot_executing, macro_action_types, Actions.IDLE.value) + + traffic_states, new_frames = traffic_step(c.traffic_states, traffic_keys) + ( + buffer_states, + buffer_birth_steps, + channel_state, + packet_seqs, + planned_packets, + successful_packets, + ) = simulate(c.buffer_states, c.buffer_birth_steps, new_frames, slot_actions, c.packet_seqs, step) + + is_ltc = jnp.arange(n) < n_drl + tx_mask = slot_actions == Actions.TX.value + tx_idx = jnp.argmax(tx_mask.astype(jnp.int32)) + tx_is_ltc = is_ltc[tx_idx] + cs_tx_same_type_now = (tx_is_ltc == is_ltc).astype(jnp.float32) + diff_type = is_ltc[:, None] != is_ltc[None, :] + tx_collision_other_now = jnp.any(diff_type & tx_mask[None, :], axis=1).astype(jnp.float32) + + arrival_hist = jnp.roll(c.arrival_hist, -1, axis=1).at[:, -1].set(new_frames.astype(jnp.float32)) + planned_tx_hist = jnp.roll(c.planned_tx_hist, -1, axis=1).at[:, -1].set(planned_packets.astype(jnp.float32)) + success_tx_hist = jnp.roll(c.success_tx_hist, -1, axis=1).at[:, -1].set(successful_packets.astype(jnp.float32)) + channel_busy_hist = jnp.roll(c.channel_busy_hist, -1).at[-1].set((channel_state != 0).astype(jnp.float32)) + collision_hist = jnp.roll(c.collision_hist, -1).at[-1].set((channel_state == -1).astype(jnp.float32)) + tx_hist = jnp.roll(c.tx_hist, -1, axis=0).at[-1].set(tx_mask.astype(jnp.int32)) + + arrival_valid = arrival_hist >= 0.0 + arrival_count = jnp.sum(arrival_valid, axis=1) + traffic_mean_arrival_rate = jnp.where( + arrival_count > 0, + jnp.sum(jnp.where(arrival_valid, arrival_hist, 0.0), axis=1) / arrival_count, + -1.0, + ) + + busy_valid = channel_busy_hist >= 0.0 + busy_count = jnp.sum(busy_valid) + channel_occupancy_pct_window = jnp.where( + busy_count > 0, + jnp.sum(jnp.where(busy_valid, channel_busy_hist, 0.0)) / busy_count, + -1.0, + ) + + coll_valid = collision_hist >= 0.0 + coll_count = jnp.sum(coll_valid) + channel_collisions_pct_window = jnp.where( + coll_count > 0, + jnp.sum(jnp.where(coll_valid, collision_hist, 0.0)) / coll_count, + -1.0, + ) + + planned_valid = planned_tx_hist >= 0.0 + planned_sum = jnp.sum(jnp.where(planned_valid, planned_tx_hist, 0.0), axis=1) + success_sum = jnp.sum(jnp.where(planned_valid, success_tx_hist, 0.0), axis=1) + back_pct = jnp.where(planned_sum > 0, success_sum / planned_sum, -1.0) + + tx_valid_rows = tx_hist[:, 0] >= 0 + unique_ltc_tx_window = jnp.where( + jnp.any(tx_valid_rows), + jnp.sum(jnp.any(jnp.where(tx_valid_rows[:, None], tx_hist[:, :n_drl], 0) > 0, axis=0).astype(jnp.float32)), + -1.0, + ) + + done_now = jnp.logical_and(slot_executing, macro_remaining == 1) + obs, slot_rewards, powers = process_output( + c.buffer_states, buffer_states, buffer_birth_steps, c.power_states, channel_state, c.obs, slot_actions, + c.terminals, reward_key, step, traffic_mean_arrival_rate, channel_occupancy_pct_window, + channel_collisions_pct_window, back_pct, unique_ltc_tx_window, cs_tx_same_type_now, + tx_collision_other_now, obs_enable_cs_tx_same_type, obs_enable_tx_collision_other, roll_obs=done_now ) - legacy_action_types, legacy_durations = decode_legacy_actions(legacy_actions_raw, legacy_tx_duration) - action_types = jnp.concatenate([drl_action_types, legacy_action_types]) - durations = jnp.concatenate([drl_durations, legacy_durations]) - max_macro = jnp.max(durations) - staged_tx = jnp.concatenate([staged_tx, c.staged_tx[n_drl:]]) - - traffic_states = c.traffic_states - buffer_states = c.buffer_states - buffer_birth_steps = c.buffer_birth_steps - packet_seqs = c.packet_seqs - obs = c.obs - powers = c.power_states - rewards = jnp.zeros_like(c.rewards) - arrival_hist = c.arrival_hist - planned_tx_hist = c.planned_tx_hist - success_tx_hist = c.success_tx_hist - channel_busy_hist = c.channel_busy_hist - collision_hist = c.collision_hist - tx_hist = c.tx_hist - channel_state = c.channel_state - last_new_frames = jnp.zeros((n,), dtype=jnp.int32) - - for inner in range(max_duration): - do_slot = inner < max_macro - - def step_slot(args): - ( - traffic_states, - buffer_states, - buffer_birth_steps, - packet_seqs, - obs, - powers, - rewards, - arrival_hist, - planned_tx_hist, - success_tx_hist, - channel_busy_hist, - collision_hist, - tx_hist, - channel_state, - last_new_frames, - ) = args - slot_traffic_keys = jax.random.split(traffic_slot_keys[inner], n) - traffic_states, new_frames = traffic_step(traffic_states, slot_traffic_keys) - slot_actions = jnp.where(inner < durations, action_types, Actions.IDLE.value) - ( - new_buffer_states, - new_buffer_birth_steps, - channel_state, - packet_seqs, - planned_packets, - successful_packets, - ) = simulate(buffer_states, buffer_birth_steps, new_frames, slot_actions, packet_seqs, step * max_duration + inner) - - is_ltc = jnp.arange(n) < n_drl - tx_mask = slot_actions == Actions.TX.value - tx_idx = jnp.argmax(tx_mask.astype(jnp.int32)) - tx_is_ltc = is_ltc[tx_idx] - cs_tx_same_type_now = (tx_is_ltc == is_ltc).astype(jnp.float32) - diff_type = is_ltc[:, None] != is_ltc[None, :] - tx_collision_other_now = jnp.any(diff_type & tx_mask[None, :], axis=1).astype(jnp.float32) - - arrival_hist = jnp.roll(arrival_hist, -1, axis=1).at[:, -1].set(new_frames.astype(jnp.float32)) - planned_tx_hist = jnp.roll(planned_tx_hist, -1, axis=1).at[:, -1].set(planned_packets.astype(jnp.float32)) - success_tx_hist = jnp.roll(success_tx_hist, -1, axis=1).at[:, -1].set(successful_packets.astype(jnp.float32)) - channel_busy_hist = jnp.roll(channel_busy_hist, -1).at[-1].set((channel_state != 0).astype(jnp.float32)) - collision_hist = jnp.roll(collision_hist, -1).at[-1].set((channel_state == -1).astype(jnp.float32)) - tx_hist = jnp.roll(tx_hist, -1, axis=0).at[-1].set(tx_mask.astype(jnp.int32)) - - arrival_valid = arrival_hist >= 0.0 - arrival_count = jnp.sum(arrival_valid, axis=1) - traffic_mean_arrival_rate = jnp.where( - arrival_count > 0, - jnp.sum(jnp.where(arrival_valid, arrival_hist, 0.0), axis=1) / arrival_count, - -1.0, - ) - - busy_valid = channel_busy_hist >= 0.0 - busy_count = jnp.sum(busy_valid) - channel_occupancy_pct_window = jnp.where( - busy_count > 0, - jnp.sum(jnp.where(busy_valid, channel_busy_hist, 0.0)) / busy_count, - -1.0, - ) - - coll_valid = collision_hist >= 0.0 - coll_count = jnp.sum(coll_valid) - channel_collisions_pct_window = jnp.where( - coll_count > 0, - jnp.sum(jnp.where(coll_valid, collision_hist, 0.0)) / coll_count, - -1.0, - ) - - planned_valid = planned_tx_hist >= 0.0 - planned_sum = jnp.sum(jnp.where(planned_valid, planned_tx_hist, 0.0), axis=1) - success_sum = jnp.sum(jnp.where(planned_valid, success_tx_hist, 0.0), axis=1) - back_pct = jnp.where(planned_sum > 0, success_sum / planned_sum, -1.0) - - tx_valid_rows = tx_hist[:, 0] >= 0 - unique_ltc_tx_window = jnp.where( - jnp.any(tx_valid_rows), - jnp.sum(jnp.any(jnp.where(tx_valid_rows[:, None], tx_hist[:, :n_drl], 0) > 0, axis=0).astype(jnp.float32)), - -1.0, - ) - - obs, slot_rewards, powers = process_output( - buffer_states, new_buffer_states, new_buffer_birth_steps, powers, channel_state, obs, slot_actions, - c.terminals, reward_slot_keys[inner], step * max_duration + inner, traffic_mean_arrival_rate, channel_occupancy_pct_window, - channel_collisions_pct_window, back_pct, unique_ltc_tx_window, cs_tx_same_type_now, - tx_collision_other_now, obs_enable_cs_tx_same_type, obs_enable_tx_collision_other, roll_obs=(inner + 1 == max_macro) - ) - rewards = rewards + slot_rewards - return ( - traffic_states, - new_buffer_states, - new_buffer_birth_steps, - packet_seqs, - obs, - powers, - rewards, - arrival_hist, - planned_tx_hist, - success_tx_hist, - channel_busy_hist, - collision_hist, - tx_hist, - channel_state, - new_frames, - ) - - ( - traffic_states, - buffer_states, - buffer_birth_steps, - packet_seqs, - obs, - powers, - rewards, - arrival_hist, - planned_tx_hist, - success_tx_hist, - channel_busy_hist, - collision_hist, - tx_hist, - channel_state, - last_new_frames, - ) = jax.lax.cond( - do_slot, - step_slot, - lambda x: x, - ( - traffic_states, - buffer_states, - buffer_birth_steps, - packet_seqs, - obs, - powers, - rewards, - arrival_hist, - planned_tx_hist, - success_tx_hist, - channel_busy_hist, - collision_hist, - tx_hist, - channel_state, - last_new_frames, - ), - ) + macro_reward_accum = c.macro_reward_accum + jnp.where(slot_executing, slot_rewards, 0.0) + rewards = jnp.where(done_now, macro_reward_accum, c.rewards) + macro_reward_accum = jnp.where(done_now, 0.0, macro_reward_accum) + macro_remaining = jnp.where(slot_executing, macro_remaining - 1, macro_remaining) terminals = jnp.logical_or(c.terminals, powers < 0) if n_drl > 0: @@ -386,12 +322,13 @@ def step_slot(args): c = Carry( drl_states, legacy_states, traffic_states, packet_seqs, buffer_states, buffer_birth_steps, - arrival_hist, planned_tx_hist, success_tx_hist, channel_busy_hist, collision_hist, tx_hist, staged_tx, powers, + arrival_hist, planned_tx_hist, success_tx_hist, channel_busy_hist, collision_hist, tx_hist, staged_tx, + macro_remaining, macro_action_types, macro_reward_accum, powers, channel_state, key, obs, raw_actions, rewards, terminals, active ) o = Output( legacy_states, obs, raw_actions, rewards, terminals, buffer_states, powers, - last_new_frames, channel_state, active, hist, bin_edges + new_frames, channel_state, active, hist, bin_edges ) yield c, o @@ -518,6 +455,9 @@ def setup_args(): collision_hist = jnp.full(obs_stats_window, -1.0, dtype=jnp.float32) tx_hist = jnp.full((obs_stats_window, n), -1, dtype=jnp.int32) staged_tx = jnp.zeros(n, dtype=jnp.int32) + macro_remaining = jnp.zeros(n, dtype=jnp.int32) + macro_action_types = jnp.full(n, Actions.IDLE.value, dtype=jnp.int32) + macro_reward_accum = jnp.zeros(n, dtype=jnp.float32) power_states = jnp.full(n, INITIAL_CAPACITY, dtype=int) channel_state = 0 obs = jnp.zeros((n, window_size, OBS_SIZE), dtype=jnp.float32) @@ -590,7 +530,8 @@ def setup_args(): rl_step_fn = jax.jit(rl_step_fn) carry = Carry( drl_states, legacy_states, traffic_states, packet_seqs, buffer_states, buffer_birth_steps, - arrival_hist, planned_tx_hist, success_tx_hist, channel_busy_hist, collision_hist, tx_hist, staged_tx, power_states, + arrival_hist, planned_tx_hist, success_tx_hist, channel_busy_hist, collision_hist, tx_hist, staged_tx, + macro_remaining, macro_action_types, macro_reward_accum, power_states, channel_state, key, obs, actions, rewards, terminals, active ) all_outputs = [] diff --git a/ltc/sim/process_output.py b/ltc/sim/process_output.py index 49177e2..e131acd 100644 --- a/ltc/sim/process_output.py +++ b/ltc/sim/process_output.py @@ -202,7 +202,8 @@ def process_output( enable_tx_collision_other, roll_obs=True, ): - return jax.vmap(process_output_i, in_axes=(0, 0, 0, 0, None, 0, 0, 0, None, None, 0, None, None, 0, None, 0, 0, None, None, None))( + roll_obs = jnp.broadcast_to(jnp.asarray(roll_obs), actions.shape) + return jax.vmap(process_output_i, in_axes=(0, 0, 0, 0, None, 0, 0, 0, None, None, 0, None, None, 0, None, 0, 0, None, None, 0))( buffer_states, new_buffer_states, new_buffer_birth_states, diff --git a/ltc/utils/scan_states.py b/ltc/utils/scan_states.py index c403563..62d01f9 100644 --- a/ltc/utils/scan_states.py +++ b/ltc/utils/scan_states.py @@ -22,6 +22,9 @@ class Carry: collision_hist: jax.Array tx_hist: jax.Array staged_tx: jax.Array + macro_remaining: jax.Array + macro_action_types: jax.Array + macro_reward_accum: jax.Array power_states: jax.Array channel_state: int key: jax.random.PRNGKey From df7601188f5fbd5e76db5e5674f24d1d19187c2f Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Tue, 21 Apr 2026 21:16:08 +0200 Subject: [PATCH 05/24] Update reward --- ltc/run.py | 65 ++++++++++++++++++++++++++++++++++++--- ltc/sim/constants.py | 10 ++++++ ltc/utils/scan_states.py | 2 ++ test/test_action_space.py | 20 +++++++++++- 4 files changed, 92 insertions(+), 5 deletions(-) diff --git a/ltc/run.py b/ltc/run.py index 0da9710..4e48dce 100644 --- a/ltc/run.py +++ b/ltc/run.py @@ -21,6 +21,17 @@ OBS_SIZE, MAX_MACRO_DURATION, LEGACY_TX_DURATION, + TX_REWARD_LENGTH, + TX_SIZE_THRESHOLD, + TX_SIZE_PENALTY_WINDOW, + TX_SIZE_PENALTY, + TX_LTC_COLLISION_PENALTY, + TX_COEX_COLLISION_PENALTY, + TX_EMPTY_BUFFER_PENALTY, + TX_MAX_RETRANSMISSION_PENALTY, + MAX_RETRANSMISSION, + ObsIdx, + TX_REWARD, ) from ltc.utils.history import build_history_filename, ensure_clean_git_worktree, get_short_commit_hash from ltc.utils.scan_states import Carry, Output @@ -167,6 +178,32 @@ def decode_legacy_actions(raw_actions, legacy_tx_duration): return action_type, duration +def upgraded_tx_slot_reward(slot_actions, channel_state, successful_packets, prev_ret_c, tx_collision_other_now): + tx_executing = slot_actions == Actions.TX.value + tx_success = jnp.logical_and(tx_executing, successful_packets > 0) + tx_collision = jnp.logical_and(tx_executing, channel_state == -1) + tx_collision_max = jnp.logical_and(tx_collision, prev_ret_c >= MAX_RETRANSMISSION) + tx_collision_normal = jnp.logical_and(tx_collision, prev_ret_c < MAX_RETRANSMISSION) + tx_collision_coex = jnp.logical_and(tx_collision_normal, tx_collision_other_now > 0.5) + tx_collision_ltc = jnp.logical_and(tx_collision_normal, tx_collision_other_now <= 0.5) + tx_empty = jnp.logical_and( + tx_executing, + jnp.logical_and(channel_state == 1, successful_packets == 0), + ) + + k_tx_slot = successful_packets.astype(jnp.float32) + k_coll_slot = tx_collision.astype(jnp.float32) + + tx_success_reward = TX_REWARD * (k_tx_slot * TX_REWARD_LENGTH - 2.0) / TX_REWARD_LENGTH + tx_collision_reward = ( + tx_collision_ltc.astype(jnp.float32) * (TX_LTC_COLLISION_PENALTY * k_coll_slot) + + tx_collision_coex.astype(jnp.float32) * (TX_COEX_COLLISION_PENALTY * k_coll_slot) + + tx_collision_max.astype(jnp.float32) * TX_MAX_RETRANSMISSION_PENALTY + ) + tx_slot_reward = jnp.where(tx_success, tx_success_reward, 0.0) + tx_collision_reward + jnp.where(tx_empty, TX_EMPTY_BUFFER_PENALTY, 0.0) + return tx_slot_reward, k_tx_slot, k_coll_slot, tx_executing + + def rl_step( drl_step, legacy_step, traffic_step, n, n_drl, n_bins=50, one_shot_step=None, one_shot_target=None, station_change_interval=None, station_change_delta=0, @@ -298,6 +335,7 @@ def rl_step_coroutine(c, step): ) done_now = jnp.logical_and(slot_executing, macro_remaining == 1) + prev_ret_c = c.obs[:, -1, ObsIdx.STATUS_RETRY_COUNTER].astype(jnp.int32) obs, slot_rewards, powers = process_output( c.buffer_states, buffer_states, buffer_birth_steps, c.power_states, channel_state, c.obs, slot_actions, c.terminals, reward_key, step, traffic_mean_arrival_rate, channel_occupancy_pct_window, @@ -305,9 +343,26 @@ def rl_step_coroutine(c, step): tx_collision_other_now, obs_enable_cs_tx_same_type, obs_enable_tx_collision_other, roll_obs=done_now ) - macro_reward_accum = c.macro_reward_accum + jnp.where(slot_executing, slot_rewards, 0.0) - rewards = jnp.where(done_now, macro_reward_accum, c.rewards) + tx_slot_reward, k_tx_slot, k_coll_slot, tx_executing = upgraded_tx_slot_reward( + slot_actions, channel_state, successful_packets, prev_ret_c, tx_collision_other_now + ) + + # TODO: Implement "Else: max 6 ms" clause once precise behavior is finalized. + slot_reward_upgraded = jnp.where(tx_executing, tx_slot_reward, slot_rewards) + + macro_tx_success_accum = c.macro_tx_success_accum + jnp.where(slot_executing, k_tx_slot, 0.0) + macro_tx_collision_accum = c.macro_tx_collision_accum + jnp.where(slot_executing, k_coll_slot, 0.0) + macro_reward_accum = c.macro_reward_accum + jnp.where(slot_executing, slot_reward_upgraded, 0.0) + + size_penalty = TX_SIZE_PENALTY * jnp.minimum( + 1.0, (macro_tx_success_accum - TX_SIZE_THRESHOLD + 1.0) / TX_SIZE_PENALTY_WINDOW + ) + size_penalty = jnp.where(macro_tx_success_accum >= TX_SIZE_THRESHOLD, size_penalty, 0.0) + rewards = jnp.where(done_now, macro_reward_accum + size_penalty, c.rewards) + macro_reward_accum = jnp.where(done_now, 0.0, macro_reward_accum) + macro_tx_success_accum = jnp.where(done_now, 0.0, macro_tx_success_accum) + macro_tx_collision_accum = jnp.where(done_now, 0.0, macro_tx_collision_accum) macro_remaining = jnp.where(slot_executing, macro_remaining - 1, macro_remaining) terminals = jnp.logical_or(c.terminals, powers < 0) @@ -323,7 +378,7 @@ def rl_step_coroutine(c, step): c = Carry( drl_states, legacy_states, traffic_states, packet_seqs, buffer_states, buffer_birth_steps, arrival_hist, planned_tx_hist, success_tx_hist, channel_busy_hist, collision_hist, tx_hist, staged_tx, - macro_remaining, macro_action_types, macro_reward_accum, powers, + macro_remaining, macro_action_types, macro_reward_accum, macro_tx_success_accum, macro_tx_collision_accum, powers, channel_state, key, obs, raw_actions, rewards, terminals, active ) o = Output( @@ -458,6 +513,8 @@ def setup_args(): macro_remaining = jnp.zeros(n, dtype=jnp.int32) macro_action_types = jnp.full(n, Actions.IDLE.value, dtype=jnp.int32) macro_reward_accum = jnp.zeros(n, dtype=jnp.float32) + macro_tx_success_accum = jnp.zeros(n, dtype=jnp.float32) + macro_tx_collision_accum = jnp.zeros(n, dtype=jnp.float32) power_states = jnp.full(n, INITIAL_CAPACITY, dtype=int) channel_state = 0 obs = jnp.zeros((n, window_size, OBS_SIZE), dtype=jnp.float32) @@ -531,7 +588,7 @@ def setup_args(): carry = Carry( drl_states, legacy_states, traffic_states, packet_seqs, buffer_states, buffer_birth_steps, arrival_hist, planned_tx_hist, success_tx_hist, channel_busy_hist, collision_hist, tx_hist, staged_tx, - macro_remaining, macro_action_types, macro_reward_accum, power_states, + macro_remaining, macro_action_types, macro_reward_accum, macro_tx_success_accum, macro_tx_collision_accum, power_states, channel_state, key, obs, actions, rewards, terminals, active ) all_outputs = [] diff --git a/ltc/sim/constants.py b/ltc/sim/constants.py index 42a8f88..90c9f4d 100644 --- a/ltc/sim/constants.py +++ b/ltc/sim/constants.py @@ -29,6 +29,16 @@ class Actions(Enum): COLLISION_PENALTY = -1.0 MAX_RETRANSMISSION_PENALTY = -1.0 +# Upgraded TX reward parameters +TX_REWARD_LENGTH = 5 +TX_SIZE_THRESHOLD = 5 +TX_SIZE_PENALTY_WINDOW = 25 +TX_SIZE_PENALTY = -1.0 +TX_LTC_COLLISION_PENALTY = -1.0 +TX_COEX_COLLISION_PENALTY = -5.0 +TX_EMPTY_BUFFER_PENALTY = EMPTY_TX_PENALTY +TX_MAX_RETRANSMISSION_PENALTY = MAX_RETRANSMISSION_PENALTY + """ Power consumption """ diff --git a/ltc/utils/scan_states.py b/ltc/utils/scan_states.py index 62d01f9..9b75856 100644 --- a/ltc/utils/scan_states.py +++ b/ltc/utils/scan_states.py @@ -25,6 +25,8 @@ class Carry: macro_remaining: jax.Array macro_action_types: jax.Array macro_reward_accum: jax.Array + macro_tx_success_accum: jax.Array + macro_tx_collision_accum: jax.Array power_states: jax.Array channel_state: int key: jax.random.PRNGKey diff --git a/test/test_action_space.py b/test/test_action_space.py index c2ad5b3..e3dbe49 100644 --- a/test/test_action_space.py +++ b/test/test_action_space.py @@ -3,7 +3,7 @@ import jax import jax.numpy as jnp -from ltc.run import decode_drl_actions, decode_legacy_actions +from ltc.run import decode_drl_actions, decode_legacy_actions, upgraded_tx_slot_reward from ltc.sim.constants import Actions @@ -45,6 +45,24 @@ def test_legacy_decode(self): self.assertTrue(jnp.array_equal(action_type, raw)) self.assertTrue(jnp.array_equal(duration, jnp.array([5, 1, 1], dtype=jnp.int32))) + def test_upgraded_tx_slot_reward_branches(self): + slot_actions = jnp.array([Actions.TX.value, Actions.TX.value, Actions.TX.value, Actions.TX.value], dtype=jnp.int32) + channel_state = jnp.array([1, -1, -1, 1], dtype=jnp.int32) + successful_packets = jnp.array([1, 0, 0, 0], dtype=jnp.int32) + prev_ret_c = jnp.array([0, 0, 8, 0], dtype=jnp.int32) + tx_collision_other_now = jnp.array([0.0, 1.0, 0.0, 0.0], dtype=jnp.float32) + rewards, k_tx, k_coll, tx_exec = upgraded_tx_slot_reward( + slot_actions, channel_state, successful_packets, prev_ret_c, tx_collision_other_now + ) + self.assertTrue(jnp.array_equal(tx_exec, jnp.ones((4,), dtype=bool))) + self.assertTrue(jnp.array_equal(k_tx, jnp.array([1.0, 0.0, 0.0, 0.0], dtype=jnp.float32))) + self.assertTrue(jnp.array_equal(k_coll, jnp.array([0.0, 1.0, 1.0, 0.0], dtype=jnp.float32))) + # success, coex collision, max retransmission, empty-buffer TX + self.assertAlmostEqual(float(rewards[0]), 0.6, places=6) + self.assertAlmostEqual(float(rewards[1]), -5.0, places=6) + self.assertAlmostEqual(float(rewards[2]), -1.0, places=6) + self.assertAlmostEqual(float(rewards[3]), -0.5, places=6) + if __name__ == "__main__": unittest.main() From 74e56bc4c954fb8d769e35a9eb00ab77ab32a23f Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Tue, 21 Apr 2026 21:24:24 +0200 Subject: [PATCH 06/24] Refactor code modified: ltc/run.py modified: ltc/sim/process_output.py modified: ltc/utils/scan_states.py new file: ltc/utils/state_structs.py --- ltc/run.py | 152 +++++++++++++++++++++++++++---------- ltc/sim/process_output.py | 129 ++++++++++++++++++++++--------- ltc/utils/scan_states.py | 15 +--- ltc/utils/state_structs.py | 50 ++++++++++++ 4 files changed, 259 insertions(+), 87 deletions(-) create mode 100644 ltc/utils/state_structs.py diff --git a/ltc/run.py b/ltc/run.py index 4e48dce..8f7345f 100644 --- a/ltc/run.py +++ b/ltc/run.py @@ -35,6 +35,7 @@ ) from ltc.utils.history import build_history_filename, ensure_clean_git_worktree, get_short_commit_hash from ltc.utils.scan_states import Carry, Output +from ltc.utils.state_structs import MacroActionState, ObsTrackerState from ltc.utils.plots import plot_all, plot_first @@ -178,6 +179,59 @@ def decode_legacy_actions(raw_actions, legacy_tx_duration): return action_type, duration +# ============================================================================ +# Macro-Action State Management +# ============================================================================ + +def _update_macro_state_from_ready_agents( + macro_state, obs_tracker, ready_drl, ready_legacy, + drl_action_types, drl_durations, drl_staged_next, + legacy_action_types, legacy_durations, + n_drl +): + """Update macro-action state for agents that are ready to choose new actions. + + Ready DRL agents: update with decoded DRL actions + Ready legacy agents: update with decoded legacy actions + Non-ready agents: keep existing macro state + """ + macro_action_types = macro_state.action_types + macro_remaining = macro_state.remaining + staged_tx = obs_tracker.staged_tx + + # DRL agents + macro_action_types = macro_action_types.at[:n_drl].set( + jnp.where(ready_drl, drl_action_types, macro_state.action_types[:n_drl]) + ) + macro_remaining = macro_remaining.at[:n_drl].set( + jnp.where(ready_drl, drl_durations, macro_state.remaining[:n_drl]) + ) + staged_tx = staged_tx.at[:n_drl].set( + jnp.where(ready_drl, drl_staged_next, obs_tracker.staged_tx[:n_drl]) + ) + + # Legacy agents + macro_action_types = macro_action_types.at[n_drl:].set( + jnp.where(ready_legacy, legacy_action_types, macro_state.action_types[n_drl:]) + ) + macro_remaining = macro_remaining.at[n_drl:].set( + jnp.where(ready_legacy, legacy_durations, macro_state.remaining[n_drl:]) + ) + + return macro_action_types, macro_remaining, staged_tx + + +def _finalize_macro_accumulators(macro_state, done_now, macro_tx_success_accum, macro_tx_collision_accum, macro_reward_accum): + """Reset macro accumulators for agents completing their macro-action. + + Agents that are done_now get their accumulators zeroed for the next macro-action. + """ + macro_tx_success_accum = jnp.where(done_now, 0.0, macro_tx_success_accum) + macro_tx_collision_accum = jnp.where(done_now, 0.0, macro_tx_collision_accum) + macro_reward_accum = jnp.where(done_now, 0.0, macro_reward_accum) + return macro_tx_success_accum, macro_tx_collision_accum, macro_reward_accum + + def upgraded_tx_slot_reward(slot_actions, channel_state, successful_packets, prev_ret_c, tx_collision_other_now): tx_executing = slot_actions == Actions.TX.value tx_success = jnp.logical_and(tx_executing, successful_packets > 0) @@ -229,7 +283,7 @@ def rl_step_coroutine(c, step): legacy_keys = jax.random.split(legacy_keys, n - n_drl) traffic_keys = jax.random.split(traffic_key, n) - ready = c.macro_remaining == 0 + ready = c.macro.remaining == 0 drl_ready = jnp.logical_and(active[:n_drl], ready[:n_drl]) legacy_ready = jnp.logical_and(active[n_drl:], ready[n_drl:]) @@ -245,29 +299,15 @@ def rl_step_coroutine(c, step): raw_actions = raw_actions.at[n_drl:].set(jnp.where(legacy_ready, legacy_actions_raw, c.actions[n_drl:])) drl_action_types, drl_durations, drl_staged_next = decode_drl_actions( - raw_actions[:n_drl], c.staged_tx[:n_drl], drl_keys, action_space_variant, max_duration + raw_actions[:n_drl], c.obs_tracker.staged_tx[:n_drl], drl_keys, action_space_variant, max_duration ) legacy_action_types, legacy_durations = decode_legacy_actions(raw_actions[n_drl:], legacy_tx_duration) - macro_action_types = c.macro_action_types - macro_remaining = c.macro_remaining - staged_tx = c.staged_tx - - macro_action_types = macro_action_types.at[:n_drl].set( - jnp.where(drl_ready, drl_action_types, c.macro_action_types[:n_drl]) - ) - macro_remaining = macro_remaining.at[:n_drl].set( - jnp.where(drl_ready, drl_durations, c.macro_remaining[:n_drl]) - ) - staged_tx = staged_tx.at[:n_drl].set( - jnp.where(drl_ready, drl_staged_next, c.staged_tx[:n_drl]) - ) - - macro_action_types = macro_action_types.at[n_drl:].set( - jnp.where(legacy_ready, legacy_action_types, c.macro_action_types[n_drl:]) - ) - macro_remaining = macro_remaining.at[n_drl:].set( - jnp.where(legacy_ready, legacy_durations, c.macro_remaining[n_drl:]) + macro_action_types, macro_remaining, staged_tx = _update_macro_state_from_ready_agents( + c.macro, c.obs_tracker, drl_ready, legacy_ready, + drl_action_types, drl_durations, drl_staged_next, + legacy_action_types, legacy_durations, + n_drl ) slot_executing = macro_remaining > 0 @@ -281,7 +321,7 @@ def rl_step_coroutine(c, step): packet_seqs, planned_packets, successful_packets, - ) = simulate(c.buffer_states, c.buffer_birth_steps, new_frames, slot_actions, c.packet_seqs, step) + ) = simulate(c.buffer_states, c.obs_tracker.buffer_birth_steps, new_frames, slot_actions, c.packet_seqs, step) is_ltc = jnp.arange(n) < n_drl tx_mask = slot_actions == Actions.TX.value @@ -291,11 +331,11 @@ def rl_step_coroutine(c, step): diff_type = is_ltc[:, None] != is_ltc[None, :] tx_collision_other_now = jnp.any(diff_type & tx_mask[None, :], axis=1).astype(jnp.float32) - arrival_hist = jnp.roll(c.arrival_hist, -1, axis=1).at[:, -1].set(new_frames.astype(jnp.float32)) - planned_tx_hist = jnp.roll(c.planned_tx_hist, -1, axis=1).at[:, -1].set(planned_packets.astype(jnp.float32)) - success_tx_hist = jnp.roll(c.success_tx_hist, -1, axis=1).at[:, -1].set(successful_packets.astype(jnp.float32)) - channel_busy_hist = jnp.roll(c.channel_busy_hist, -1).at[-1].set((channel_state != 0).astype(jnp.float32)) - collision_hist = jnp.roll(c.collision_hist, -1).at[-1].set((channel_state == -1).astype(jnp.float32)) + arrival_hist = jnp.roll(c.obs_tracker.arrival_hist, -1, axis=1).at[:, -1].set(new_frames.astype(jnp.float32)) + planned_tx_hist = jnp.roll(c.obs_tracker.planned_tx_hist, -1, axis=1).at[:, -1].set(planned_packets.astype(jnp.float32)) + success_tx_hist = jnp.roll(c.obs_tracker.success_tx_hist, -1, axis=1).at[:, -1].set(successful_packets.astype(jnp.float32)) + channel_busy_hist = jnp.roll(c.obs_tracker.channel_busy_hist, -1).at[-1].set((channel_state != 0).astype(jnp.float32)) + collision_hist = jnp.roll(c.obs_tracker.collision_hist, -1).at[-1].set((channel_state == -1).astype(jnp.float32)) tx_hist = jnp.roll(c.tx_hist, -1, axis=0).at[-1].set(tx_mask.astype(jnp.int32)) arrival_valid = arrival_hist >= 0.0 @@ -350,9 +390,9 @@ def rl_step_coroutine(c, step): # TODO: Implement "Else: max 6 ms" clause once precise behavior is finalized. slot_reward_upgraded = jnp.where(tx_executing, tx_slot_reward, slot_rewards) - macro_tx_success_accum = c.macro_tx_success_accum + jnp.where(slot_executing, k_tx_slot, 0.0) - macro_tx_collision_accum = c.macro_tx_collision_accum + jnp.where(slot_executing, k_coll_slot, 0.0) - macro_reward_accum = c.macro_reward_accum + jnp.where(slot_executing, slot_reward_upgraded, 0.0) + macro_tx_success_accum = c.macro.tx_success_accum + jnp.where(slot_executing, k_tx_slot, 0.0) + macro_tx_collision_accum = c.macro.tx_collision_accum + jnp.where(slot_executing, k_coll_slot, 0.0) + macro_reward_accum = c.macro.reward_accum + jnp.where(slot_executing, slot_reward_upgraded, 0.0) size_penalty = TX_SIZE_PENALTY * jnp.minimum( 1.0, (macro_tx_success_accum - TX_SIZE_THRESHOLD + 1.0) / TX_SIZE_PENALTY_WINDOW @@ -360,9 +400,9 @@ def rl_step_coroutine(c, step): size_penalty = jnp.where(macro_tx_success_accum >= TX_SIZE_THRESHOLD, size_penalty, 0.0) rewards = jnp.where(done_now, macro_reward_accum + size_penalty, c.rewards) - macro_reward_accum = jnp.where(done_now, 0.0, macro_reward_accum) - macro_tx_success_accum = jnp.where(done_now, 0.0, macro_tx_success_accum) - macro_tx_collision_accum = jnp.where(done_now, 0.0, macro_tx_collision_accum) + macro_tx_success_accum, macro_tx_collision_accum, macro_reward_accum = _finalize_macro_accumulators( + c.macro, done_now, macro_tx_success_accum, macro_tx_collision_accum, macro_reward_accum + ) macro_remaining = jnp.where(slot_executing, macro_remaining - 1, macro_remaining) terminals = jnp.logical_or(c.terminals, powers < 0) @@ -375,11 +415,26 @@ def rl_step_coroutine(c, step): else: hist, bin_edges = None, None + macro_state = MacroActionState( + remaining=macro_remaining, + action_types=macro_action_types, + reward_accum=macro_reward_accum, + tx_success_accum=macro_tx_success_accum, + tx_collision_accum=macro_tx_collision_accum, + ) + obs_tracker_state = ObsTrackerState( + buffer_birth_steps=c.obs_tracker.buffer_birth_steps, + arrival_hist=arrival_hist, + planned_tx_hist=planned_tx_hist, + success_tx_hist=success_tx_hist, + channel_busy_hist=channel_busy_hist, + collision_hist=collision_hist, + staged_tx=staged_tx, + ) c = Carry( - drl_states, legacy_states, traffic_states, packet_seqs, buffer_states, buffer_birth_steps, - arrival_hist, planned_tx_hist, success_tx_hist, channel_busy_hist, collision_hist, tx_hist, staged_tx, - macro_remaining, macro_action_types, macro_reward_accum, macro_tx_success_accum, macro_tx_collision_accum, powers, - channel_state, key, obs, raw_actions, rewards, terminals, active + drl_states, legacy_states, traffic_states, packet_seqs, buffer_states, + tx_hist, powers, channel_state, key, obs, raw_actions, rewards, terminals, active, + macro_state, obs_tracker_state ) o = Output( legacy_states, obs, raw_actions, rewards, terminals, buffer_states, powers, @@ -585,11 +640,26 @@ def setup_args(): legacy_tx_duration=LEGACY_TX_DURATION, ) rl_step_fn = jax.jit(rl_step_fn) + macro_state = MacroActionState( + remaining=macro_remaining, + action_types=macro_action_types, + reward_accum=macro_reward_accum, + tx_success_accum=macro_tx_success_accum, + tx_collision_accum=macro_tx_collision_accum, + ) + obs_tracker_state = ObsTrackerState( + buffer_birth_steps=buffer_birth_steps, + arrival_hist=arrival_hist, + planned_tx_hist=planned_tx_hist, + success_tx_hist=success_tx_hist, + channel_busy_hist=channel_busy_hist, + collision_hist=collision_hist, + staged_tx=staged_tx, + ) carry = Carry( - drl_states, legacy_states, traffic_states, packet_seqs, buffer_states, buffer_birth_steps, - arrival_hist, planned_tx_hist, success_tx_hist, channel_busy_hist, collision_hist, tx_hist, staged_tx, - macro_remaining, macro_action_types, macro_reward_accum, macro_tx_success_accum, macro_tx_collision_accum, power_states, - channel_state, key, obs, actions, rewards, terminals, active + drl_states, legacy_states, traffic_states, packet_seqs, buffer_states, + tx_hist, power_states, channel_state, key, obs, actions, rewards, terminals, active, + macro_state, obs_tracker_state ) all_outputs = [] diff --git a/ltc/sim/process_output.py b/ltc/sim/process_output.py index e131acd..8bb369f 100644 --- a/ltc/sim/process_output.py +++ b/ltc/sim/process_output.py @@ -5,6 +5,39 @@ from ltc.sim.sim import is_buffer_empty +# ============================================================================ +# Reward Computation Helpers +# ============================================================================ + +def _compute_cs_reward(buffer_state, channel_state, no_tx, ret_c, key): + """Compute reward for CS (channel sensing) action. + + CS reward branches on whether buffer is empty or full: + - Empty buffer: EMPTY_BUFFER_REWARD + - Full buffer with short idle: NO_TX_REWARD + - Full buffer with long idle: NO_TX_PENALTY (scaled) + """ + return jax.lax.cond(is_buffer_empty(buffer_state), idle_empty_buffer, idle_full_buffer, + (buffer_state, ret_c, no_tx, key)) + + +def _compute_tx_reward(buffer_state, channel_state, ret_c, no_tx, key): + """Compute reward for TX (transmission) action. + + TX reward branches on channel outcome: + - Collision with collision_count < MAX_RETRANSMISSION: COLLISION_PENALTY, ret_c+1 + - Collision with collision_count >= MAX_RETRANSMISSION: MAX_RETRANSMISSION_PENALTY, reset + - Success with empty buffer: EMPTY_TX_PENALTY + - Success with full buffer: TX_REWARD, ret_c=0 + """ + args = (buffer_state, ret_c, channel_state, no_tx, key) + return jax.lax.cond(channel_state == 1, transmission_without_collision, transmission_with_collision, args) + + +# ============================================================================ +# Legacy Reward Branch Functions (organized for clarity) +# ============================================================================ + def no_transmission(args): _, buffer_state, _, _, _, _ = args return jax.lax.cond(is_buffer_empty(buffer_state), idle_empty_buffer, idle_full_buffer, args) @@ -94,6 +127,61 @@ def _oldest_packet_age_norm(buffer_state, buffer_birth_state, current_step): return jnp.where(age >= 0, age.astype(jnp.float32) / queue_size, -1.0) +# ============================================================================ +# Observation Construction +# ============================================================================ + +def _build_observation_entry( + buffer_state, new_buffer_state, new_buffer_birth_state, action, channel_state, + current_step, traffic_mean_arrival_rate, channel_occupancy_pct_window, + channel_collisions_pct_window, back_pct, unique_ltc_tx_window, + cs_tx_same_type_now, tx_collision_other_now, + enable_cs_tx_same_type, enable_tx_collision_other, + ret_c, no_tx +): + """Build the observation entry as a 14-element feature vector.""" + buffer_count = jnp.sum((new_buffer_state != EMPTY_PACKET_ID).astype(jnp.int32)) + buffer_occupancy_pct = buffer_count.astype(jnp.float32) / new_buffer_state.shape[0] + oldest_packet_age_norm = _oldest_packet_age_norm(new_buffer_state, new_buffer_birth_state, current_step) + + cs_busy = jnp.where(action == Actions.CS.value, channel_state != 0, -1).astype(jnp.float32) + cs_tx_same_type = jnp.where( + enable_cs_tx_same_type, + jnp.where(jnp.logical_and(action == Actions.CS.value, channel_state == 1), cs_tx_same_type_now, -1.0), + -1.0, + ) + tx_collision_other = jnp.where( + enable_tx_collision_other, + jnp.where(jnp.logical_and(action == Actions.TX.value, channel_state == -1), tx_collision_other_now, -1.0), + -1.0, + ) + + return jnp.array( + [ + buffer_occupancy_pct, + buffer_count.astype(jnp.float32), + oldest_packet_age_norm, + traffic_mean_arrival_rate, + cs_busy, + cs_tx_same_type.astype(jnp.float32), + tx_collision_other.astype(jnp.float32), + channel_occupancy_pct_window, + channel_collisions_pct_window, + action.astype(jnp.float32), + back_pct, + ret_c.astype(jnp.float32), + no_tx.astype(jnp.float32), + unique_ltc_tx_window, + ], + dtype=jnp.float32, + ) + + +# ============================================================================ +# Main Per-Agent Output Processing +# ============================================================================ + + def process_output_i( buffer_state, new_buffer_state, @@ -135,40 +223,13 @@ def process_output_i( ) ) - buffer_count = jnp.sum((new_buffer_state != EMPTY_PACKET_ID).astype(jnp.int32)) - buffer_occupancy_pct = buffer_count.astype(jnp.float32) / new_buffer_state.shape[0] - oldest_packet_age_norm = _oldest_packet_age_norm(new_buffer_state, new_buffer_birth_state, current_step) - - cs_busy = jnp.where(action == Actions.CS.value, channel_state != 0, -1).astype(jnp.float32) - cs_tx_same_type = jnp.where( - enable_cs_tx_same_type, - jnp.where(jnp.logical_and(action == Actions.CS.value, channel_state == 1), cs_tx_same_type_now, -1.0), - -1.0, - ) - tx_collision_other = jnp.where( - enable_tx_collision_other, - jnp.where(jnp.logical_and(action == Actions.TX.value, channel_state == -1), tx_collision_other_now, -1.0), - -1.0, - ) - - obs_t = jnp.array( - [ - buffer_occupancy_pct, - buffer_count.astype(jnp.float32), - oldest_packet_age_norm, - traffic_mean_arrival_rate, - cs_busy, - cs_tx_same_type.astype(jnp.float32), - tx_collision_other.astype(jnp.float32), - channel_occupancy_pct_window, - channel_collisions_pct_window, - action.astype(jnp.float32), - back_pct, - ret_c.astype(jnp.float32), - no_tx.astype(jnp.float32), - unique_ltc_tx_window, - ], - dtype=jnp.float32, + obs_t = _build_observation_entry( + buffer_state, new_buffer_state, new_buffer_birth_state, action, channel_state, + current_step, traffic_mean_arrival_rate, channel_occupancy_pct_window, + channel_collisions_pct_window, back_pct, unique_ltc_tx_window, + cs_tx_same_type_now, tx_collision_other_now, + enable_cs_tx_same_type, enable_tx_collision_other, + ret_c, no_tx ) obs = jax.lax.cond( roll_obs, diff --git a/ltc/utils/scan_states.py b/ltc/utils/scan_states.py index 9b75856..f0cac94 100644 --- a/ltc/utils/scan_states.py +++ b/ltc/utils/scan_states.py @@ -4,6 +4,7 @@ from reinforced_lib.agents import AgentState from ltc.sim import ModelState +from ltc.utils.state_structs import MacroActionState, ObsTrackerState @jax.tree_util.register_dataclass @@ -14,19 +15,7 @@ class Carry: traffic_states: ModelState packet_seqs: jax.Array buffer_states: jax.Array - buffer_birth_steps: jax.Array - arrival_hist: jax.Array - planned_tx_hist: jax.Array - success_tx_hist: jax.Array - channel_busy_hist: jax.Array - collision_hist: jax.Array tx_hist: jax.Array - staged_tx: jax.Array - macro_remaining: jax.Array - macro_action_types: jax.Array - macro_reward_accum: jax.Array - macro_tx_success_accum: jax.Array - macro_tx_collision_accum: jax.Array power_states: jax.Array channel_state: int key: jax.random.PRNGKey @@ -35,6 +24,8 @@ class Carry: rewards: jax.Array terminals: jax.Array active: jax.Array + macro: MacroActionState + obs_tracker: ObsTrackerState @jax.tree_util.register_dataclass diff --git a/ltc/utils/state_structs.py b/ltc/utils/state_structs.py new file mode 100644 index 0000000..3b4d0f6 --- /dev/null +++ b/ltc/utils/state_structs.py @@ -0,0 +1,50 @@ +"""Organized state structures for the macro-action simulator. + +This module defines nested dataclasses for logically grouping state fields +within the Carry object, improving readability and maintainability. +""" + +from dataclasses import dataclass +import jax + + +@jax.tree_util.register_dataclass +@dataclass +class MacroActionState: + """State for tracking macro-action execution per agent. + + Fields: + remaining: Steps left in current macro-action for each agent. + action_types: Primitive action type (TX/CS/IDLE) for current macro. + reward_accum: Accumulated reward across macro-action slots. + tx_success_accum: Count of successful TX subframes in current macro. + tx_collision_accum: Count of collided TX subframes in current macro. + """ + remaining: "jax.Array" + action_types: "jax.Array" + reward_accum: "jax.Array" + tx_success_accum: "jax.Array" + tx_collision_accum: "jax.Array" + + +@jax.tree_util.register_dataclass +@dataclass +class ObsTrackerState: + """State for tracking observation metrics and packet histories. + + Fields: + buffer_birth_steps: Arrival step for each packet in queue. + arrival_hist: Rolling window of per-slot arrival counts. + planned_tx_hist: Rolling window of planned TX attempts per agent. + success_tx_hist: Rolling window of successful TX per agent. + channel_busy_hist: Rolling window of channel busy readings. + collision_hist: Rolling window of collision outcomes. + staged_tx: Per-agent count of staged TX actions (for TX_STAGE/COMMIT). + """ + buffer_birth_steps: "jax.Array" + arrival_hist: "jax.Array" + planned_tx_hist: "jax.Array" + success_tx_hist: "jax.Array" + channel_busy_hist: "jax.Array" + collision_hist: "jax.Array" + staged_tx: "jax.Array" From 5ec7c29f5c37adf4b0ce880daf818023e489ac3f Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Tue, 21 Apr 2026 22:16:28 +0200 Subject: [PATCH 07/24] Refactor code --- ltc/run.py | 68 ++++------------ ltc/sim/constants.py | 1 - ltc/sim/process_output.py | 134 +++++++------------------------- ltc/sim/sim.py | 21 ----- ltc/utils/plots.py | 2 +- ltc/utils/scan_states.py | 45 ----------- ltc/utils/state_structs.py | 50 ------------ ltc/utils/structs.py | 94 ++++++++++++++++++++++ test/sim/test_process_output.py | 26 ++++--- 9 files changed, 159 insertions(+), 282 deletions(-) delete mode 100644 ltc/utils/scan_states.py delete mode 100644 ltc/utils/state_structs.py create mode 100644 ltc/utils/structs.py diff --git a/ltc/run.py b/ltc/run.py index 8f7345f..d93b6ce 100644 --- a/ltc/run.py +++ b/ltc/run.py @@ -14,32 +14,12 @@ from ltc.agents import BayesianDDQN, DCF, QNetwork, StochasticVariationalNetwork from ltc.sim import InitialStateConf, cox_traffic, process_output, simulate -from ltc.sim.constants import ( - INITIAL_CAPACITY, - Actions, - EMPTY_PACKET_ID, - OBS_SIZE, - MAX_MACRO_DURATION, - LEGACY_TX_DURATION, - TX_REWARD_LENGTH, - TX_SIZE_THRESHOLD, - TX_SIZE_PENALTY_WINDOW, - TX_SIZE_PENALTY, - TX_LTC_COLLISION_PENALTY, - TX_COEX_COLLISION_PENALTY, - TX_EMPTY_BUFFER_PENALTY, - TX_MAX_RETRANSMISSION_PENALTY, - MAX_RETRANSMISSION, - ObsIdx, - TX_REWARD, -) +from ltc.sim.constants import * from ltc.utils.history import build_history_filename, ensure_clean_git_worktree, get_short_commit_hash -from ltc.utils.scan_states import Carry, Output -from ltc.utils.state_structs import MacroActionState, ObsTrackerState +from ltc.utils.structs import MacroActionState, ObsTrackerState, ActionDecodingResults, Carry, Output from ltc.utils.plots import plot_all, plot_first - def init_agents(agent, key, n): keys = jax.random.split(key, n) states = jax.vmap(agent.init)(keys) @@ -179,53 +159,35 @@ def decode_legacy_actions(raw_actions, legacy_tx_duration): return action_type, duration -# ============================================================================ -# Macro-Action State Management -# ============================================================================ - def _update_macro_state_from_ready_agents( macro_state, obs_tracker, ready_drl, ready_legacy, - drl_action_types, drl_durations, drl_staged_next, - legacy_action_types, legacy_durations, - n_drl + decoded_actions, n_drl ): - """Update macro-action state for agents that are ready to choose new actions. - - Ready DRL agents: update with decoded DRL actions - Ready legacy agents: update with decoded legacy actions - Non-ready agents: keep existing macro state - """ macro_action_types = macro_state.action_types macro_remaining = macro_state.remaining staged_tx = obs_tracker.staged_tx - # DRL agents macro_action_types = macro_action_types.at[:n_drl].set( - jnp.where(ready_drl, drl_action_types, macro_state.action_types[:n_drl]) + jnp.where(ready_drl, decoded_actions.drl_action_types, macro_state.action_types[:n_drl]) ) macro_remaining = macro_remaining.at[:n_drl].set( - jnp.where(ready_drl, drl_durations, macro_state.remaining[:n_drl]) + jnp.where(ready_drl, decoded_actions.drl_durations, macro_state.remaining[:n_drl]) ) staged_tx = staged_tx.at[:n_drl].set( - jnp.where(ready_drl, drl_staged_next, obs_tracker.staged_tx[:n_drl]) + jnp.where(ready_drl, decoded_actions.drl_staged_next, obs_tracker.staged_tx[:n_drl]) ) - # Legacy agents macro_action_types = macro_action_types.at[n_drl:].set( - jnp.where(ready_legacy, legacy_action_types, macro_state.action_types[n_drl:]) + jnp.where(ready_legacy, decoded_actions.legacy_action_types, macro_state.action_types[n_drl:]) ) macro_remaining = macro_remaining.at[n_drl:].set( - jnp.where(ready_legacy, legacy_durations, macro_state.remaining[n_drl:]) + jnp.where(ready_legacy, decoded_actions.legacy_durations, macro_state.remaining[n_drl:]) ) return macro_action_types, macro_remaining, staged_tx def _finalize_macro_accumulators(macro_state, done_now, macro_tx_success_accum, macro_tx_collision_accum, macro_reward_accum): - """Reset macro accumulators for agents completing their macro-action. - - Agents that are done_now get their accumulators zeroed for the next macro-action. - """ macro_tx_success_accum = jnp.where(done_now, 0.0, macro_tx_success_accum) macro_tx_collision_accum = jnp.where(done_now, 0.0, macro_tx_collision_accum) macro_reward_accum = jnp.where(done_now, 0.0, macro_reward_accum) @@ -303,11 +265,17 @@ def rl_step_coroutine(c, step): ) legacy_action_types, legacy_durations = decode_legacy_actions(raw_actions[n_drl:], legacy_tx_duration) + decoded_actions = ActionDecodingResults( + drl_action_types=drl_action_types, + drl_durations=drl_durations, + drl_staged_next=drl_staged_next, + legacy_action_types=legacy_action_types, + legacy_durations=legacy_durations, + ) + macro_action_types, macro_remaining, staged_tx = _update_macro_state_from_ready_agents( c.macro, c.obs_tracker, drl_ready, legacy_ready, - drl_action_types, drl_durations, drl_staged_next, - legacy_action_types, legacy_durations, - n_drl + decoded_actions, n_drl ) slot_executing = macro_remaining > 0 @@ -387,7 +355,6 @@ def rl_step_coroutine(c, step): slot_actions, channel_state, successful_packets, prev_ret_c, tx_collision_other_now ) - # TODO: Implement "Else: max 6 ms" clause once precise behavior is finalized. slot_reward_upgraded = jnp.where(tx_executing, tx_slot_reward, slot_rewards) macro_tx_success_accum = c.macro.tx_success_accum + jnp.where(slot_executing, k_tx_slot, 0.0) @@ -454,7 +421,6 @@ def pre_rl_fn(*args, **kwargs): def post_rl_fn(intermediate, *args, **kwargs): gen = rl_step_coroutine(*args, **kwargs) - # Unused computations will be DCE-ed _ = next(gen) return gen.send(intermediate) diff --git a/ltc/sim/constants.py b/ltc/sim/constants.py index 90c9f4d..bdabf3b 100644 --- a/ltc/sim/constants.py +++ b/ltc/sim/constants.py @@ -29,7 +29,6 @@ class Actions(Enum): COLLISION_PENALTY = -1.0 MAX_RETRANSMISSION_PENALTY = -1.0 -# Upgraded TX reward parameters TX_REWARD_LENGTH = 5 TX_SIZE_THRESHOLD = 5 TX_SIZE_PENALTY_WINDOW = 25 diff --git a/ltc/sim/process_output.py b/ltc/sim/process_output.py index 8bb369f..5fde586 100644 --- a/ltc/sim/process_output.py +++ b/ltc/sim/process_output.py @@ -3,41 +3,9 @@ from ltc.sim.constants import * from ltc.sim.sim import is_buffer_empty +from ltc.utils.structs import ObsFeatureInputs, ObsFeatureConfig -# ============================================================================ -# Reward Computation Helpers -# ============================================================================ - -def _compute_cs_reward(buffer_state, channel_state, no_tx, ret_c, key): - """Compute reward for CS (channel sensing) action. - - CS reward branches on whether buffer is empty or full: - - Empty buffer: EMPTY_BUFFER_REWARD - - Full buffer with short idle: NO_TX_REWARD - - Full buffer with long idle: NO_TX_PENALTY (scaled) - """ - return jax.lax.cond(is_buffer_empty(buffer_state), idle_empty_buffer, idle_full_buffer, - (buffer_state, ret_c, no_tx, key)) - - -def _compute_tx_reward(buffer_state, channel_state, ret_c, no_tx, key): - """Compute reward for TX (transmission) action. - - TX reward branches on channel outcome: - - Collision with collision_count < MAX_RETRANSMISSION: COLLISION_PENALTY, ret_c+1 - - Collision with collision_count >= MAX_RETRANSMISSION: MAX_RETRANSMISSION_PENALTY, reset - - Success with empty buffer: EMPTY_TX_PENALTY - - Success with full buffer: TX_REWARD, ret_c=0 - """ - args = (buffer_state, ret_c, channel_state, no_tx, key) - return jax.lax.cond(channel_state == 1, transmission_without_collision, transmission_with_collision, args) - - -# ============================================================================ -# Legacy Reward Branch Functions (organized for clarity) -# ============================================================================ - def no_transmission(args): _, buffer_state, _, _, _, _ = args return jax.lax.cond(is_buffer_empty(buffer_state), idle_empty_buffer, idle_full_buffer, args) @@ -127,59 +95,42 @@ def _oldest_packet_age_norm(buffer_state, buffer_birth_state, current_step): return jnp.where(age >= 0, age.astype(jnp.float32) / queue_size, -1.0) -# ============================================================================ -# Observation Construction -# ============================================================================ - def _build_observation_entry( buffer_state, new_buffer_state, new_buffer_birth_state, action, channel_state, - current_step, traffic_mean_arrival_rate, channel_occupancy_pct_window, - channel_collisions_pct_window, back_pct, unique_ltc_tx_window, - cs_tx_same_type_now, tx_collision_other_now, - enable_cs_tx_same_type, enable_tx_collision_other, - ret_c, no_tx + current_step, obs_features, obs_config, ret_c, no_tx ): - """Build the observation entry as a 14-element feature vector.""" buffer_count = jnp.sum((new_buffer_state != EMPTY_PACKET_ID).astype(jnp.int32)) buffer_occupancy_pct = buffer_count.astype(jnp.float32) / new_buffer_state.shape[0] oldest_packet_age_norm = _oldest_packet_age_norm(new_buffer_state, new_buffer_birth_state, current_step) cs_busy = jnp.where(action == Actions.CS.value, channel_state != 0, -1).astype(jnp.float32) cs_tx_same_type = jnp.where( - enable_cs_tx_same_type, - jnp.where(jnp.logical_and(action == Actions.CS.value, channel_state == 1), cs_tx_same_type_now, -1.0), + obs_config.enable_cs_tx_same_type, + jnp.where(jnp.logical_and(action == Actions.CS.value, channel_state == 1), obs_features.cs_tx_same_type_now, -1.0), -1.0, ) tx_collision_other = jnp.where( - enable_tx_collision_other, - jnp.where(jnp.logical_and(action == Actions.TX.value, channel_state == -1), tx_collision_other_now, -1.0), + obs_config.enable_tx_collision_other, + jnp.where(jnp.logical_and(action == Actions.TX.value, channel_state == -1), obs_features.tx_collision_other_now, -1.0), -1.0, ) - return jnp.array( - [ - buffer_occupancy_pct, - buffer_count.astype(jnp.float32), - oldest_packet_age_norm, - traffic_mean_arrival_rate, - cs_busy, - cs_tx_same_type.astype(jnp.float32), - tx_collision_other.astype(jnp.float32), - channel_occupancy_pct_window, - channel_collisions_pct_window, - action.astype(jnp.float32), - back_pct, - ret_c.astype(jnp.float32), - no_tx.astype(jnp.float32), - unique_ltc_tx_window, - ], - dtype=jnp.float32, - ) - - -# ============================================================================ -# Main Per-Agent Output Processing -# ============================================================================ + return jnp.array([ + buffer_occupancy_pct, + buffer_count.astype(jnp.float32), + oldest_packet_age_norm, + obs_features.traffic_mean_arrival_rate, + cs_busy, + cs_tx_same_type.astype(jnp.float32), + tx_collision_other.astype(jnp.float32), + obs_features.channel_occupancy_pct_window, + obs_features.channel_collisions_pct_window, + action.astype(jnp.float32), + obs_features.back_pct, + ret_c.astype(jnp.float32), + no_tx.astype(jnp.float32), + obs_features.unique_ltc_tx_window, + ], dtype=jnp.float32) def process_output_i( @@ -193,15 +144,8 @@ def process_output_i( terminal, key, current_step, - traffic_mean_arrival_rate, - channel_occupancy_pct_window, - channel_collisions_pct_window, - back_pct, - unique_ltc_tx_window, - cs_tx_same_type_now, - tx_collision_other_now, - enable_cs_tx_same_type, - enable_tx_collision_other, + obs_features, + obs_config, roll_obs, ): ret_c = obs[-1, ObsIdx.STATUS_RETRY_COUNTER].astype(jnp.int32) @@ -225,11 +169,7 @@ def process_output_i( obs_t = _build_observation_entry( buffer_state, new_buffer_state, new_buffer_birth_state, action, channel_state, - current_step, traffic_mean_arrival_rate, channel_occupancy_pct_window, - channel_collisions_pct_window, back_pct, unique_ltc_tx_window, - cs_tx_same_type_now, tx_collision_other_now, - enable_cs_tx_same_type, enable_tx_collision_other, - ret_c, no_tx + current_step, obs_features, obs_config, ret_c, no_tx ) obs = jax.lax.cond( roll_obs, @@ -252,19 +192,12 @@ def process_output( terminals, key, current_step, - traffic_mean_arrival_rate, - channel_occupancy_pct_window, - channel_collisions_pct_window, - back_pct, - unique_ltc_tx_window, - cs_tx_same_type_now, - tx_collision_other_now, - enable_cs_tx_same_type, - enable_tx_collision_other, + obs_features, + obs_config, roll_obs=True, ): roll_obs = jnp.broadcast_to(jnp.asarray(roll_obs), actions.shape) - return jax.vmap(process_output_i, in_axes=(0, 0, 0, 0, None, 0, 0, 0, None, None, 0, None, None, 0, None, 0, 0, None, None, 0))( + return jax.vmap(process_output_i, in_axes=(0, 0, 0, 0, None, 0, 0, 0, None, None, 0, None, 0))( buffer_states, new_buffer_states, new_buffer_birth_states, @@ -275,14 +208,7 @@ def process_output( terminals, key, current_step, - traffic_mean_arrival_rate, - channel_occupancy_pct_window, - channel_collisions_pct_window, - back_pct, - unique_ltc_tx_window, - cs_tx_same_type_now, - tx_collision_other_now, - enable_cs_tx_same_type, - enable_tx_collision_other, + obs_features, + obs_config, roll_obs, ) diff --git a/ltc/sim/sim.py b/ltc/sim/sim.py index abdacaf..f51cf1f 100644 --- a/ltc/sim/sim.py +++ b/ltc/sim/sim.py @@ -5,13 +5,6 @@ def channel_state_selector(actions): - """ - Returns a state of the channel due to how many STA transmit in the same time. - :param actions: vector of stations that transmit at a given time - :return: - 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) return jnp.where( ones_count > 1, -1, @@ -29,15 +22,11 @@ def is_buffer_empty(buffer_state): def remove_packets_from_queue(queue_state, remove_mask): - """ - Removes packets selected by remove_mask and compacts remaining packets to queue head. - """ n = queue_state.shape[0] valid = queue_state != EMPTY_PACKET_ID remove_mask = jnp.asarray(remove_mask, dtype=bool) keep = valid & ~remove_mask - # Stable compaction using argsort over static-size keys. idx = jnp.arange(n) sort_key = jnp.where(keep, idx, n + idx) order = jnp.argsort(sort_key) @@ -52,12 +41,6 @@ def remove_transmitted_packets(buffer_states, tx_masks): def enqueue_generated_packets(buffer_state, buffer_birth_state, new_frames, station_id, packet_seq, current_step): - """ - Queue semantics: - - transmitted packets are removed elsewhere, - - arrivals are appended at queue tail, - - if overflow occurs, oldest packets are dropped. - """ q = buffer_state.shape[0] new_frames = jnp.maximum(new_frames.astype(jnp.int32), 0) occupied = jnp.sum((buffer_state != EMPTY_PACKET_ID).astype(jnp.int32)) @@ -94,10 +77,6 @@ def add_new_frames(buffer_states, buffer_birth_steps, new_frames, packet_seqs, c def simulate(buffer_states, buffer_birth_steps, new_frames, actions, packet_seqs, current_step): - """ - One simulation step with fixed-size packet queues. - On successful TX, one packet from queue head is removed for transmitting station. - """ channel_state = channel_state_selector(actions) pre_nonempty = jnp.any(buffer_states != EMPTY_PACKET_ID, axis=1) success_tx = (actions == Actions.TX.value) & (channel_state == 1) diff --git a/ltc/utils/plots.py b/ltc/utils/plots.py index 9731885..188ad92 100644 --- a/ltc/utils/plots.py +++ b/ltc/utils/plots.py @@ -8,7 +8,7 @@ import matplotlib.pyplot as plt import numpy as np -from ltc.utils.scan_states import Output +from ltc.utils.structs import Output from ltc.utils.history import parse_history_filename, unpack_history from ltc.sim.constants import NO_TX_REWARD, Actions, INITIAL_CAPACITY, EMPTY_PACKET_ID diff --git a/ltc/utils/scan_states.py b/ltc/utils/scan_states.py deleted file mode 100644 index f0cac94..0000000 --- a/ltc/utils/scan_states.py +++ /dev/null @@ -1,45 +0,0 @@ -from dataclasses import dataclass - -import jax -from reinforced_lib.agents import AgentState - -from ltc.sim import ModelState -from ltc.utils.state_structs import MacroActionState, ObsTrackerState - - -@jax.tree_util.register_dataclass -@dataclass -class Carry: - drl_states: AgentState - legacy_states: AgentState - traffic_states: ModelState - packet_seqs: jax.Array - buffer_states: jax.Array - tx_hist: jax.Array - power_states: jax.Array - channel_state: int - key: jax.random.PRNGKey - obs: jax.Array - actions: jax.Array - rewards: jax.Array - terminals: jax.Array - active: jax.Array - macro: MacroActionState - obs_tracker: ObsTrackerState - - -@jax.tree_util.register_dataclass -@dataclass -class Output: - legacy_states: AgentState - observations: jax.Array - actions: jax.Array - rewards: jax.Array - terminals: jax.Array - buffer_states: jax.Array - power_states: jax.Array - new_frames: jax.Array - channel_state: int - active: jax.Array - weights_histogram: jax.Array - weights_bin_edges: jax.Array diff --git a/ltc/utils/state_structs.py b/ltc/utils/state_structs.py deleted file mode 100644 index 3b4d0f6..0000000 --- a/ltc/utils/state_structs.py +++ /dev/null @@ -1,50 +0,0 @@ -"""Organized state structures for the macro-action simulator. - -This module defines nested dataclasses for logically grouping state fields -within the Carry object, improving readability and maintainability. -""" - -from dataclasses import dataclass -import jax - - -@jax.tree_util.register_dataclass -@dataclass -class MacroActionState: - """State for tracking macro-action execution per agent. - - Fields: - remaining: Steps left in current macro-action for each agent. - action_types: Primitive action type (TX/CS/IDLE) for current macro. - reward_accum: Accumulated reward across macro-action slots. - tx_success_accum: Count of successful TX subframes in current macro. - tx_collision_accum: Count of collided TX subframes in current macro. - """ - remaining: "jax.Array" - action_types: "jax.Array" - reward_accum: "jax.Array" - tx_success_accum: "jax.Array" - tx_collision_accum: "jax.Array" - - -@jax.tree_util.register_dataclass -@dataclass -class ObsTrackerState: - """State for tracking observation metrics and packet histories. - - Fields: - buffer_birth_steps: Arrival step for each packet in queue. - arrival_hist: Rolling window of per-slot arrival counts. - planned_tx_hist: Rolling window of planned TX attempts per agent. - success_tx_hist: Rolling window of successful TX per agent. - channel_busy_hist: Rolling window of channel busy readings. - collision_hist: Rolling window of collision outcomes. - staged_tx: Per-agent count of staged TX actions (for TX_STAGE/COMMIT). - """ - buffer_birth_steps: "jax.Array" - arrival_hist: "jax.Array" - planned_tx_hist: "jax.Array" - success_tx_hist: "jax.Array" - channel_busy_hist: "jax.Array" - collision_hist: "jax.Array" - staged_tx: "jax.Array" diff --git a/ltc/utils/structs.py b/ltc/utils/structs.py new file mode 100644 index 0000000..bae4681 --- /dev/null +++ b/ltc/utils/structs.py @@ -0,0 +1,94 @@ +from dataclasses import dataclass +import jax +from reinforced_lib.agents import AgentState + +from ltc.sim import ModelState + + +@jax.tree_util.register_dataclass +@dataclass +class MacroActionState: + remaining: jax.Array + action_types: jax.Array + reward_accum: jax.Array + tx_success_accum: jax.Array + tx_collision_accum: jax.Array + + +@jax.tree_util.register_dataclass +@dataclass +class ObsTrackerState: + buffer_birth_steps: jax.Array + arrival_hist: jax.Array + planned_tx_hist: jax.Array + success_tx_hist: jax.Array + channel_busy_hist: jax.Array + collision_hist: jax.Array + staged_tx: jax.Array + + +@jax.tree_util.register_dataclass +@dataclass +class ObsFeatureInputs: + traffic_mean_arrival_rate: jax.Array + channel_occupancy_pct_window: jax.Array + channel_collisions_pct_window: jax.Array + back_pct: jax.Array + unique_ltc_tx_window: jax.Array + cs_tx_same_type_now: jax.Array + tx_collision_other_now: jax.Array + + +@jax.tree_util.register_dataclass +@dataclass +class ObsFeatureConfig: + enable_cs_tx_same_type: bool + enable_tx_collision_other: bool + + +@jax.tree_util.register_dataclass +@dataclass +class ActionDecodingResults: + drl_action_types: jax.Array + drl_durations: jax.Array + drl_staged_next: jax.Array + legacy_action_types: jax.Array + legacy_durations: jax.Array + + +@jax.tree_util.register_dataclass +@dataclass +class Carry: + drl_states: AgentState + legacy_states: AgentState + traffic_states: ModelState + packet_seqs: jax.Array + buffer_states: jax.Array + tx_hist: jax.Array + power_states: jax.Array + channel_state: int + key: jax.random.PRNGKey + obs: jax.Array + actions: jax.Array + rewards: jax.Array + terminals: jax.Array + active: jax.Array + macro: MacroActionState + obs_tracker: ObsTrackerState + + +@jax.tree_util.register_dataclass +@dataclass +class Output: + legacy_states: AgentState + observations: jax.Array + actions: jax.Array + rewards: jax.Array + terminals: jax.Array + buffer_states: jax.Array + power_states: jax.Array + new_frames: jax.Array + channel_state: int + active: jax.Array + weights_histogram: jax.Array + weights_bin_edges: jax.Array diff --git a/test/sim/test_process_output.py b/test/sim/test_process_output.py index 15718ae..6cebf01 100644 --- a/test/sim/test_process_output.py +++ b/test/sim/test_process_output.py @@ -5,6 +5,7 @@ from ltc.sim.constants import * from ltc.sim.process_output import * +from ltc.utils.structs import ObsFeatureInputs, ObsFeatureConfig KEY = jax.random.PRNGKey(0) EMPTY = EMPTY_PACKET_ID @@ -87,6 +88,20 @@ def test_process_rl_output(self): expected_R_0 = -1.0 expected_power = power_states + jnp.array([-TX_CONSUMPTION, -CS_CONSUMPTION, -TX_CONSUMPTION, -TX_CONSUMPTION]) + obs_features = ObsFeatureInputs( + traffic_mean_arrival_rate=jnp.array([1.0, 2.0, 3.0, 4.0], dtype=jnp.float32), + channel_occupancy_pct_window=jnp.full((4,), 0.25, dtype=jnp.float32), + channel_collisions_pct_window=jnp.full((4,), 0.5, dtype=jnp.float32), + back_pct=jnp.array([0.0, 0.6, 1.0, 0.0], dtype=jnp.float32), + unique_ltc_tx_window=jnp.full((4,), 2.0, dtype=jnp.float32), + cs_tx_same_type_now=jnp.array([0.0, 1.0, 0.0, 0.0], dtype=jnp.float32), + tx_collision_other_now=jnp.array([1.0, 0.0, 1.0, 1.0], dtype=jnp.float32), + ) + obs_config = ObsFeatureConfig( + enable_cs_tx_same_type=True, + enable_tx_collision_other=True, + ) + result_obs, result_R, power = process_output( buffer_states=buffer_states, new_buffer_states=new_buffer_states, @@ -98,15 +113,8 @@ def test_process_rl_output(self): terminals=terminals, key=KEY, current_step=10, - traffic_mean_arrival_rate=jnp.array([1.0, 2.0, 3.0, 4.0], dtype=jnp.float32), - channel_occupancy_pct_window=0.25, - channel_collisions_pct_window=0.5, - back_pct=jnp.array([0.0, 0.6, 1.0, 0.0], dtype=jnp.float32), - unique_ltc_tx_window=2.0, - cs_tx_same_type_now=jnp.array([0.0, 1.0, 0.0, 0.0], dtype=jnp.float32), - tx_collision_other_now=jnp.array([1.0, 0.0, 1.0, 1.0], dtype=jnp.float32), - enable_cs_tx_same_type=True, - enable_tx_collision_other=True, + obs_features=obs_features, + obs_config=obs_config, ) self.assertEqual(result_obs.shape[-1], OBS_SIZE) From 98425920416192a07c8cd01e89828a76a843ece3 Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Tue, 21 Apr 2026 22:19:25 +0200 Subject: [PATCH 08/24] Remove file --- ltc/sim/observation.py | 1 - 1 file changed, 1 deletion(-) delete mode 100644 ltc/sim/observation.py diff --git a/ltc/sim/observation.py b/ltc/sim/observation.py deleted file mode 100644 index 00bc499..0000000 --- a/ltc/sim/observation.py +++ /dev/null @@ -1 +0,0 @@ -from ltc.sim.constants import ObsIdx, OBS_SIZE From afcb1be369658ffdb427190decf0f6c86a940caf Mon Sep 17 00:00:00 2001 From: kaszczech Date: Wed, 22 Apr 2026 08:12:48 +0200 Subject: [PATCH 09/24] Multi-slot tx --- ltc/run.py | 21 ++++++++++++++++----- 1 file changed, 16 insertions(+), 5 deletions(-) diff --git a/ltc/run.py b/ltc/run.py index d93b6ce..02779d0 100644 --- a/ltc/run.py +++ b/ltc/run.py @@ -16,7 +16,7 @@ from ltc.sim import InitialStateConf, cox_traffic, process_output, simulate from ltc.sim.constants import * from ltc.utils.history import build_history_filename, ensure_clean_git_worktree, get_short_commit_hash -from ltc.utils.structs import MacroActionState, ObsTrackerState, ActionDecodingResults, Carry, Output +from ltc.utils.structs import MacroActionState, ObsTrackerState, ActionDecodingResults, Carry, Output, ObsFeatureInputs, ObsFeatureConfig from ltc.utils.plots import plot_all, plot_first @@ -344,11 +344,22 @@ def rl_step_coroutine(c, step): done_now = jnp.logical_and(slot_executing, macro_remaining == 1) prev_ret_c = c.obs[:, -1, ObsIdx.STATUS_RETRY_COUNTER].astype(jnp.int32) + obs_features = ObsFeatureInputs( + traffic_mean_arrival_rate=traffic_mean_arrival_rate, + channel_occupancy_pct_window=jnp.broadcast_to(channel_occupancy_pct_window, (n,)), + channel_collisions_pct_window=jnp.broadcast_to(channel_collisions_pct_window, (n,)), + back_pct=back_pct, + unique_ltc_tx_window=jnp.broadcast_to(unique_ltc_tx_window, (n,)), + cs_tx_same_type_now=cs_tx_same_type_now, + tx_collision_other_now=tx_collision_other_now, + ) + obs_config = ObsFeatureConfig( + enable_cs_tx_same_type=obs_enable_cs_tx_same_type, + enable_tx_collision_other=obs_enable_tx_collision_other, + ) obs, slot_rewards, powers = process_output( c.buffer_states, buffer_states, buffer_birth_steps, c.power_states, channel_state, c.obs, slot_actions, - c.terminals, reward_key, step, traffic_mean_arrival_rate, channel_occupancy_pct_window, - channel_collisions_pct_window, back_pct, unique_ltc_tx_window, cs_tx_same_type_now, - tx_collision_other_now, obs_enable_cs_tx_same_type, obs_enable_tx_collision_other, roll_obs=done_now + c.terminals, reward_key, step, obs_features, obs_config, roll_obs=done_now ) tx_slot_reward, k_tx_slot, k_coll_slot, tx_executing = upgraded_tx_slot_reward( @@ -390,7 +401,7 @@ def rl_step_coroutine(c, step): tx_collision_accum=macro_tx_collision_accum, ) obs_tracker_state = ObsTrackerState( - buffer_birth_steps=c.obs_tracker.buffer_birth_steps, + buffer_birth_steps=buffer_birth_steps, arrival_hist=arrival_hist, planned_tx_hist=planned_tx_hist, success_tx_hist=success_tx_hist, From 7128de6832ba43dcc3898354a518e01b6c595230 Mon Sep 17 00:00:00 2001 From: kaszczech Date: Wed, 22 Apr 2026 10:16:09 +0200 Subject: [PATCH 10/24] Add TX_SLOTS time scaling: TX takes 3x longer than CS mini-slots --- ltc/run.py | 19 ++++++++++++++++--- ltc/sim/constants.py | 6 ++++++ ltc/sim/sim.py | 10 ++++++---- test/sim/test_sim.py | 4 +++- 4 files changed, 31 insertions(+), 8 deletions(-) diff --git a/ltc/run.py b/ltc/run.py index 02779d0..3f4afd9 100644 --- a/ltc/run.py +++ b/ltc/run.py @@ -170,8 +170,13 @@ def _update_macro_state_from_ready_agents( macro_action_types = macro_action_types.at[:n_drl].set( jnp.where(ready_drl, decoded_actions.drl_action_types, macro_state.action_types[:n_drl]) ) + drl_scaled_durations = jnp.where( + decoded_actions.drl_action_types == Actions.TX.value, + decoded_actions.drl_durations * TX_SLOTS, + decoded_actions.drl_durations, + ) macro_remaining = macro_remaining.at[:n_drl].set( - jnp.where(ready_drl, decoded_actions.drl_durations, macro_state.remaining[:n_drl]) + jnp.where(ready_drl, drl_scaled_durations, macro_state.remaining[:n_drl]) ) staged_tx = staged_tx.at[:n_drl].set( jnp.where(ready_drl, decoded_actions.drl_staged_next, obs_tracker.staged_tx[:n_drl]) @@ -180,8 +185,13 @@ def _update_macro_state_from_ready_agents( macro_action_types = macro_action_types.at[n_drl:].set( jnp.where(ready_legacy, decoded_actions.legacy_action_types, macro_state.action_types[n_drl:]) ) + legacy_scaled_durations = jnp.where( + decoded_actions.legacy_action_types == Actions.TX.value, + decoded_actions.legacy_durations * TX_SLOTS, + decoded_actions.legacy_durations, + ) macro_remaining = macro_remaining.at[n_drl:].set( - jnp.where(ready_legacy, decoded_actions.legacy_durations, macro_state.remaining[n_drl:]) + jnp.where(ready_legacy, legacy_scaled_durations, macro_state.remaining[n_drl:]) ) return macro_action_types, macro_remaining, staged_tx @@ -281,6 +291,9 @@ def rl_step_coroutine(c, step): slot_executing = macro_remaining > 0 slot_actions = jnp.where(slot_executing, macro_action_types, Actions.IDLE.value) + # TX transmits a packet only at the first mini-slot of each TX unit + tx_packet_mask = (macro_remaining % TX_SLOTS == 0) & (slot_actions == Actions.TX.value) + traffic_states, new_frames = traffic_step(c.traffic_states, traffic_keys) ( buffer_states, @@ -289,7 +302,7 @@ def rl_step_coroutine(c, step): packet_seqs, planned_packets, successful_packets, - ) = simulate(c.buffer_states, c.obs_tracker.buffer_birth_steps, new_frames, slot_actions, c.packet_seqs, step) + ) = simulate(c.buffer_states, c.obs_tracker.buffer_birth_steps, new_frames, slot_actions, c.packet_seqs, step, tx_packet_mask) is_ltc = jnp.arange(n) < n_drl tx_mask = slot_actions == Actions.TX.value diff --git a/ltc/sim/constants.py b/ltc/sim/constants.py index bdabf3b..7ef3c15 100644 --- a/ltc/sim/constants.py +++ b/ltc/sim/constants.py @@ -61,6 +61,12 @@ class Actions(Enum): MAX_MACRO_DURATION = 5 LEGACY_TX_DURATION = 5 +""" +Time scaling: one TX slot occupies TX_SLOTS mini-slots. +During one TX, CS agents can observe TX_SLOTS times. +""" +TX_SLOTS = 3 + class ObsIdx: BUFFER_OCCUPANCY_PCT = 0 diff --git a/ltc/sim/sim.py b/ltc/sim/sim.py index f51cf1f..7d6f72a 100644 --- a/ltc/sim/sim.py +++ b/ltc/sim/sim.py @@ -76,15 +76,17 @@ def add_new_frames(buffer_states, buffer_birth_steps, new_frames, packet_seqs, c ) -def simulate(buffer_states, buffer_birth_steps, new_frames, actions, packet_seqs, current_step): +def simulate(buffer_states, buffer_birth_steps, new_frames, actions, packet_seqs, current_step, tx_packet_mask): channel_state = channel_state_selector(actions) pre_nonempty = jnp.any(buffer_states != EMPTY_PACKET_ID, axis=1) success_tx = (actions == Actions.TX.value) & (channel_state == 1) - planned_packets = (actions == Actions.TX.value).astype(jnp.int32) - successful_packets = (success_tx & pre_nonempty).astype(jnp.int32) + # tx_packet_mask: True only in the first mini-slot of each TX unit (actual packet send) + packet_tx = success_tx & tx_packet_mask + planned_packets = ((actions == Actions.TX.value) & tx_packet_mask).astype(jnp.int32) + successful_packets = (packet_tx & pre_nonempty).astype(jnp.int32) queue_size = buffer_states.shape[1] head_mask = jnp.arange(queue_size) == 0 - tx_masks = success_tx[:, None] & head_mask[None, :] + tx_masks = packet_tx[:, None] & head_mask[None, :] buffer_states = remove_transmitted_packets(buffer_states, tx_masks) buffer_birth_steps = remove_transmitted_packets(buffer_birth_steps, tx_masks) diff --git a/test/sim/test_sim.py b/test/sim/test_sim.py index ff816eb..5d4c73a 100644 --- a/test/sim/test_sim.py +++ b/test/sim/test_sim.py @@ -78,6 +78,7 @@ def test_simulate(self): new_frames = jnp.array([0, 2, 1, 0], dtype=jnp.int32) actions = jnp.array([Actions.TX.value, Actions.CS.value, Actions.CS.value, Actions.CS.value]) + tx_packet_mask = jnp.array([True, False, False, False]) ( result_buffer_states, result_buffer_birth_steps, @@ -86,7 +87,8 @@ def test_simulate(self): planned_packets, successful_packets, ) = simulate( - buffer_states, buffer_birth_steps, new_frames, actions, packet_seqs, current_step=10 + buffer_states, buffer_birth_steps, new_frames, actions, packet_seqs, current_step=10, + tx_packet_mask=tx_packet_mask, ) self.assertEqual(result_channel_state, 1) From 2225e81a5f0c73f9145099a3b9716baf79c4af3c Mon Sep 17 00:00:00 2001 From: kaszczech Date: Wed, 22 Apr 2026 10:17:00 +0200 Subject: [PATCH 11/24] Add TX_SLOTS time scaling for multi-slot TX --- ltc/utils/structs.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/ltc/utils/structs.py b/ltc/utils/structs.py index bae4681..b35d09d 100644 --- a/ltc/utils/structs.py +++ b/ltc/utils/structs.py @@ -1,9 +1,8 @@ from dataclasses import dataclass +from typing import Any import jax from reinforced_lib.agents import AgentState -from ltc.sim import ModelState - @jax.tree_util.register_dataclass @dataclass @@ -61,7 +60,7 @@ class ActionDecodingResults: class Carry: drl_states: AgentState legacy_states: AgentState - traffic_states: ModelState + traffic_states: Any packet_seqs: jax.Array buffer_states: jax.Array tx_hist: jax.Array From 7d53028f36dde7e5bed6164ba9c87498ac4610fb Mon Sep 17 00:00:00 2001 From: kaszczech Date: Wed, 22 Apr 2026 11:19:25 +0200 Subject: [PATCH 12/24] Reward fix --- ltc/run.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/ltc/run.py b/ltc/run.py index 3f4afd9..2f64a83 100644 --- a/ltc/run.py +++ b/ltc/run.py @@ -204,7 +204,7 @@ def _finalize_macro_accumulators(macro_state, done_now, macro_tx_success_accum, return macro_tx_success_accum, macro_tx_collision_accum, macro_reward_accum -def upgraded_tx_slot_reward(slot_actions, channel_state, successful_packets, prev_ret_c, tx_collision_other_now): +def upgraded_tx_slot_reward(slot_actions, channel_state, successful_packets, prev_ret_c, tx_collision_other_now, tx_packet_mask): tx_executing = slot_actions == Actions.TX.value tx_success = jnp.logical_and(tx_executing, successful_packets > 0) tx_collision = jnp.logical_and(tx_executing, channel_state == -1) @@ -213,7 +213,7 @@ def upgraded_tx_slot_reward(slot_actions, channel_state, successful_packets, pre tx_collision_coex = jnp.logical_and(tx_collision_normal, tx_collision_other_now > 0.5) tx_collision_ltc = jnp.logical_and(tx_collision_normal, tx_collision_other_now <= 0.5) tx_empty = jnp.logical_and( - tx_executing, + tx_executing & tx_packet_mask, jnp.logical_and(channel_state == 1, successful_packets == 0), ) @@ -376,7 +376,7 @@ def rl_step_coroutine(c, step): ) tx_slot_reward, k_tx_slot, k_coll_slot, tx_executing = upgraded_tx_slot_reward( - slot_actions, channel_state, successful_packets, prev_ret_c, tx_collision_other_now + slot_actions, channel_state, successful_packets, prev_ret_c, tx_collision_other_now, tx_packet_mask ) slot_reward_upgraded = jnp.where(tx_executing, tx_slot_reward, slot_rewards) From d20aad90fe85ec3072e581831ea6e4aa28bc4b7c Mon Sep 17 00:00:00 2001 From: kaszczech Date: Wed, 22 Apr 2026 13:20:30 +0200 Subject: [PATCH 13/24] Successful packet tracking --- ltc/run.py | 4 ++-- ltc/sim/constants.py | 18 +++++++++--------- ltc/utils/structs.py | 1 + 3 files changed, 12 insertions(+), 11 deletions(-) diff --git a/ltc/run.py b/ltc/run.py index 2f64a83..1acb757 100644 --- a/ltc/run.py +++ b/ltc/run.py @@ -207,7 +207,7 @@ def _finalize_macro_accumulators(macro_state, done_now, macro_tx_success_accum, def upgraded_tx_slot_reward(slot_actions, channel_state, successful_packets, prev_ret_c, tx_collision_other_now, tx_packet_mask): tx_executing = slot_actions == Actions.TX.value tx_success = jnp.logical_and(tx_executing, successful_packets > 0) - tx_collision = jnp.logical_and(tx_executing, channel_state == -1) + tx_collision = jnp.logical_and(tx_executing & tx_packet_mask, channel_state == -1) tx_collision_max = jnp.logical_and(tx_collision, prev_ret_c >= MAX_RETRANSMISSION) tx_collision_normal = jnp.logical_and(tx_collision, prev_ret_c < MAX_RETRANSMISSION) tx_collision_coex = jnp.logical_and(tx_collision_normal, tx_collision_other_now > 0.5) @@ -429,7 +429,7 @@ def rl_step_coroutine(c, step): ) o = Output( legacy_states, obs, raw_actions, rewards, terminals, buffer_states, powers, - new_frames, channel_state, active, hist, bin_edges + new_frames, channel_state, active, hist, bin_edges, successful_packets ) yield c, o diff --git a/ltc/sim/constants.py b/ltc/sim/constants.py index 7ef3c15..50dd128 100644 --- a/ltc/sim/constants.py +++ b/ltc/sim/constants.py @@ -10,13 +10,19 @@ class Actions(Enum): IDLE = 2 +""" +Time scaling: one TX slot occupies TX_SLOTS mini-slots. +During one TX, CS agents can observe TX_SLOTS times. +""" +TX_SLOTS = 3 + """ Simulation parameters """ MAX_RETRANSMISSION = 8 -SAFE_IDLE_PERIOD = 25 +SAFE_IDLE_PERIOD = 25 * TX_SLOTS SAFE_IDLE_PERIOD_STD = 3 -PENALIZED_IDLE_PERIOD = 25 +PENALIZED_IDLE_PERIOD = 25 * TX_SLOTS """ Rewards and penalties @@ -31,7 +37,7 @@ class Actions(Enum): TX_REWARD_LENGTH = 5 TX_SIZE_THRESHOLD = 5 -TX_SIZE_PENALTY_WINDOW = 25 +TX_SIZE_PENALTY_WINDOW = 25 * TX_SLOTS TX_SIZE_PENALTY = -1.0 TX_LTC_COLLISION_PENALTY = -1.0 TX_COEX_COLLISION_PENALTY = -5.0 @@ -61,12 +67,6 @@ class Actions(Enum): MAX_MACRO_DURATION = 5 LEGACY_TX_DURATION = 5 -""" -Time scaling: one TX slot occupies TX_SLOTS mini-slots. -During one TX, CS agents can observe TX_SLOTS times. -""" -TX_SLOTS = 3 - class ObsIdx: BUFFER_OCCUPANCY_PCT = 0 diff --git a/ltc/utils/structs.py b/ltc/utils/structs.py index b35d09d..8f58c44 100644 --- a/ltc/utils/structs.py +++ b/ltc/utils/structs.py @@ -91,3 +91,4 @@ class Output: active: jax.Array weights_histogram: jax.Array weights_bin_edges: jax.Array + successful_packets: jax.Array From 741e7ad057e935f3f31e01b39b8dd1ffe8c2f6db Mon Sep 17 00:00:00 2001 From: kaszczech Date: Thu, 23 Apr 2026 09:16:21 +0200 Subject: [PATCH 14/24] Successful packet tracking --- ltc/run.py | 42 ++++++++++++++++++++++++++++++------------ ltc/sim/__init__.py | 2 +- ltc/sim/constants.py | 2 +- ltc/sim/sim.py | 11 +++++------ ltc/utils/structs.py | 1 + test/sim/test_sim.py | 3 ++- 6 files changed, 40 insertions(+), 21 deletions(-) diff --git a/ltc/run.py b/ltc/run.py index 1acb757..c3b84a1 100644 --- a/ltc/run.py +++ b/ltc/run.py @@ -13,7 +13,7 @@ from tqdm import trange from ltc.agents import BayesianDDQN, DCF, QNetwork, StochasticVariationalNetwork -from ltc.sim import InitialStateConf, cox_traffic, process_output, simulate +from ltc.sim import InitialStateConf, cox_traffic, process_output, simulate, channel_state_selector from ltc.sim.constants import * from ltc.utils.history import build_history_filename, ensure_clean_git_worktree, get_short_commit_hash from ltc.utils.structs import MacroActionState, ObsTrackerState, ActionDecodingResults, Carry, Output, ObsFeatureInputs, ObsFeatureConfig @@ -204,23 +204,24 @@ def _finalize_macro_accumulators(macro_state, done_now, macro_tx_success_accum, return macro_tx_success_accum, macro_tx_collision_accum, macro_reward_accum -def upgraded_tx_slot_reward(slot_actions, channel_state, successful_packets, prev_ret_c, tx_collision_other_now, tx_packet_mask): +def upgraded_tx_slot_reward(slot_actions, successful_packets, prev_ret_c, tx_collision_other_now, tx_packet_mask, sub_collision_flag): tx_executing = slot_actions == Actions.TX.value - tx_success = jnp.logical_and(tx_executing, successful_packets > 0) - tx_collision = jnp.logical_and(tx_executing & tx_packet_mask, channel_state == -1) + tx_committed = tx_packet_mask + tx_success = jnp.logical_and(tx_committed, jnp.logical_and(~sub_collision_flag, successful_packets > 0)) + tx_collision = jnp.logical_and(tx_committed, sub_collision_flag) tx_collision_max = jnp.logical_and(tx_collision, prev_ret_c >= MAX_RETRANSMISSION) tx_collision_normal = jnp.logical_and(tx_collision, prev_ret_c < MAX_RETRANSMISSION) tx_collision_coex = jnp.logical_and(tx_collision_normal, tx_collision_other_now > 0.5) tx_collision_ltc = jnp.logical_and(tx_collision_normal, tx_collision_other_now <= 0.5) tx_empty = jnp.logical_and( - tx_executing & tx_packet_mask, - jnp.logical_and(channel_state == 1, successful_packets == 0), + jnp.logical_and(tx_committed, ~sub_collision_flag), + successful_packets == 0, ) k_tx_slot = successful_packets.astype(jnp.float32) k_coll_slot = tx_collision.astype(jnp.float32) - tx_success_reward = TX_REWARD * (k_tx_slot * TX_REWARD_LENGTH - 2.0) / TX_REWARD_LENGTH + tx_success_reward = TX_REWARD * k_tx_slot tx_collision_reward = ( tx_collision_ltc.astype(jnp.float32) * (TX_LTC_COLLISION_PENALTY * k_coll_slot) + tx_collision_coex.astype(jnp.float32) * (TX_COEX_COLLISION_PENALTY * k_coll_slot) @@ -291,8 +292,18 @@ def rl_step_coroutine(c, step): slot_executing = macro_remaining > 0 slot_actions = jnp.where(slot_executing, macro_action_types, Actions.IDLE.value) - # TX transmits a packet only at the first mini-slot of each TX unit - tx_packet_mask = (macro_remaining % TX_SLOTS == 0) & (slot_actions == Actions.TX.value) + # Sub-window collision tracking: one packet = TX_SLOTS mini-slots + tx_executing_pre = slot_actions == Actions.TX.value + channel_state_prelim = channel_state_selector(slot_actions) + sub_window_start = (macro_remaining % TX_SLOTS == 0) & slot_executing & tx_executing_pre + sub_collision_flag = jnp.where(sub_window_start, False, c.macro.sub_collision_flag) + sub_collision_flag = jnp.where( + tx_executing_pre & (channel_state_prelim == -1), + True, + sub_collision_flag, + ) + tx_packet_mask = (macro_remaining % TX_SLOTS == 1) & slot_executing & tx_executing_pre + tx_success_mask = tx_packet_mask & ~sub_collision_flag traffic_states, new_frames = traffic_step(c.traffic_states, traffic_keys) ( @@ -302,7 +313,8 @@ def rl_step_coroutine(c, step): packet_seqs, planned_packets, successful_packets, - ) = simulate(c.buffer_states, c.obs_tracker.buffer_birth_steps, new_frames, slot_actions, c.packet_seqs, step, tx_packet_mask) + ) = simulate(c.buffer_states, c.obs_tracker.buffer_birth_steps, new_frames, slot_actions, c.packet_seqs, step, + tx_packet_mask=tx_packet_mask, tx_success_mask=tx_success_mask) is_ltc = jnp.arange(n) < n_drl tx_mask = slot_actions == Actions.TX.value @@ -376,7 +388,7 @@ def rl_step_coroutine(c, step): ) tx_slot_reward, k_tx_slot, k_coll_slot, tx_executing = upgraded_tx_slot_reward( - slot_actions, channel_state, successful_packets, prev_ret_c, tx_collision_other_now, tx_packet_mask + slot_actions, successful_packets, prev_ret_c, tx_collision_other_now, tx_packet_mask, sub_collision_flag ) slot_reward_upgraded = jnp.where(tx_executing, tx_slot_reward, slot_rewards) @@ -385,11 +397,14 @@ def rl_step_coroutine(c, step): macro_tx_collision_accum = c.macro.tx_collision_accum + jnp.where(slot_executing, k_coll_slot, 0.0) macro_reward_accum = c.macro.reward_accum + jnp.where(slot_executing, slot_reward_upgraded, 0.0) + k_total = macro_tx_success_accum + length_scale = (k_total * TX_REWARD_LENGTH - 2.0) / TX_REWARD_LENGTH + length_bonus = jnp.where(k_total > 0, TX_REWARD * length_scale, 0.0) size_penalty = TX_SIZE_PENALTY * jnp.minimum( 1.0, (macro_tx_success_accum - TX_SIZE_THRESHOLD + 1.0) / TX_SIZE_PENALTY_WINDOW ) size_penalty = jnp.where(macro_tx_success_accum >= TX_SIZE_THRESHOLD, size_penalty, 0.0) - rewards = jnp.where(done_now, macro_reward_accum + size_penalty, c.rewards) + rewards = jnp.where(done_now, macro_reward_accum + length_bonus + size_penalty, c.rewards) macro_tx_success_accum, macro_tx_collision_accum, macro_reward_accum = _finalize_macro_accumulators( c.macro, done_now, macro_tx_success_accum, macro_tx_collision_accum, macro_reward_accum @@ -412,6 +427,7 @@ def rl_step_coroutine(c, step): reward_accum=macro_reward_accum, tx_success_accum=macro_tx_success_accum, tx_collision_accum=macro_tx_collision_accum, + sub_collision_flag=sub_collision_flag, ) obs_tracker_state = ObsTrackerState( buffer_birth_steps=buffer_birth_steps, @@ -560,6 +576,7 @@ def setup_args(): macro_reward_accum = jnp.zeros(n, dtype=jnp.float32) macro_tx_success_accum = jnp.zeros(n, dtype=jnp.float32) macro_tx_collision_accum = jnp.zeros(n, dtype=jnp.float32) + macro_sub_collision_flag = jnp.zeros(n, dtype=bool) power_states = jnp.full(n, INITIAL_CAPACITY, dtype=int) channel_state = 0 obs = jnp.zeros((n, window_size, OBS_SIZE), dtype=jnp.float32) @@ -636,6 +653,7 @@ def setup_args(): reward_accum=macro_reward_accum, tx_success_accum=macro_tx_success_accum, tx_collision_accum=macro_tx_collision_accum, + sub_collision_flag=macro_sub_collision_flag, ) obs_tracker_state = ObsTrackerState( buffer_birth_steps=buffer_birth_steps, diff --git a/ltc/sim/__init__.py b/ltc/sim/__init__.py index 325b814..1b9ba0d 100644 --- a/ltc/sim/__init__.py +++ b/ltc/sim/__init__.py @@ -1,3 +1,3 @@ from ltc.sim.traffic import InitialStateConf, ModelState, TrafficModel, cox_traffic from ltc.sim.process_output import process_output -from ltc.sim.sim import simulate +from ltc.sim.sim import simulate, channel_state_selector diff --git a/ltc/sim/constants.py b/ltc/sim/constants.py index 50dd128..4b378b4 100644 --- a/ltc/sim/constants.py +++ b/ltc/sim/constants.py @@ -14,7 +14,7 @@ class Actions(Enum): Time scaling: one TX slot occupies TX_SLOTS mini-slots. During one TX, CS agents can observe TX_SLOTS times. """ -TX_SLOTS = 3 +TX_SLOTS = 20 """ Simulation parameters diff --git a/ltc/sim/sim.py b/ltc/sim/sim.py index 7d6f72a..a669836 100644 --- a/ltc/sim/sim.py +++ b/ltc/sim/sim.py @@ -76,14 +76,13 @@ def add_new_frames(buffer_states, buffer_birth_steps, new_frames, packet_seqs, c ) -def simulate(buffer_states, buffer_birth_steps, new_frames, actions, packet_seqs, current_step, tx_packet_mask): +def simulate(buffer_states, buffer_birth_steps, new_frames, actions, packet_seqs, current_step, + tx_packet_mask, tx_success_mask): channel_state = channel_state_selector(actions) pre_nonempty = jnp.any(buffer_states != EMPTY_PACKET_ID, axis=1) - success_tx = (actions == Actions.TX.value) & (channel_state == 1) - # tx_packet_mask: True only in the first mini-slot of each TX unit (actual packet send) - packet_tx = success_tx & tx_packet_mask - planned_packets = ((actions == Actions.TX.value) & tx_packet_mask).astype(jnp.int32) - successful_packets = (packet_tx & pre_nonempty).astype(jnp.int32) + packet_tx = tx_success_mask & pre_nonempty + planned_packets = tx_packet_mask.astype(jnp.int32) + successful_packets = packet_tx.astype(jnp.int32) queue_size = buffer_states.shape[1] head_mask = jnp.arange(queue_size) == 0 tx_masks = packet_tx[:, None] & head_mask[None, :] diff --git a/ltc/utils/structs.py b/ltc/utils/structs.py index 8f58c44..7b91c50 100644 --- a/ltc/utils/structs.py +++ b/ltc/utils/structs.py @@ -12,6 +12,7 @@ class MacroActionState: reward_accum: jax.Array tx_success_accum: jax.Array tx_collision_accum: jax.Array + sub_collision_flag: jax.Array # True if any collision in current TX_SLOTS sub-window @jax.tree_util.register_dataclass diff --git a/test/sim/test_sim.py b/test/sim/test_sim.py index 5d4c73a..96bbbee 100644 --- a/test/sim/test_sim.py +++ b/test/sim/test_sim.py @@ -79,6 +79,7 @@ def test_simulate(self): actions = jnp.array([Actions.TX.value, Actions.CS.value, Actions.CS.value, Actions.CS.value]) tx_packet_mask = jnp.array([True, False, False, False]) + tx_success_mask = jnp.array([True, False, False, False]) ( result_buffer_states, result_buffer_birth_steps, @@ -88,7 +89,7 @@ def test_simulate(self): successful_packets, ) = simulate( buffer_states, buffer_birth_steps, new_frames, actions, packet_seqs, current_step=10, - tx_packet_mask=tx_packet_mask, + tx_packet_mask=tx_packet_mask, tx_success_mask=tx_success_mask, ) self.assertEqual(result_channel_state, 1) From 99731ed4e462843084599cc022c4b92056bfbd3d Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Thu, 30 Apr 2026 08:05:05 +0200 Subject: [PATCH 15/24] Refactor code, reward calculation, and action space --- ltc/run.py | 533 ++++++++++++++------------------ ltc/sim/constants.py | 9 +- ltc/sim/process_output.py | 134 ++++++-- ltc/sim/sim.py | 4 +- ltc/utils/history.py | 6 +- ltc/utils/structs.py | 13 +- pyproject.toml | 2 +- test/sim/test_action_space.py | 87 ++++++ test/sim/test_process_output.py | 44 ++- test/test_action_space.py | 68 ---- 10 files changed, 474 insertions(+), 426 deletions(-) create mode 100644 test/sim/test_action_space.py delete mode 100644 test/test_action_space.py diff --git a/ltc/run.py b/ltc/run.py index c3b84a1..33a7c35 100644 --- a/ltc/run.py +++ b/ltc/run.py @@ -16,7 +16,7 @@ from ltc.sim import InitialStateConf, cox_traffic, process_output, simulate, channel_state_selector from ltc.sim.constants import * from ltc.utils.history import build_history_filename, ensure_clean_git_worktree, get_short_commit_hash -from ltc.utils.structs import MacroActionState, ObsTrackerState, ActionDecodingResults, Carry, Output, ObsFeatureInputs, ObsFeatureConfig +from ltc.utils.structs import * from ltc.utils.plots import plot_all, plot_first @@ -96,38 +96,21 @@ def schedule_active_stations( def action_space_size(action_space_variant, max_duration): - if action_space_variant in { - 'txcs_duration_set', - 'txcs_to_discrete_duration', - 'txcs_to_continuous_duration', - 'txcs_to_geometric_duration', - }: + if action_space_variant == 'txcs_duration_set': return 2 * max_duration if action_space_variant == 'tx_stage_commit': return 3 raise ValueError(f'Unknown action-space variant: {action_space_variant}') -def decode_drl_actions(raw_actions, staged_tx, keys, action_space_variant, max_duration): +def decode_drl_actions(raw_actions, staged_tx, action_space_variant, max_duration): raw_actions = raw_actions.astype(jnp.int32) - if action_space_variant in { - 'txcs_duration_set', - 'txcs_to_discrete_duration', - 'txcs_to_continuous_duration', - 'txcs_to_geometric_duration', - }: + if action_space_variant == 'txcs_duration_set': is_tx = raw_actions < max_duration duration_idx = jnp.where(is_tx, raw_actions, raw_actions - max_duration) action_type = jnp.where(is_tx, Actions.TX.value, Actions.CS.value) - - if action_space_variant == 'txcs_to_geometric_duration': - p = (duration_idx.astype(jnp.float32) + 1.0) / (max_duration + 1.0) - duration = jax.vmap(lambda k, pp: jax.random.geometric(k, pp))(keys, p).astype(jnp.int32) - duration = jnp.clip(duration, 1, max_duration) - else: - duration = duration_idx + 1 - + duration = duration_idx + 1 return action_type.astype(jnp.int32), duration.astype(jnp.int32), staged_tx if action_space_variant == 'tx_stage_commit': @@ -159,7 +142,100 @@ def decode_legacy_actions(raw_actions, legacy_tx_duration): return action_type, duration -def _update_macro_state_from_ready_agents( +def window_mean(hist, axis=None): + valid = hist >= 0.0 + count = jnp.sum(valid, axis=axis) + return jnp.where(count > 0, jnp.sum(jnp.where(valid, hist, 0.0), axis=axis) / count, -1.0) + + +def roll_append(arr, val): + return jnp.roll(arr, -1, axis=-1).at[..., -1].set(val) + + +def update_packet_tracking( + macro, slot_executing, slot_actions, macro_remaining, + tx_collision_other_now, agent_ready_now, ret_c, buffer_states, +): + tx_executing_pre = slot_actions == Actions.TX.value + channel_state_prelim = channel_state_selector(slot_actions) + + sub_window_start = (macro_remaining % TX_SLOTS == 0) & slot_executing & tx_executing_pre + sub_collision_flag = jnp.where(sub_window_start, False, macro.sub_collision_flag) + sub_collision_flag = jnp.where( + tx_executing_pre & (channel_state_prelim == -1), True, sub_collision_flag, + ) + sub_collision_other_flag = jnp.where(sub_window_start, False, macro.sub_collision_other_flag) + sub_collision_other_flag = jnp.where( + tx_executing_pre & (channel_state_prelim == -1) & (tx_collision_other_now > 0.5), + True, sub_collision_other_flag, + ) + tx_packet_mask = (macro_remaining % TX_SLOTS == 1) & slot_executing & tx_executing_pre + tx_success_mask = tx_packet_mask & ~sub_collision_flag + + collision_with_other = tx_collision_other_now > 0.5 + is_first_tx_slot = agent_ready_now & tx_executing_pre + header_collision = jnp.where(is_first_tx_slot, channel_state_prelim == -1, macro.header_collision) + header_collision_other = jnp.where( + is_first_tx_slot, (channel_state_prelim == -1) & collision_with_other, macro.header_collision_other, + ) + done_now = jnp.logical_and(slot_executing, macro_remaining == 1) + ack_collision_now = done_now & tx_executing_pre & (channel_state_prelim == -1) + ack_collision_other_now = ack_collision_now & collision_with_other + + initial_ret_c = jnp.where(is_first_tx_slot, ret_c, macro.initial_ret_c) + tx_started_empty = jnp.where( + is_first_tx_slot, jnp.all(buffer_states == EMPTY_PACKET_ID, axis=1), macro.tx_started_empty, + ) + + return ( + sub_collision_flag, sub_collision_other_flag, + tx_packet_mask, tx_success_mask, + header_collision, header_collision_other, + done_now, ack_collision_now, ack_collision_other_now, + initial_ret_c, tx_started_empty, + ) + + +def compute_obs_features( + obs_tracker, tx_hist, n_drl, + new_frames, planned_packets, successful_packets, channel_state, tx_mask, + cs_tx_same_type_now, tx_collision_other_now, +): + n = obs_tracker.arrival_hist.shape[0] + + arrival_hist = roll_append(obs_tracker.arrival_hist, new_frames.astype(jnp.float32)) + planned_tx_hist = roll_append(obs_tracker.planned_tx_hist, planned_packets.astype(jnp.float32)) + success_tx_hist = roll_append(obs_tracker.success_tx_hist, successful_packets.astype(jnp.float32)) + channel_busy_hist = roll_append(obs_tracker.channel_busy_hist, (channel_state != 0).astype(jnp.float32)) + collision_hist = roll_append(obs_tracker.collision_hist, (channel_state == -1).astype(jnp.float32)) + tx_hist = jnp.roll(tx_hist, -1, axis=0).at[-1].set(tx_mask.astype(jnp.int32)) + + planned_valid = planned_tx_hist >= 0.0 + planned_sum = jnp.sum(jnp.where(planned_valid, planned_tx_hist, 0.0), axis=1) + success_sum = jnp.sum(jnp.where(planned_valid, success_tx_hist, 0.0), axis=1) + back_pct = jnp.where(planned_sum > 0, success_sum / planned_sum, -1.0) + + tx_valid_rows = tx_hist[:, 0] >= 0 + unique_ltc_tx_window = jnp.where( + jnp.any(tx_valid_rows), + jnp.sum(jnp.any(jnp.where(tx_valid_rows[:, None], tx_hist[:, :n_drl], 0) > 0, axis=0).astype(jnp.float32)), + -1.0, + ) + + features = ObsFeatureInputs( + traffic_mean_arrival_rate=window_mean(arrival_hist, axis=1), + channel_occupancy_pct_window=jnp.broadcast_to(window_mean(channel_busy_hist), (n,)), + channel_collisions_pct_window=jnp.broadcast_to(window_mean(collision_hist), (n,)), + back_pct=back_pct, + unique_ltc_tx_window=jnp.broadcast_to(unique_ltc_tx_window, (n,)), + cs_tx_same_type_now=cs_tx_same_type_now, + tx_collision_other_now=tx_collision_other_now, + ) + + return (arrival_hist, planned_tx_hist, success_tx_hist, channel_busy_hist, collision_hist, tx_hist), features + + +def update_macro_state_from_ready_agents( macro_state, obs_tracker, ready_drl, ready_legacy, decoded_actions, n_drl ): @@ -197,40 +273,6 @@ def _update_macro_state_from_ready_agents( return macro_action_types, macro_remaining, staged_tx -def _finalize_macro_accumulators(macro_state, done_now, macro_tx_success_accum, macro_tx_collision_accum, macro_reward_accum): - macro_tx_success_accum = jnp.where(done_now, 0.0, macro_tx_success_accum) - macro_tx_collision_accum = jnp.where(done_now, 0.0, macro_tx_collision_accum) - macro_reward_accum = jnp.where(done_now, 0.0, macro_reward_accum) - return macro_tx_success_accum, macro_tx_collision_accum, macro_reward_accum - - -def upgraded_tx_slot_reward(slot_actions, successful_packets, prev_ret_c, tx_collision_other_now, tx_packet_mask, sub_collision_flag): - tx_executing = slot_actions == Actions.TX.value - tx_committed = tx_packet_mask - tx_success = jnp.logical_and(tx_committed, jnp.logical_and(~sub_collision_flag, successful_packets > 0)) - tx_collision = jnp.logical_and(tx_committed, sub_collision_flag) - tx_collision_max = jnp.logical_and(tx_collision, prev_ret_c >= MAX_RETRANSMISSION) - tx_collision_normal = jnp.logical_and(tx_collision, prev_ret_c < MAX_RETRANSMISSION) - tx_collision_coex = jnp.logical_and(tx_collision_normal, tx_collision_other_now > 0.5) - tx_collision_ltc = jnp.logical_and(tx_collision_normal, tx_collision_other_now <= 0.5) - tx_empty = jnp.logical_and( - jnp.logical_and(tx_committed, ~sub_collision_flag), - successful_packets == 0, - ) - - k_tx_slot = successful_packets.astype(jnp.float32) - k_coll_slot = tx_collision.astype(jnp.float32) - - tx_success_reward = TX_REWARD * k_tx_slot - tx_collision_reward = ( - tx_collision_ltc.astype(jnp.float32) * (TX_LTC_COLLISION_PENALTY * k_coll_slot) - + tx_collision_coex.astype(jnp.float32) * (TX_COEX_COLLISION_PENALTY * k_coll_slot) - + tx_collision_max.astype(jnp.float32) * TX_MAX_RETRANSMISSION_PENALTY - ) - tx_slot_reward = jnp.where(tx_success, tx_success_reward, 0.0) + tx_collision_reward + jnp.where(tx_empty, TX_EMPTY_BUFFER_PENALTY, 0.0) - return tx_slot_reward, k_tx_slot, k_coll_slot, tx_executing - - def rl_step( drl_step, legacy_step, traffic_step, n, n_drl, n_bins=50, one_shot_step=None, one_shot_target=None, station_change_interval=None, station_change_delta=0, @@ -238,7 +280,7 @@ def rl_step( obs_enable_cs_tx_same_type=False, obs_enable_tx_collision_other=False, action_space_variant='txcs_duration_set', max_duration=5, legacy_tx_duration=5, ): - def rl_step_coroutine(c, step): + def rl_step_fn(c, step): active = schedule_active_stations( c.active, step, @@ -260,7 +302,7 @@ def rl_step_coroutine(c, step): drl_ready = jnp.logical_and(active[:n_drl], ready[:n_drl]) legacy_ready = jnp.logical_and(active[n_drl:], ready[n_drl:]) - drl_states, drl_actions_raw = yield drl_step( + drl_states, drl_actions_raw = drl_step( c.drl_states, drl_keys, c.obs[:n_drl], c.actions[:n_drl], c.rewards[:n_drl], c.terminals[:n_drl], drl_ready ) legacy_states, legacy_actions_raw = legacy_step( @@ -272,7 +314,7 @@ def rl_step_coroutine(c, step): raw_actions = raw_actions.at[n_drl:].set(jnp.where(legacy_ready, legacy_actions_raw, c.actions[n_drl:])) drl_action_types, drl_durations, drl_staged_next = decode_drl_actions( - raw_actions[:n_drl], c.obs_tracker.staged_tx[:n_drl], drl_keys, action_space_variant, max_duration + raw_actions[:n_drl], c.obs_tracker.staged_tx[:n_drl], action_space_variant, max_duration ) legacy_action_types, legacy_durations = decode_legacy_actions(raw_actions[n_drl:], legacy_tx_duration) @@ -284,26 +326,33 @@ def rl_step_coroutine(c, step): legacy_durations=legacy_durations, ) - macro_action_types, macro_remaining, staged_tx = _update_macro_state_from_ready_agents( + macro_action_types, macro_remaining, staged_tx = update_macro_state_from_ready_agents( c.macro, c.obs_tracker, drl_ready, legacy_ready, decoded_actions, n_drl ) slot_executing = macro_remaining > 0 slot_actions = jnp.where(slot_executing, macro_action_types, Actions.IDLE.value) + tx_mask = slot_actions == Actions.TX.value - # Sub-window collision tracking: one packet = TX_SLOTS mini-slots - tx_executing_pre = slot_actions == Actions.TX.value - channel_state_prelim = channel_state_selector(slot_actions) - sub_window_start = (macro_remaining % TX_SLOTS == 0) & slot_executing & tx_executing_pre - sub_collision_flag = jnp.where(sub_window_start, False, c.macro.sub_collision_flag) - sub_collision_flag = jnp.where( - tx_executing_pre & (channel_state_prelim == -1), - True, - sub_collision_flag, + is_ltc = jnp.arange(n) < n_drl + tx_idx = jnp.argmax(tx_mask.astype(jnp.int32)) + cs_tx_same_type_now = (is_ltc[tx_idx] == is_ltc).astype(jnp.float32) + diff_type = is_ltc[:, None] != is_ltc[None, :] + tx_collision_other_now = jnp.any(diff_type & tx_mask[None, :], axis=1).astype(jnp.float32) + + agent_ready_now = jnp.concatenate([drl_ready, legacy_ready]) + ret_c = c.obs[:, -1, ObsIdx.STATUS_RETRY_COUNTER].astype(jnp.int32) + ( + sub_collision_flag, sub_collision_other_flag, + tx_packet_mask, tx_success_mask, + header_collision, header_collision_other, + done_now, ack_collision_now, ack_collision_other_now, + initial_ret_c, tx_started_empty, + ) = update_packet_tracking( + c.macro, slot_executing, slot_actions, macro_remaining, + tx_collision_other_now, agent_ready_now, ret_c, c.buffer_states, ) - tx_packet_mask = (macro_remaining % TX_SLOTS == 1) & slot_executing & tx_executing_pre - tx_success_mask = tx_packet_mask & ~sub_collision_flag traffic_states, new_frames = traffic_step(c.traffic_states, traffic_keys) ( @@ -316,99 +365,44 @@ def rl_step_coroutine(c, step): ) = simulate(c.buffer_states, c.obs_tracker.buffer_birth_steps, new_frames, slot_actions, c.packet_seqs, step, tx_packet_mask=tx_packet_mask, tx_success_mask=tx_success_mask) - is_ltc = jnp.arange(n) < n_drl - tx_mask = slot_actions == Actions.TX.value - tx_idx = jnp.argmax(tx_mask.astype(jnp.int32)) - tx_is_ltc = is_ltc[tx_idx] - cs_tx_same_type_now = (tx_is_ltc == is_ltc).astype(jnp.float32) - diff_type = is_ltc[:, None] != is_ltc[None, :] - tx_collision_other_now = jnp.any(diff_type & tx_mask[None, :], axis=1).astype(jnp.float32) - - arrival_hist = jnp.roll(c.obs_tracker.arrival_hist, -1, axis=1).at[:, -1].set(new_frames.astype(jnp.float32)) - planned_tx_hist = jnp.roll(c.obs_tracker.planned_tx_hist, -1, axis=1).at[:, -1].set(planned_packets.astype(jnp.float32)) - success_tx_hist = jnp.roll(c.obs_tracker.success_tx_hist, -1, axis=1).at[:, -1].set(successful_packets.astype(jnp.float32)) - channel_busy_hist = jnp.roll(c.obs_tracker.channel_busy_hist, -1).at[-1].set((channel_state != 0).astype(jnp.float32)) - collision_hist = jnp.roll(c.obs_tracker.collision_hist, -1).at[-1].set((channel_state == -1).astype(jnp.float32)) - tx_hist = jnp.roll(c.tx_hist, -1, axis=0).at[-1].set(tx_mask.astype(jnp.int32)) - - arrival_valid = arrival_hist >= 0.0 - arrival_count = jnp.sum(arrival_valid, axis=1) - traffic_mean_arrival_rate = jnp.where( - arrival_count > 0, - jnp.sum(jnp.where(arrival_valid, arrival_hist, 0.0), axis=1) / arrival_count, - -1.0, - ) - - busy_valid = channel_busy_hist >= 0.0 - busy_count = jnp.sum(busy_valid) - channel_occupancy_pct_window = jnp.where( - busy_count > 0, - jnp.sum(jnp.where(busy_valid, channel_busy_hist, 0.0)) / busy_count, - -1.0, - ) - - coll_valid = collision_hist >= 0.0 - coll_count = jnp.sum(coll_valid) - channel_collisions_pct_window = jnp.where( - coll_count > 0, - jnp.sum(jnp.where(coll_valid, collision_hist, 0.0)) / coll_count, - -1.0, - ) + (arrival_hist, planned_tx_hist, success_tx_hist, channel_busy_hist, collision_hist, tx_hist), obs_features = \ + compute_obs_features( + c.obs_tracker, c.tx_hist, n_drl, + new_frames, planned_packets, successful_packets, channel_state, tx_mask, + cs_tx_same_type_now, tx_collision_other_now, + ) - planned_valid = planned_tx_hist >= 0.0 - planned_sum = jnp.sum(jnp.where(planned_valid, planned_tx_hist, 0.0), axis=1) - success_sum = jnp.sum(jnp.where(planned_valid, success_tx_hist, 0.0), axis=1) - back_pct = jnp.where(planned_sum > 0, success_sum / planned_sum, -1.0) + packet_collided = tx_packet_mask & sub_collision_flag + k_tx_accum = c.macro.tx_success_accum + successful_packets.astype(jnp.float32) + k_coll_ltc = c.macro.k_coll_ltc + jnp.where(packet_collided & ~sub_collision_other_flag, 1.0, 0.0) + k_coll_coex = c.macro.k_coll_coex + jnp.where(packet_collided & sub_collision_other_flag, 1.0, 0.0) - tx_valid_rows = tx_hist[:, 0] >= 0 - unique_ltc_tx_window = jnp.where( - jnp.any(tx_valid_rows), - jnp.sum(jnp.any(jnp.where(tx_valid_rows[:, None], tx_hist[:, :n_drl], 0) > 0, axis=0).astype(jnp.float32)), - -1.0, - ) - - done_now = jnp.logical_and(slot_executing, macro_remaining == 1) - prev_ret_c = c.obs[:, -1, ObsIdx.STATUS_RETRY_COUNTER].astype(jnp.int32) - obs_features = ObsFeatureInputs( - traffic_mean_arrival_rate=traffic_mean_arrival_rate, - channel_occupancy_pct_window=jnp.broadcast_to(channel_occupancy_pct_window, (n,)), - channel_collisions_pct_window=jnp.broadcast_to(channel_collisions_pct_window, (n,)), - back_pct=back_pct, - unique_ltc_tx_window=jnp.broadcast_to(unique_ltc_tx_window, (n,)), - cs_tx_same_type_now=cs_tx_same_type_now, - tx_collision_other_now=tx_collision_other_now, - ) obs_config = ObsFeatureConfig( enable_cs_tx_same_type=obs_enable_cs_tx_same_type, enable_tx_collision_other=obs_enable_tx_collision_other, ) - obs, slot_rewards, powers = process_output( + obs, rewards_new, powers = process_output( c.buffer_states, buffer_states, buffer_birth_steps, c.power_states, channel_state, c.obs, slot_actions, - c.terminals, reward_key, step, obs_features, obs_config, roll_obs=done_now + c.terminals, reward_key, step, obs_features, obs_config, + roll_obs=done_now, + done_now=done_now, + k_tx=k_tx_accum, + k_coll_ltc=k_coll_ltc, + k_coll_coex=k_coll_coex, + header_collision=header_collision, + header_collision_other=header_collision_other, + ack_collision=ack_collision_now, + ack_collision_other=ack_collision_other_now, + initial_ret_c=initial_ret_c, + tx_started_empty=tx_started_empty, ) - tx_slot_reward, k_tx_slot, k_coll_slot, tx_executing = upgraded_tx_slot_reward( - slot_actions, successful_packets, prev_ret_c, tx_collision_other_now, tx_packet_mask, sub_collision_flag - ) - - slot_reward_upgraded = jnp.where(tx_executing, tx_slot_reward, slot_rewards) + rewards = jnp.where(done_now, rewards_new, c.rewards) - macro_tx_success_accum = c.macro.tx_success_accum + jnp.where(slot_executing, k_tx_slot, 0.0) - macro_tx_collision_accum = c.macro.tx_collision_accum + jnp.where(slot_executing, k_coll_slot, 0.0) - macro_reward_accum = c.macro.reward_accum + jnp.where(slot_executing, slot_reward_upgraded, 0.0) - - k_total = macro_tx_success_accum - length_scale = (k_total * TX_REWARD_LENGTH - 2.0) / TX_REWARD_LENGTH - length_bonus = jnp.where(k_total > 0, TX_REWARD * length_scale, 0.0) - size_penalty = TX_SIZE_PENALTY * jnp.minimum( - 1.0, (macro_tx_success_accum - TX_SIZE_THRESHOLD + 1.0) / TX_SIZE_PENALTY_WINDOW - ) - size_penalty = jnp.where(macro_tx_success_accum >= TX_SIZE_THRESHOLD, size_penalty, 0.0) - rewards = jnp.where(done_now, macro_reward_accum + length_bonus + size_penalty, c.rewards) - - macro_tx_success_accum, macro_tx_collision_accum, macro_reward_accum = _finalize_macro_accumulators( - c.macro, done_now, macro_tx_success_accum, macro_tx_collision_accum, macro_reward_accum - ) + # Reset per-macro accumulators when action completes + k_tx_accum = jnp.where(done_now, 0.0, k_tx_accum) + k_coll_ltc = jnp.where(done_now, 0.0, k_coll_ltc) + k_coll_coex = jnp.where(done_now, 0.0, k_coll_coex) macro_remaining = jnp.where(slot_executing, macro_remaining - 1, macro_remaining) terminals = jnp.logical_or(c.terminals, powers < 0) @@ -424,10 +418,15 @@ def rl_step_coroutine(c, step): macro_state = MacroActionState( remaining=macro_remaining, action_types=macro_action_types, - reward_accum=macro_reward_accum, - tx_success_accum=macro_tx_success_accum, - tx_collision_accum=macro_tx_collision_accum, + tx_success_accum=k_tx_accum, + k_coll_ltc=k_coll_ltc, + k_coll_coex=k_coll_coex, sub_collision_flag=sub_collision_flag, + sub_collision_other_flag=sub_collision_other_flag, + header_collision=header_collision, + header_collision_other=header_collision_other, + initial_ret_c=initial_ret_c, + tx_started_empty=tx_started_empty, ) obs_tracker_state = ObsTrackerState( buffer_birth_steps=buffer_birth_steps, @@ -447,27 +446,12 @@ def rl_step_coroutine(c, step): legacy_states, obs, raw_actions, rewards, terminals, buffer_states, powers, new_frames, channel_state, active, hist, bin_edges, successful_packets ) - yield c, o - - def rl_step_fn(*args, **kwargs): - gen = rl_step_coroutine(*args, **kwargs) - intercepted = next(gen) - return gen.send(intercepted) - - def pre_rl_fn(*args, **kwargs): - gen = rl_step_coroutine(*args, **kwargs) - intercepted = next(gen) - return intercepted - - def post_rl_fn(intermediate, *args, **kwargs): - gen = rl_step_coroutine(*args, **kwargs) - _ = next(gen) - return gen.send(intermediate) + return c, o - return rl_step_fn, pre_rl_fn, post_rl_fn + return rl_step_fn -def setup_args(): +if __name__ == '__main__': parser = argparse.ArgumentParser(description="Run the RL network simulation with configurable parameters.") parser.add_argument('--n', type=int, default=5, help='Initial number of agents in the simulation.') parser.add_argument('--n_final', type=int, help='Final number of agents in the simulation.') @@ -478,19 +462,7 @@ def setup_args(): parser.add_argument('--obs_stats_window', type=int, default=1000, help='Window length for rolling observation statistics.') parser.add_argument('--obs_enable_cs_tx_same_type', action='store_true', default=False, help='Enable CS me-vs-other transmitter feature.') parser.add_argument('--obs_enable_tx_collision_other', action='store_true', default=False, help='Enable TX collision-involves-other feature.') - parser.add_argument( - '--action_space_variant', - type=str, - default='txcs_duration_set', - choices=[ - 'txcs_duration_set', - 'txcs_to_discrete_duration', - 'txcs_to_continuous_duration', - 'txcs_to_geometric_duration', - 'tx_stage_commit', - ], - help='Macro action-space variant.', - ) + parser.add_argument('--action_space_variant', type=str, default='txcs_duration_set', choices=['txcs_duration_set', 'tx_stage_commit'], help='Macro action-space variant.') parser.add_argument('--seed', type=int, default=42, help='Random seed for reproducibility.') parser.add_argument('--save_plots', action='store_true', default=False, help='Whether to save the generated plots.') parser.add_argument('--loc', type=float, default=5.0, help='loc traffic generator parameter.') @@ -502,89 +474,46 @@ def setup_args(): parser.add_argument('--station_change_stop_step', type=int, help='Global step when periodic station changes stop.') parser.add_argument('--traffic_type', type=str, default='saturated', choices=['constant', 'saturated', 'bursty', 'custom'],help="Traffic model to use: 'constant', 'saturated', 'bursty', or 'custom'.") parser.add_argument('--legacy_type', type=str, default='tdma', choices=['q-aloha', 'eb-aloha', 'fw-aloha', 'tdma'], help="Legacy agent type to use: 'q-aloha', 'eb-aloha', 'fw-aloha', or 'tdma'.") + parser.add_argument('--skip_git_check', action='store_true', default=False, help='Skip clean git worktree check.') args = parser.parse_args() - return args - -if __name__ == '__main__': - args = setup_args() - ensure_clean_git_worktree() + if not args.skip_git_check: + ensure_clean_git_worktree() commit_hash = get_short_commit_hash() n_init = args.n - has_n_final = args.n_final is not None - n_final = args.n_final if has_n_final else n_init + n_final = args.n_final if args.n_final is not None else n_init n = max(n_init, n_final) - n_epochs = args.n_epochs - n_steps = args.n_steps - total_steps = n_epochs * n_steps - window_size = args.window_size - queue_size = args.queue_size - obs_stats_window = args.obs_stats_window - seed = args.seed - traffic_type = args.traffic_type - - loc = args.loc - scale = args.scale - f3dB = args.f3dB - station_change_interval = args.station_change_interval - station_change_delta = args.station_change_delta - station_change_start_step = args.station_change_start_step - station_change_stop_step = args.station_change_stop_step - obs_enable_cs_tx_same_type = args.obs_enable_cs_tx_same_type - obs_enable_tx_collision_other = args.obs_enable_tx_collision_other - action_space_variant = args.action_space_variant + total_steps = args.n_epochs * args.n_steps - if queue_size <= 0: + if args.queue_size <= 0: raise ValueError('--queue_size must be positive.') - if obs_stats_window <= 0: + if args.obs_stats_window <= 0: raise ValueError('--obs_stats_window must be positive.') - if station_change_interval is not None and station_change_interval <= 0: + station_change_start_step = args.station_change_start_step + if args.station_change_interval is not None and args.station_change_interval <= 0: raise ValueError('--station_change_interval must be positive.') - if station_change_interval is None: - if station_change_start_step is not None or station_change_stop_step is not None: + if args.station_change_interval is None: + if args.station_change_start_step is not None or args.station_change_stop_step is not None: raise ValueError('--station_change_start_step/--station_change_stop_step require --station_change_interval.') else: - if station_change_delta == 0: + if args.station_change_delta == 0: raise ValueError('--station_change_delta must be non-zero when --station_change_interval is used.') if station_change_start_step is None: - station_change_start_step = station_change_interval + station_change_start_step = args.station_change_interval one_shot_step = None one_shot_target = None - station_change_target = n_final if has_n_final else None - if station_change_interval is None and n_final != n_init: + station_change_target = n_final if args.n_final is not None else None + if args.station_change_interval is None and n_final != n_init: one_shot_step = total_steps // 2 one_shot_target = n_final - key = jax.random.key(seed) - num_actions = action_space_size(action_space_variant, MAX_MACRO_DURATION) - actions = jnp.zeros(n, dtype=int) - packet_seqs = jnp.zeros(n, dtype=jnp.int32) - buffer_states = jnp.full((n, queue_size), EMPTY_PACKET_ID, dtype=jnp.int32) - buffer_birth_steps = jnp.full((n, queue_size), -1, dtype=jnp.int32) - arrival_hist = jnp.full((n, obs_stats_window), -1.0, dtype=jnp.float32) - planned_tx_hist = jnp.full((n, obs_stats_window), -1.0, dtype=jnp.float32) - success_tx_hist = jnp.full((n, obs_stats_window), -1.0, dtype=jnp.float32) - channel_busy_hist = jnp.full(obs_stats_window, -1.0, dtype=jnp.float32) - collision_hist = jnp.full(obs_stats_window, -1.0, dtype=jnp.float32) - tx_hist = jnp.full((obs_stats_window, n), -1, dtype=jnp.int32) - staged_tx = jnp.zeros(n, dtype=jnp.int32) - macro_remaining = jnp.zeros(n, dtype=jnp.int32) - macro_action_types = jnp.full(n, Actions.IDLE.value, dtype=jnp.int32) - macro_reward_accum = jnp.zeros(n, dtype=jnp.float32) - macro_tx_success_accum = jnp.zeros(n, dtype=jnp.float32) - macro_tx_collision_accum = jnp.zeros(n, dtype=jnp.float32) - macro_sub_collision_flag = jnp.zeros(n, dtype=bool) - power_states = jnp.full(n, INITIAL_CAPACITY, dtype=int) - channel_state = 0 - obs = jnp.zeros((n, window_size, OBS_SIZE), dtype=jnp.float32) - rewards = jnp.zeros(n) - terminals = jnp.full(n, False, dtype=bool) - active = jnp.ones(n, dtype=bool).at[n_init:].set(False) - - lr_schedule = optax.cosine_decay_schedule(init_value=1e-4, decay_steps=60000, alpha=0.01) + key = jax.random.key(args.seed) + num_actions = action_space_size(args.action_space_variant, MAX_MACRO_DURATION) + + lr_schedule = optax.cosine_decay_schedule(init_value=1e-4, decay_steps=total_steps, alpha=0.01) optimizer = optax.chain( optax.clip_by_global_norm(1.0), optax.adam(lr_schedule, b1=0.95, b2=0.95) @@ -592,7 +521,7 @@ def setup_args(): drl = BayesianDDQN( q_network=StochasticVariationalNetwork(QNetwork(num_actions, num_layers=1, dim=64, num_heads=4)), - obs_space_shape=obs.shape[1:], + obs_space_shape=(args.window_size, OBS_SIZE), act_space_size=num_actions, optimizer=optimizer, experience_replay_buffer_size=30000, @@ -614,70 +543,82 @@ def setup_args(): key, init_key = jax.random.split(key) - if traffic_type == 'constant': + if args.traffic_type == 'constant': traffic = cox_traffic(f3dB=1.0, loc=-1.0, scale=0.0, initial_state=InitialStateConf.ZERO) - elif traffic_type == 'saturated': + elif args.traffic_type == 'saturated': traffic = cox_traffic(f3dB=1.0, loc=5.0, scale=0.0, initial_state=InitialStateConf.ZERO) - elif traffic_type == 'bursty': + elif args.traffic_type == 'bursty': traffic = cox_traffic(f3dB=0.1, loc=-5.0, scale=5.0, initial_state=InitialStateConf.ZERO) - elif traffic_type == 'custom': - traffic = cox_traffic(f3dB=f3dB, loc=loc, scale=scale, initial_state=InitialStateConf.ZERO) + elif args.traffic_type == 'custom': + traffic = cox_traffic(f3dB=args.f3dB, loc=args.loc, scale=args.scale, initial_state=InitialStateConf.ZERO) else: - raise ValueError(f'Unknown traffic type: {traffic_type}') - + raise ValueError(f'Unknown traffic type: {args.traffic_type}') traffic_states, traffic_step = init_traffic(traffic, init_key, n) - rl_step_fn, _, _ = rl_step( - drl_step, - legacy_step, - traffic_step, - n, - n, + rl_step_fn = jax.jit(rl_step( + drl_step, legacy_step, traffic_step, n, n, one_shot_step=one_shot_step, one_shot_target=one_shot_target, - station_change_interval=station_change_interval, - station_change_delta=station_change_delta, + station_change_interval=args.station_change_interval, + station_change_delta=args.station_change_delta, station_change_start_step=station_change_start_step, - station_change_stop_step=station_change_stop_step, + station_change_stop_step=args.station_change_stop_step, station_change_target=station_change_target, - obs_enable_cs_tx_same_type=obs_enable_cs_tx_same_type, - obs_enable_tx_collision_other=obs_enable_tx_collision_other, - action_space_variant=action_space_variant, + obs_enable_cs_tx_same_type=args.obs_enable_cs_tx_same_type, + obs_enable_tx_collision_other=args.obs_enable_tx_collision_other, + action_space_variant=args.action_space_variant, max_duration=MAX_MACRO_DURATION, legacy_tx_duration=LEGACY_TX_DURATION, - ) - rl_step_fn = jax.jit(rl_step_fn) - macro_state = MacroActionState( - remaining=macro_remaining, - action_types=macro_action_types, - reward_accum=macro_reward_accum, - tx_success_accum=macro_tx_success_accum, - tx_collision_accum=macro_tx_collision_accum, - sub_collision_flag=macro_sub_collision_flag, - ) - obs_tracker_state = ObsTrackerState( - buffer_birth_steps=buffer_birth_steps, - arrival_hist=arrival_hist, - planned_tx_hist=planned_tx_hist, - success_tx_hist=success_tx_hist, - channel_busy_hist=channel_busy_hist, - collision_hist=collision_hist, - staged_tx=staged_tx, - ) + )) + carry = Carry( - drl_states, legacy_states, traffic_states, packet_seqs, buffer_states, - tx_hist, power_states, channel_state, key, obs, actions, rewards, terminals, active, - macro_state, obs_tracker_state + drl_states=drl_states, + legacy_states=legacy_states, + traffic_states=traffic_states, + packet_seqs=jnp.zeros(n, dtype=jnp.int32), + buffer_states=jnp.full((n, args.queue_size), EMPTY_PACKET_ID, dtype=jnp.int32), + tx_hist=jnp.full((args.obs_stats_window, n), -1, dtype=jnp.int32), + power_states=jnp.full(n, INITIAL_CAPACITY, dtype=int), + channel_state=0, + key=key, + obs=jnp.zeros((n, args.window_size, OBS_SIZE), dtype=jnp.float32), + actions=jnp.zeros(n, dtype=int), + rewards=jnp.zeros(n), + terminals=jnp.full(n, False, dtype=bool), + active=jnp.ones(n, dtype=bool).at[n_init:].set(False), + macro=MacroActionState( + remaining=jnp.zeros(n, dtype=jnp.int32), + action_types=jnp.full(n, Actions.IDLE.value, dtype=jnp.int32), + tx_success_accum=jnp.zeros(n, dtype=jnp.float32), + k_coll_ltc=jnp.zeros(n, dtype=jnp.float32), + k_coll_coex=jnp.zeros(n, dtype=jnp.float32), + sub_collision_flag=jnp.zeros(n, dtype=bool), + sub_collision_other_flag=jnp.zeros(n, dtype=bool), + header_collision=jnp.zeros(n, dtype=bool), + header_collision_other=jnp.zeros(n, dtype=bool), + initial_ret_c=jnp.zeros(n, dtype=jnp.int32), + tx_started_empty=jnp.zeros(n, dtype=bool), + ), + obs_tracker=ObsTrackerState( + buffer_birth_steps=jnp.full((n, args.queue_size), -1, dtype=jnp.int32), + arrival_hist=jnp.full((n, args.obs_stats_window), -1.0, dtype=jnp.float32), + planned_tx_hist=jnp.full((n, args.obs_stats_window), -1.0, dtype=jnp.float32), + success_tx_hist=jnp.full((n, args.obs_stats_window), -1.0, dtype=jnp.float32), + channel_busy_hist=jnp.full(args.obs_stats_window, -1.0, dtype=jnp.float32), + collision_hist=jnp.full(args.obs_stats_window, -1.0, dtype=jnp.float32), + staged_tx=jnp.zeros(n, dtype=jnp.int32), + ), ) + all_outputs = [] - for epoch in trange(n_epochs): - global_steps = jnp.arange(epoch * n_steps, (epoch + 1) * n_steps, dtype=jnp.int32) + for epoch in trange(args.n_epochs): + global_steps = jnp.arange(epoch * args.n_steps, (epoch + 1) * args.n_steps, dtype=jnp.int32) carry, output = jax.lax.scan(rl_step_fn, carry, xs=global_steps) all_outputs.append(output) all_outputs = jax.tree.map(lambda *x: jnp.stack(x), *all_outputs) - filename = build_history_filename(n_init, n_final, seed, commit_hash) + filename = build_history_filename(n_init, n_final, args.seed, commit_hash) with lz4.frame.open(filename, 'wb') as f: cloudpickle.dump((carry.drl_states, all_outputs, vars(args)), f) diff --git a/ltc/sim/constants.py b/ltc/sim/constants.py index 4b378b4..c0bec8a 100644 --- a/ltc/sim/constants.py +++ b/ltc/sim/constants.py @@ -31,18 +31,15 @@ class Actions(Enum): EMPTY_BUFFER_REWARD = 0.5 NO_TX_REWARD = 0.0 NO_TX_PENALTY = -1.0 -EMPTY_TX_PENALTY = -0.5 COLLISION_PENALTY = -1.0 -MAX_RETRANSMISSION_PENALTY = -1.0 -TX_REWARD_LENGTH = 5 TX_SIZE_THRESHOLD = 5 -TX_SIZE_PENALTY_WINDOW = 25 * TX_SLOTS +TX_SIZE_PENALTY_WINDOW = 5 TX_SIZE_PENALTY = -1.0 TX_LTC_COLLISION_PENALTY = -1.0 TX_COEX_COLLISION_PENALTY = -5.0 -TX_EMPTY_BUFFER_PENALTY = EMPTY_TX_PENALTY -TX_MAX_RETRANSMISSION_PENALTY = MAX_RETRANSMISSION_PENALTY +TX_EMPTY_BUFFER_PENALTY = -0.5 +TX_MAX_RETRANSMISSION_PENALTY = -1.0 """ Power consumption diff --git a/ltc/sim/process_output.py b/ltc/sim/process_output.py index 5fde586..c769d37 100644 --- a/ltc/sim/process_output.py +++ b/ltc/sim/process_output.py @@ -3,7 +3,6 @@ from ltc.sim.constants import * from ltc.sim.sim import is_buffer_empty -from ltc.utils.structs import ObsFeatureInputs, ObsFeatureConfig def no_transmission(args): @@ -44,50 +43,74 @@ def no_transmission_long(args): def transmission(args): _, _, _, channel_state, _, _ = args - return jax.lax.cond(channel_state == 1, transmission_without_collision, transmission_with_collision, args) + return jax.lax.cond(channel_state == 1, reset_counters, transmission_with_collision, args) def transmission_with_collision(args): _, _, ret_c, _, _, _ = args - return jax.lax.cond(ret_c < MAX_RETRANSMISSION, retransmission, max_retransmission_collision, args) + return jax.lax.cond(ret_c < MAX_RETRANSMISSION, retransmission, reset_counters, args) -def max_retransmission_collision(_): - reward = MAX_RETRANSMISSION_PENALTY +def reset_counters(_): ret_c = 0 no_tx = 0 - return reward, ret_c, no_tx + return 0.0, ret_c, no_tx def retransmission(args): _, _, ret_c, _, _, _ = args - reward = COLLISION_PENALTY ret_c = ret_c + 1 no_tx = 0 - return reward, ret_c, no_tx + return 0.0, ret_c, no_tx -def transmission_without_collision(args): - _, buffer_state, _, _, _, _ = args - return jax.lax.cond(~is_buffer_empty(buffer_state), successful_transmission, empty_buffer_transmission, args) +def tx_macro_reward( + k_tx, k_coll_ltc, k_coll_coex, tx_started_empty, initial_ret_c, + header_collision, ack_collision, header_collision_other, ack_collision_other, +): + hdr_ack = header_collision | ack_collision + k_tx_eff = jnp.where(hdr_ack, 0.0, k_tx) + # When header/ACK collides, attribute all k_tx frames as collisions of the same type + collision_with_other = (header_collision & header_collision_other) | (ack_collision & ack_collision_other) + k_coll_ltc_eff = k_coll_ltc + jnp.where(hdr_ack & ~collision_with_other, k_tx, 0.0) + k_coll_coex_eff = k_coll_coex + jnp.where(hdr_ack & collision_with_other, k_tx, 0.0) -def empty_buffer_transmission(_): - reward = EMPTY_TX_PENALTY - ret_c = 0 - no_tx = 0 - return reward, ret_c, no_tx + k_coll = k_coll_ltc_eff + k_coll_coex_eff + any_collision = k_coll > 0 + success_reward = jnp.where( + k_tx_eff > 0, + TX_REWARD * (k_tx_eff * TX_SLOTS - 2.0) / TX_SLOTS, + 0.0, + ) -def successful_transmission(args): - _, _, ret_c, _, _, _ = args - reward = TX_REWARD - ret_c = 0 - no_tx = 0 - return reward, ret_c, no_tx + size_penalty = jnp.where( + k_tx_eff >= TX_SIZE_THRESHOLD, + TX_SIZE_PENALTY * jnp.minimum(1.0, (k_tx_eff - TX_SIZE_THRESHOLD + 1.0) / TX_SIZE_PENALTY_WINDOW), + 0.0, + ) + + collision_reward = jnp.where( + any_collision, + jnp.where( + initial_ret_c >= MAX_RETRANSMISSION, + TX_MAX_RETRANSMISSION_PENALTY, + TX_LTC_COLLISION_PENALTY * k_coll_ltc_eff + TX_COEX_COLLISION_PENALTY * k_coll_coex_eff, + ), + 0.0, + ) + empty_penalty = jnp.where( + jnp.logical_and(tx_started_empty, ~any_collision), + TX_EMPTY_BUFFER_PENALTY, + 0.0, + ) -def _oldest_packet_age_norm(buffer_state, buffer_birth_state, current_step): + return success_reward + size_penalty + collision_reward + empty_penalty + + +def oldest_packet_age_norm(buffer_state, buffer_birth_state, current_step): valid = buffer_state != EMPTY_PACKET_ID oldest_birth = jnp.min(jnp.where(valid, buffer_birth_state, current_step)) age = jnp.where(jnp.any(valid), current_step - oldest_birth, -1) @@ -95,13 +118,13 @@ def _oldest_packet_age_norm(buffer_state, buffer_birth_state, current_step): return jnp.where(age >= 0, age.astype(jnp.float32) / queue_size, -1.0) -def _build_observation_entry( +def build_observation_entry( buffer_state, new_buffer_state, new_buffer_birth_state, action, channel_state, current_step, obs_features, obs_config, ret_c, no_tx ): buffer_count = jnp.sum((new_buffer_state != EMPTY_PACKET_ID).astype(jnp.int32)) buffer_occupancy_pct = buffer_count.astype(jnp.float32) / new_buffer_state.shape[0] - oldest_packet_age_norm = _oldest_packet_age_norm(new_buffer_state, new_buffer_birth_state, current_step) + packet_age_norm = oldest_packet_age_norm(new_buffer_state, new_buffer_birth_state, current_step) cs_busy = jnp.where(action == Actions.CS.value, channel_state != 0, -1).astype(jnp.float32) cs_tx_same_type = jnp.where( @@ -118,7 +141,7 @@ def _build_observation_entry( return jnp.array([ buffer_occupancy_pct, buffer_count.astype(jnp.float32), - oldest_packet_age_norm, + packet_age_norm, obs_features.traffic_mean_arrival_rate, cs_busy, cs_tx_same_type.astype(jnp.float32), @@ -147,12 +170,33 @@ def process_output_i( obs_features, obs_config, roll_obs, + done_now, + k_tx, + k_coll_ltc, + k_coll_coex, + header_collision, + header_collision_other, + ack_collision, + ack_collision_other, + initial_ret_c, + tx_started_empty, ): ret_c = obs[-1, ObsIdx.STATUS_RETRY_COUNTER].astype(jnp.int32) no_tx = obs[-1, ObsIdx.STATUS_NO_TX_COUNTER].astype(jnp.int32) args = (action, buffer_state, ret_c, channel_state, no_tx, key) - reward, ret_c, no_tx = jax.lax.cond(action == Actions.TX.value, transmission, no_transmission, args) + # Per-slot call: updates ret_c/no_tx for observation; reward only used for CS/IDLE + slot_reward, ret_c, no_tx = jax.lax.cond(action == Actions.TX.value, transmission, no_transmission, args) + + macro_reward = tx_macro_reward( + k_tx, k_coll_ltc, k_coll_coex, tx_started_empty, initial_ret_c, + header_collision, ack_collision, header_collision_other, ack_collision_other, + ) + reward = jnp.where( + action == Actions.TX.value, + jnp.where(done_now, macro_reward, 0.0), + slot_reward, + ) reward = jnp.where(terminal, 0., reward) channel_state = jnp.where(action == Actions.CS.value, channel_state, -1) @@ -167,7 +211,7 @@ def process_output_i( ) ) - obs_t = _build_observation_entry( + obs_t = build_observation_entry( buffer_state, new_buffer_state, new_buffer_birth_state, action, channel_state, current_step, obs_features, obs_config, ret_c, no_tx ) @@ -195,9 +239,31 @@ def process_output( obs_features, obs_config, roll_obs=True, + done_now=False, + k_tx=None, + k_coll_ltc=None, + k_coll_coex=None, + header_collision=None, + header_collision_other=None, + ack_collision=None, + ack_collision_other=None, + initial_ret_c=None, + tx_started_empty=None, ): + n = actions.shape[0] roll_obs = jnp.broadcast_to(jnp.asarray(roll_obs), actions.shape) - return jax.vmap(process_output_i, in_axes=(0, 0, 0, 0, None, 0, 0, 0, None, None, 0, None, 0))( + done_now = jnp.broadcast_to(jnp.asarray(done_now), actions.shape) + k_tx = jnp.broadcast_to(jnp.asarray(k_tx if k_tx is not None else 0.0, dtype=jnp.float32), (n,)) + k_coll_ltc = jnp.broadcast_to(jnp.asarray(k_coll_ltc if k_coll_ltc is not None else 0.0, dtype=jnp.float32), (n,)) + k_coll_coex = jnp.broadcast_to(jnp.asarray(k_coll_coex if k_coll_coex is not None else 0.0, dtype=jnp.float32), (n,)) + header_collision = jnp.broadcast_to(jnp.asarray(header_collision if header_collision is not None else False), (n,)) + header_collision_other = jnp.broadcast_to(jnp.asarray(header_collision_other if header_collision_other is not None else False), (n,)) + ack_collision = jnp.broadcast_to(jnp.asarray(ack_collision if ack_collision is not None else False), (n,)) + ack_collision_other = jnp.broadcast_to(jnp.asarray(ack_collision_other if ack_collision_other is not None else False), (n,)) + initial_ret_c = jnp.broadcast_to(jnp.asarray(initial_ret_c if initial_ret_c is not None else 0, dtype=jnp.int32), (n,)) + tx_started_empty = jnp.broadcast_to(jnp.asarray(tx_started_empty if tx_started_empty is not None else False), (n,)) + + return jax.vmap(process_output_i, in_axes=(0, 0, 0, 0, None, 0, 0, 0, None, None, 0, None, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0))( buffer_states, new_buffer_states, new_buffer_birth_states, @@ -211,4 +277,14 @@ def process_output( obs_features, obs_config, roll_obs, + done_now, + k_tx, + k_coll_ltc, + k_coll_coex, + header_collision, + header_collision_other, + ack_collision, + ack_collision_other, + initial_ret_c, + tx_started_empty, ) diff --git a/ltc/sim/sim.py b/ltc/sim/sim.py index a669836..e276bd0 100644 --- a/ltc/sim/sim.py +++ b/ltc/sim/sim.py @@ -12,7 +12,7 @@ def channel_state_selector(actions): ) -def _cantor_pairing(a, b): +def cantor_pairing(a, b): s = a + b return (s * (s + 1)) // 2 + b @@ -60,7 +60,7 @@ def enqueue_generated_packets(buffer_state, buffer_birth_state, new_frames, stat new_rel = idx - kept_old take_new = jnp.logical_and(idx >= kept_old, new_rel < kept_new) new_local_ids = new_start + jnp.maximum(new_rel, 0) - new_packets = _cantor_pairing(station_id.astype(jnp.int32), new_local_ids) + new_packets = cantor_pairing(station_id.astype(jnp.int32), new_local_ids) new_births = jnp.full((q,), current_step, dtype=jnp.int32) next_state = jnp.where(take_old, old_packets, jnp.where(take_new, new_packets, EMPTY_PACKET_ID)) diff --git a/ltc/utils/history.py b/ltc/utils/history.py index f1dc106..4b731ca 100644 --- a/ltc/utils/history.py +++ b/ltc/utils/history.py @@ -5,7 +5,7 @@ HISTORY_SUFFIX = ".pkl.lz4" -def _run_git_command(*args: str) -> str: +def run_git_command(*args: str) -> str: result = subprocess.run( ["git", *args], check=True, @@ -16,7 +16,7 @@ def _run_git_command(*args: str) -> str: def ensure_clean_git_worktree() -> None: - status = _run_git_command("status", "--porcelain", "--untracked-files=no") + status = run_git_command("status", "--porcelain", "--untracked-files=no") if status: raise RuntimeError( "Refusing to run with tracked uncommitted changes. Commit or stash changes first." @@ -24,7 +24,7 @@ def ensure_clean_git_worktree() -> None: def get_short_commit_hash() -> str: - return _run_git_command("rev-parse", "--short", "HEAD") + return run_git_command("rev-parse", "--short", "HEAD") def build_history_filename(n: int, n_drl: int, seed: int, commit_hash: str | None = None) -> str: diff --git a/ltc/utils/structs.py b/ltc/utils/structs.py index 7b91c50..a285515 100644 --- a/ltc/utils/structs.py +++ b/ltc/utils/structs.py @@ -9,10 +9,15 @@ class MacroActionState: remaining: jax.Array action_types: jax.Array - reward_accum: jax.Array - tx_success_accum: jax.Array - tx_collision_accum: jax.Array - sub_collision_flag: jax.Array # True if any collision in current TX_SLOTS sub-window + tx_success_accum: jax.Array # k_tx: successful subframes + k_coll_ltc: jax.Array # collided subframes vs LTC + k_coll_coex: jax.Array # collided subframes vs coex + sub_collision_flag: jax.Array # True if any collision in current TX_SLOTS sub-window + sub_collision_other_flag: jax.Array # True if any coex collision in current TX_SLOTS sub-window + header_collision: jax.Array # True if first mini-slot of macro TX collided + header_collision_other: jax.Array # True if the header collision was with a non-LTC station + initial_ret_c: jax.Array # ret_c captured at macro TX start + tx_started_empty: jax.Array # True if buffer was empty when macro TX began @jax.tree_util.register_dataclass diff --git a/pyproject.toml b/pyproject.toml index 93dfc63..470fd8c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "ltc" -version = "0.1.0" +version = "0.2.0" requires-python = ">=3.12, <4" description = "Learning to communicate" diff --git a/test/sim/test_action_space.py b/test/sim/test_action_space.py new file mode 100644 index 0000000..2afacdf --- /dev/null +++ b/test/sim/test_action_space.py @@ -0,0 +1,87 @@ +import unittest + +import jax.numpy as jnp + +from ltc.run import decode_drl_actions, decode_legacy_actions +from ltc.sim.constants import Actions, TX_REWARD, TX_SLOTS, TX_LTC_COLLISION_PENALTY, TX_COEX_COLLISION_PENALTY, TX_MAX_RETRANSMISSION_PENALTY, TX_EMPTY_BUFFER_PENALTY, MAX_RETRANSMISSION, TX_SIZE_THRESHOLD, TX_SIZE_PENALTY, TX_SIZE_PENALTY_WINDOW +from ltc.sim.process_output import tx_macro_reward + + +class ActionSpaceDecodeTestCase(unittest.TestCase): + def test_duration_set_decode(self): + raw = jnp.array([0, 4, 5, 9], dtype=jnp.int32) + action_type, duration, staged = decode_drl_actions( + raw, jnp.zeros((4,), dtype=jnp.int32), "txcs_duration_set", 5 + ) + self.assertTrue(jnp.array_equal(action_type, jnp.array([Actions.TX.value, Actions.TX.value, Actions.CS.value, Actions.CS.value]))) + self.assertTrue(jnp.array_equal(duration, jnp.array([1, 5, 1, 5]))) + self.assertTrue(jnp.array_equal(staged, jnp.zeros((4,), dtype=jnp.int32))) + + def test_stage_commit_decode(self): + raw = jnp.array([1, 2, 2], dtype=jnp.int32) # stage, commit, commit + staged = jnp.array([0, 2, 1], dtype=jnp.int32) + action_type, duration, staged_next = decode_drl_actions( + raw, staged, "tx_stage_commit", 5 + ) + self.assertTrue(jnp.array_equal(action_type, jnp.array([Actions.CS.value, Actions.TX.value, Actions.TX.value]))) + self.assertTrue(jnp.array_equal(duration, jnp.array([1, 2, 1]))) + self.assertTrue(jnp.array_equal(staged_next, jnp.array([1, 0, 0]))) + + def test_legacy_decode(self): + raw = jnp.array([Actions.TX.value, Actions.CS.value, Actions.IDLE.value], dtype=jnp.int32) + action_type, duration = decode_legacy_actions(raw, legacy_tx_duration=5) + self.assertTrue(jnp.array_equal(action_type, raw)) + self.assertTrue(jnp.array_equal(duration, jnp.array([5, 1, 1], dtype=jnp.int32))) + + def _macro_reward(self, **kwargs): + defaults = dict( + k_tx=jnp.array(0.0), k_coll_ltc=jnp.array(0.0), k_coll_coex=jnp.array(0.0), + tx_started_empty=jnp.array(False), initial_ret_c=jnp.array(0), + header_collision=jnp.array(False), ack_collision=jnp.array(False), + header_collision_other=jnp.array(False), ack_collision_other=jnp.array(False), + ) + defaults.update(kwargs) + return tx_macro_reward(**defaults) + + def test_tx_macro_reward_branches(self): + # success: k_tx=1, no collision + r = self._macro_reward(k_tx=jnp.array(1.0)) + expected = TX_REWARD * (1.0 * TX_SLOTS - 2.0) / TX_SLOTS + self.assertAlmostEqual(float(r), expected, places=6) + + # coex collision, ret_c < MAX + r = self._macro_reward(k_coll_coex=jnp.array(1.0)) + self.assertAlmostEqual(float(r), TX_COEX_COLLISION_PENALTY * 1.0, places=6) + + # LTC collision, max retransmission + r = self._macro_reward(k_coll_ltc=jnp.array(1.0), initial_ret_c=jnp.array(MAX_RETRANSMISSION)) + self.assertAlmostEqual(float(r), TX_MAX_RETRANSMISSION_PENALTY, places=6) + + # empty buffer, no collision + r = self._macro_reward(tx_started_empty=jnp.array(True)) + self.assertAlmostEqual(float(r), TX_EMPTY_BUFFER_PENALTY, places=6) + + # header collision (LTC): k_tx=3 successful frames all become LTC collisions + r = self._macro_reward( + k_tx=jnp.array(3.0), + header_collision=jnp.array(True), header_collision_other=jnp.array(False), + ) + self.assertAlmostEqual(float(r), TX_LTC_COLLISION_PENALTY * 3.0, places=6) + + # header collision (coex): k_tx=2 successful frames become coex collisions + r = self._macro_reward( + k_tx=jnp.array(2.0), + header_collision=jnp.array(True), header_collision_other=jnp.array(True), + ) + self.assertAlmostEqual(float(r), TX_COEX_COLLISION_PENALTY * 2.0, places=6) + + # ACK collision (LTC) with mixed mid-frame: k_tx=1 successful also becomes LTC collision + r = self._macro_reward( + k_tx=jnp.array(1.0), k_coll_ltc=jnp.array(1.0), + ack_collision=jnp.array(True), ack_collision_other=jnp.array(False), + ) + self.assertAlmostEqual(float(r), TX_LTC_COLLISION_PENALTY * 2.0, places=6) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/sim/test_process_output.py b/test/sim/test_process_output.py index 6cebf01..1673d97 100644 --- a/test/sim/test_process_output.py +++ b/test/sim/test_process_output.py @@ -18,46 +18,34 @@ def test_no_transmission(self): self.assertEqual(reward, 0.0) self.assertEqual(r, 5) - def test_transmission_without_collision_empty_buffer(self): - args = (Actions.TX.value, jnp.array([EMPTY, EMPTY, EMPTY]), 3, 1, 1, KEY) - 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 = (Actions.TX.value, jnp.array([5, EMPTY, EMPTY]), 2, 1, 1, KEY) - reward, r, no_tx = transmission_without_collision(args) - self.assertAlmostEqual(reward, TX_REWARD) - self.assertEqual(r, 0) - def test_transmission_with_collision_retransmission(self): args = (Actions.TX.value, jnp.array([5, EMPTY, EMPTY]), 2, -1, 1, KEY) reward, r, no_tx = transmission_with_collision(args) - self.assertEqual(reward, COLLISION_PENALTY) + self.assertEqual(reward, 0.0) self.assertEqual(r, 3) def test_transmission_with_collision_max_retransmission(self): args = (Actions.TX.value, jnp.array([5, EMPTY, EMPTY]), MAX_RETRANSMISSION, -1, 1, KEY) reward, r, no_tx = transmission_with_collision(args) - self.assertEqual(reward, MAX_RETRANSMISSION_PENALTY) + self.assertEqual(reward, 0.0) self.assertEqual(r, 0) def test_transmission_collision_path(self): args = (Actions.TX.value, jnp.array([5, EMPTY, EMPTY]), 7, -1, 1, KEY) reward, r, no_tx = transmission(args) - self.assertEqual(reward, COLLISION_PENALTY) + self.assertEqual(reward, 0.0) self.assertEqual(r, 8) def test_transmission_successful_path(self): args = (Actions.TX.value, jnp.array([5, EMPTY, EMPTY]), 1, 1, 1, KEY) reward, r, no_tx = transmission(args) - self.assertAlmostEqual(reward, TX_REWARD) + self.assertEqual(reward, 0.0) self.assertEqual(r, 0) def test_transmission_empty_buffer(self): args = (Actions.TX.value, jnp.array([EMPTY, EMPTY, EMPTY]), 1, 1, 1, KEY) reward, r, no_tx = transmission(args) - self.assertEqual(reward, EMPTY_TX_PENALTY) + self.assertEqual(reward, 0.0) self.assertEqual(r, 0) def test_process_rl_output(self): @@ -102,6 +90,18 @@ def test_process_rl_output(self): enable_tx_collision_other=True, ) + # Agent 0: TX, collision, ret_c=8 >= MAX_RETRANSMISSION, coex collision → TX_MAX_RETRANSMISSION_PENALTY + done_now = jnp.array([True, True, True, True]) + k_tx = jnp.zeros(4, dtype=jnp.float32) + k_coll_ltc = jnp.zeros(4, dtype=jnp.float32) + k_coll_coex = jnp.array([1.0, 0.0, 1.0, 1.0], dtype=jnp.float32) # coex collision for TX agents + header_collision = jnp.zeros(4, dtype=bool) + header_collision_other = jnp.zeros(4, dtype=bool) + ack_collision = jnp.zeros(4, dtype=bool) + ack_collision_other = jnp.zeros(4, dtype=bool) + initial_ret_c = jnp.array([8, 0, 8, 8], dtype=jnp.int32) + tx_started_empty = jnp.array([True, False, False, True], dtype=bool) + result_obs, result_R, power = process_output( buffer_states=buffer_states, new_buffer_states=new_buffer_states, @@ -115,6 +115,16 @@ def test_process_rl_output(self): current_step=10, obs_features=obs_features, obs_config=obs_config, + done_now=done_now, + k_tx=k_tx, + k_coll_ltc=k_coll_ltc, + k_coll_coex=k_coll_coex, + header_collision=header_collision, + header_collision_other=header_collision_other, + ack_collision=ack_collision, + ack_collision_other=ack_collision_other, + initial_ret_c=initial_ret_c, + tx_started_empty=tx_started_empty, ) self.assertEqual(result_obs.shape[-1], OBS_SIZE) diff --git a/test/test_action_space.py b/test/test_action_space.py deleted file mode 100644 index e3dbe49..0000000 --- a/test/test_action_space.py +++ /dev/null @@ -1,68 +0,0 @@ -import unittest - -import jax -import jax.numpy as jnp - -from ltc.run import decode_drl_actions, decode_legacy_actions, upgraded_tx_slot_reward -from ltc.sim.constants import Actions - - -class ActionSpaceDecodeTestCase(unittest.TestCase): - def test_duration_set_decode(self): - raw = jnp.array([0, 4, 5, 9], dtype=jnp.int32) - keys = jax.random.split(jax.random.PRNGKey(0), 4) - action_type, duration, staged = decode_drl_actions( - raw, jnp.zeros((4,), dtype=jnp.int32), keys, "txcs_duration_set", 5 - ) - self.assertTrue(jnp.array_equal(action_type, jnp.array([Actions.TX.value, Actions.TX.value, Actions.CS.value, Actions.CS.value]))) - self.assertTrue(jnp.array_equal(duration, jnp.array([1, 5, 1, 5]))) - self.assertTrue(jnp.array_equal(staged, jnp.zeros((4,), dtype=jnp.int32))) - - def test_stage_commit_decode(self): - raw = jnp.array([1, 2, 2], dtype=jnp.int32) # stage, commit, commit - staged = jnp.array([0, 2, 1], dtype=jnp.int32) - keys = jax.random.split(jax.random.PRNGKey(1), 3) - action_type, duration, staged_next = decode_drl_actions( - raw, staged, keys, "tx_stage_commit", 5 - ) - self.assertTrue(jnp.array_equal(action_type, jnp.array([Actions.CS.value, Actions.TX.value, Actions.TX.value]))) - self.assertTrue(jnp.array_equal(duration, jnp.array([1, 2, 1]))) - self.assertTrue(jnp.array_equal(staged_next, jnp.array([1, 0, 0]))) - - def test_geometric_decode_range(self): - raw = jnp.array([0, 9], dtype=jnp.int32) - keys = jax.random.split(jax.random.PRNGKey(2), 2) - action_type, duration, _ = decode_drl_actions( - raw, jnp.zeros((2,), dtype=jnp.int32), keys, "txcs_to_geometric_duration", 5 - ) - self.assertTrue(jnp.array_equal(action_type, jnp.array([Actions.TX.value, Actions.CS.value]))) - self.assertTrue(jnp.all(duration >= 1)) - self.assertTrue(jnp.all(duration <= 5)) - - def test_legacy_decode(self): - raw = jnp.array([Actions.TX.value, Actions.CS.value, Actions.IDLE.value], dtype=jnp.int32) - action_type, duration = decode_legacy_actions(raw, legacy_tx_duration=5) - self.assertTrue(jnp.array_equal(action_type, raw)) - self.assertTrue(jnp.array_equal(duration, jnp.array([5, 1, 1], dtype=jnp.int32))) - - def test_upgraded_tx_slot_reward_branches(self): - slot_actions = jnp.array([Actions.TX.value, Actions.TX.value, Actions.TX.value, Actions.TX.value], dtype=jnp.int32) - channel_state = jnp.array([1, -1, -1, 1], dtype=jnp.int32) - successful_packets = jnp.array([1, 0, 0, 0], dtype=jnp.int32) - prev_ret_c = jnp.array([0, 0, 8, 0], dtype=jnp.int32) - tx_collision_other_now = jnp.array([0.0, 1.0, 0.0, 0.0], dtype=jnp.float32) - rewards, k_tx, k_coll, tx_exec = upgraded_tx_slot_reward( - slot_actions, channel_state, successful_packets, prev_ret_c, tx_collision_other_now - ) - self.assertTrue(jnp.array_equal(tx_exec, jnp.ones((4,), dtype=bool))) - self.assertTrue(jnp.array_equal(k_tx, jnp.array([1.0, 0.0, 0.0, 0.0], dtype=jnp.float32))) - self.assertTrue(jnp.array_equal(k_coll, jnp.array([0.0, 1.0, 1.0, 0.0], dtype=jnp.float32))) - # success, coex collision, max retransmission, empty-buffer TX - self.assertAlmostEqual(float(rewards[0]), 0.6, places=6) - self.assertAlmostEqual(float(rewards[1]), -5.0, places=6) - self.assertAlmostEqual(float(rewards[2]), -1.0, places=6) - self.assertAlmostEqual(float(rewards[3]), -0.5, places=6) - - -if __name__ == "__main__": - unittest.main() From 78bb71aaaddae5a7af6db493e52c9a955db12ea6 Mon Sep 17 00:00:00 2001 From: kaszczech Date: Wed, 6 May 2026 08:47:23 +0200 Subject: [PATCH 16/24] Plotting actions, channel state and rewards --- ltc/utils/plot_actions.py | 393 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 393 insertions(+) create mode 100644 ltc/utils/plot_actions.py diff --git a/ltc/utils/plot_actions.py b/ltc/utils/plot_actions.py new file mode 100644 index 0000000..1c6e4a4 --- /dev/null +++ b/ltc/utils/plot_actions.py @@ -0,0 +1,393 @@ +import argparse + +import cloudpickle +import lz4.frame +import matplotlib.patches as mpatches +import matplotlib.pyplot as plt +import numpy as np + +from ltc.sim.constants import Actions, MAX_MACRO_DURATION +from ltc.utils.history import parse_history_filename, unpack_history + +TX_COLOR = np.array([0.2, 0.75, 0.2 ]) # green +CS_COLOR = np.array([1.0, 0.85, 0.0 ]) # yellow +IDLE_COLOR = np.array([0.92, 0.92, 0.92]) # light gray +CH_SUCCESS = np.array([0.2, 0.75, 0.2 ]) # green +CH_IDLE = np.array([1.0, 0.85, 0.0 ]) # yellow +CH_COLLISION= np.array([0.9, 0.15, 0.15]) # red + + +def _x_ticks(ax, n_steps): + ticks = np.arange(0, n_steps, max(1, n_steps // 20)) + ax.set_xticks(ticks) + ax.set_xticklabels(ticks, fontsize=7, rotation=45) + + +def _vlines(ax, n_steps, every=10): + for x in range(every, n_steps, every): + ax.axvline(x - 0.5, color='black', linewidth=0.4, alpha=0.4) + + +def decode_action_type(raw_action, action_space_variant, max_duration): + if action_space_variant in { + 'txcs_duration_set', + 'txcs_to_discrete_duration', + 'txcs_to_continuous_duration', + 'txcs_to_geometric_duration', + }: + return np.where(raw_action < max_duration, Actions.TX.value, Actions.CS.value) + if action_space_variant == 'tx_stage_commit': + return np.where(raw_action == 0, Actions.CS.value, Actions.TX.value) + return raw_action + + +# ── 1. Actions ──────────────────────────────────────────────────────────────── + +def plot_action_slots(history, metadata, n, n_drl, seed, epoch=0, max_steps=300): + action_space_variant = (metadata or {}).get('action_space_variant', 'txcs_duration_set') + max_duration = (metadata or {}).get('max_duration', MAX_MACRO_DURATION) + + raw = np.array(history.actions[epoch, :max_steps]) # (steps, n_agents) + n_steps, n_agents = raw.shape + has_slot = hasattr(history, 'slot_actions') and history.slot_actions is not None + if has_slot: + types = np.array(history.slot_actions[epoch, :max_steps]).T + else: + types = decode_action_type(raw.T, action_space_variant, max_duration) + + grid = np.zeros((n_agents, n_steps, 3)) + grid[types == Actions.TX.value] = TX_COLOR + grid[types == Actions.CS.value] = CS_COLOR + grid[types == Actions.IDLE.value] = IDLE_COLOR + + fig, ax = plt.subplots(figsize=(max(14, n_steps / 30), max(2, n_agents * 0.7 + 1.2))) + ax.imshow(grid, aspect='auto', interpolation='nearest', origin='upper') + + ax.set_yticks(range(n_agents)) + ax.set_yticklabels( + [f'DRL {i}' if i < n_drl else f'Legacy {i - n_drl}' for i in range(n_agents)], + fontsize=9, + ) + _x_ticks(ax, n_steps) + ax.set_xlabel('Slot', fontsize=10) + ax.set_ylabel('Agent', fontsize=10) + ax.set_title(f'Actions — epoch {epoch}, {n_steps} slots [{action_space_variant}]', fontsize=11) + ax.legend(handles=[ + mpatches.Patch(color=TX_COLOR, label='TX'), + mpatches.Patch(color=CS_COLOR, label='CS'), + mpatches.Patch(color=IDLE_COLOR, label='IDLE'), + ], loc='upper right', fontsize=9) + + plt.tight_layout() + out = f'action_slots_{n}_{n_drl}_{seed}_e{epoch}.png' + plt.savefig(out, dpi=150, bbox_inches='tight') + print(f'Saved: {out}') + plt.show() + + +# ── 2. Channel state ────────────────────────────────────────────────────────── + +def plot_channel_slots(history, n, n_drl, seed, epoch=0, max_steps=300): + ch = np.array(history.channel_state[epoch, :max_steps]) # (steps,) values: -1, 0, 1 + n_steps = len(ch) + + grid = np.zeros((1, n_steps, 3)) + grid[0, ch == 1] = CH_SUCCESS + grid[0, ch == 0] = CH_IDLE + grid[0, ch == -1] = CH_COLLISION + + fig, ax = plt.subplots(figsize=(max(14, n_steps / 30), 1.4)) + ax.imshow(grid, aspect='auto', interpolation='nearest', origin='upper') + + ax.set_yticks([0]) + ax.set_yticklabels(['Channel'], fontsize=9) + _x_ticks(ax, n_steps) + ax.set_xlabel('Slot', fontsize=10) + ax.set_title(f'Channel state — epoch {epoch}, {n_steps} slots', fontsize=11) + ax.legend(handles=[ + mpatches.Patch(color=CH_SUCCESS, label='TX success'), + mpatches.Patch(color=CH_IDLE, label='Idle'), + mpatches.Patch(color=CH_COLLISION, label='Collision'), + ], loc='upper right', fontsize=9) + + plt.tight_layout() + out = f'channel_slots_{n}_{n_drl}_{seed}_e{epoch}.png' + plt.savefig(out, dpi=150, bbox_inches='tight') + print(f'Saved: {out}') + plt.show() + + +# ── 3. Rewards ──────────────────────────────────────────────────────────────── + +def plot_reward_slots(history, n, n_drl, seed, epoch=0, max_steps=300): + rewards = np.array(history.rewards[epoch, :max_steps]) # (steps, n_agents) + n_steps, n_agents = rewards.shape + + # diverging colormap: red=negative, white=0, green=positive + vmax = max(np.abs(rewards).max(), 1e-6) + + fig, axes = plt.subplots( + n_agents, 1, + figsize=(max(14, n_steps / 30), max(3, n_agents * 1.1 + 0.8)), + sharex=True, + ) + if n_agents == 1: + axes = [axes] + + for i, ax in enumerate(axes): + r = rewards[:, i].reshape(1, -1) + im = ax.imshow(r, aspect='auto', interpolation='nearest', + cmap='RdYlGn', vmin=-vmax, vmax=vmax, origin='upper') + label = f'DRL {i}' if i < n_drl else f'Legacy {i - n_drl}' + ax.set_yticks([0]) + ax.set_yticklabels([label], fontsize=9) + if i == n_agents - 1: + _x_ticks(ax, n_steps) + ax.set_xlabel('Slot', fontsize=10) + + fig.colorbar(im, ax=axes, orientation='vertical', fraction=0.015, pad=0.01, label='Reward') + fig.suptitle(f'Rewards per agent — epoch {epoch}, {n_steps} slots', fontsize=11, y=1.01) + + plt.tight_layout() + out = f'reward_slots_{n}_{n_drl}_{seed}_e{epoch}.png' + plt.savefig(out, dpi=150, bbox_inches='tight') + print(f'Saved: {out}') + plt.show() + + +# ── 4. Combined ─────────────────────────────────────────────────────────────── + +def plot_combined(history, metadata, n, n_drl, seed, epoch=0, max_steps=300, grid_every=10): + action_space_variant = (metadata or {}).get('action_space_variant', 'txcs_duration_set') + max_duration = (metadata or {}).get('max_duration', MAX_MACRO_DURATION) + + raw = np.array(history.actions[epoch, :max_steps]) # (steps, n_agents) + has_slot = hasattr(history, 'slot_actions') and history.slot_actions is not None + slot_act = np.array(history.slot_actions[epoch, :max_steps]) if has_slot else None + ch = np.array(history.channel_state[epoch, :max_steps]) # (steps,) + rewards = np.array(history.rewards[epoch, :max_steps]) # (steps, n_agents) + has_packets = hasattr(history, 'successful_packets') and history.successful_packets is not None + pkts = np.array(history.successful_packets[epoch, :max_steps]) if has_packets else None + + n_steps, n_agents = raw.shape + + # row layout: actions×n_agents | channel×1 | rewards×n_agents | packets×n_agents + n_rows = n_agents + 1 + n_agents + (n_agents if has_packets else 0) + heights = [1] * n_agents + [1] + [1] * n_agents + ([1] * n_agents if has_packets else []) + + fig, axes = plt.subplots( + n_rows, 1, + figsize=(max(16, n_steps / 25), n_rows * 0.55 + 1.0), + gridspec_kw={'height_ratios': heights, 'hspace': 0.08}, + sharex=True, + ) + + # ── actions ── (slot_actions if available, else decoded raw macro-actions) + if has_slot: + types = slot_act.T # (n_agents, n_steps) — already TX/CS/IDLE values + else: + types = decode_action_type(raw.T, action_space_variant, max_duration) + action_grid = np.zeros((n_agents, n_steps, 3)) + action_grid[types == Actions.TX.value] = TX_COLOR + action_grid[types == Actions.CS.value] = CS_COLOR + action_grid[types == Actions.IDLE.value] = IDLE_COLOR + + for i in range(n_agents): + ax = axes[i] + ax.imshow(action_grid[i:i+1], aspect='auto', interpolation='nearest', origin='upper') + _vlines(ax, n_steps, grid_every) + label = f'DRL {i}' if i < n_drl else f'Legacy {i - n_drl}' + ax.set_yticks([0]) + ax.set_yticklabels([label], fontsize=8) + ax.tick_params(left=False) + if i == 0: + ax.set_title( + f'Simulation overview — epoch {epoch}, {n_steps} slots [{action_space_variant}]', + fontsize=11, pad=6, + ) + + # divider label + axes[n_agents - 1].spines['bottom'].set_linewidth(1.5) + + # ── channel ── + ch_grid = np.zeros((1, n_steps, 3)) + ch_grid[0, ch == 1] = CH_SUCCESS + ch_grid[0, ch == 0] = CH_IDLE + ch_grid[0, ch == -1] = CH_COLLISION + + ax_ch = axes[n_agents] + ax_ch.imshow(ch_grid, aspect='auto', interpolation='nearest', origin='upper') + _vlines(ax_ch, n_steps, grid_every) + ax_ch.set_yticks([0]) + ax_ch.set_yticklabels(['Channel'], fontsize=8) + ax_ch.tick_params(left=False) + ax_ch.spines['bottom'].set_linewidth(1.5) + + # ── rewards ── + vmax = max(np.abs(rewards).max(), 1e-6) + reward_ims = [] + for i in range(n_agents): + ax = axes[n_agents + 1 + i] + r = rewards[:, i].reshape(1, -1) + im = ax.imshow(r, aspect='auto', interpolation='nearest', + cmap='RdYlGn', vmin=-vmax, vmax=vmax, origin='upper') + _vlines(ax, n_steps, grid_every) + reward_ims.append(im) + label = f'R DRL {i}' if i < n_drl else f'R Leg {i - n_drl}' + ax.set_yticks([0]) + ax.set_yticklabels([label], fontsize=8) + ax.tick_params(left=False) + + fig.colorbar(reward_ims[-1], ax=axes[n_agents + 1:n_agents + 1 + n_agents], orientation='vertical', + fraction=0.012, pad=0.01, label='Reward') + + # ── packets ── + if has_packets: + PKT_COLOR = np.array([0.2, 0.75, 0.2]) + NO_PKT_COLOR = np.array([0.93, 0.93, 0.93]) + for i in range(n_agents): + ax = axes[n_agents + 1 + n_agents + i] + pkt_grid = np.where(pkts[:, i:i+1].T == 1, 1.0, 0.0) + rgb = np.where(pkts[:, i:i+1].T[..., None] == 1, PKT_COLOR, NO_PKT_COLOR) + ax.imshow(rgb.reshape(1, n_steps, 3), aspect='auto', interpolation='nearest', origin='upper') + _vlines(ax, n_steps, grid_every) + label = f'Pkt DRL {i}' if i < n_drl else f'Pkt Leg {i - n_drl}' + ax.set_yticks([0]) + ax.set_yticklabels([label], fontsize=8) + ax.tick_params(left=False) + + # ── x ticks only on last row ── + _x_ticks(axes[-1], n_steps) + axes[-1].set_xlabel('Slot', fontsize=10) + + # ── shared legend (actions + channel) ── + legend_handles = [ + mpatches.Patch(color=TX_COLOR, label='TX'), + mpatches.Patch(color=CS_COLOR, label='CS'), + mpatches.Patch(color=IDLE_COLOR, label='IDLE'), + mpatches.Patch(color=CH_SUCCESS, label='Ch: success'), + mpatches.Patch(color=CH_IDLE, label='Ch: idle'), + mpatches.Patch(color=CH_COLLISION, label='Ch: collision'), + ] + if has_packets: + legend_handles.append(mpatches.Patch(color=np.array([0.2, 0.75, 0.2]), label='Packet sent')) + legend_handles.append(mpatches.Patch(color=np.array([0.93, 0.93, 0.93]), label='No packet')) + fig.legend(handles=legend_handles, loc='upper right', fontsize=8, ncol=2, bbox_to_anchor=(0.88, 1.0)) + + out = f'overview_{n}_{n_drl}_{seed}_e{epoch}.png' + plt.savefig(out, dpi=150, bbox_inches='tight') + print(f'Saved: {out}') + plt.show() + + +# ── 5. Raw data inspector ───────────────────────────────────────────────────── + +def inspect(history, metadata, n, n_drl, seed, epoch=0, max_steps=20): + from ltc.sim.constants import ObsIdx + + n_epochs = history.actions.shape[0] + n_steps_total = history.actions.shape[1] + n_agents = history.actions.shape[2] + + print("=" * 70) + print(f"FILE: n={n} n_drl={n_drl} seed={seed}") + print(f"SHAPE: {n_epochs} epochs × {n_steps_total} steps × {n_agents} agents") + print(f"ARGS: {metadata}") + print("=" * 70) + + raw = np.array(history.actions[epoch, :max_steps]) # (steps, agents) + ch = np.array(history.channel_state[epoch, :max_steps]) # (steps,) + rewards = np.array(history.rewards[epoch, :max_steps]) # (steps, agents) + obs = np.array(history.observations[epoch, :max_steps]) # (steps, agents, win, obs_size) + buf = np.array(history.buffer_states[epoch, :max_steps]) # (steps, agents, queue) + has_pkts = hasattr(history, 'successful_packets') and history.successful_packets is not None + pkts_raw = np.array(history.successful_packets[epoch, :max_steps]) if has_pkts else None + has_slot = hasattr(history, 'slot_actions') and history.slot_actions is not None + slot_raw = np.array(history.slot_actions[epoch, :max_steps]) if has_slot else None + + action_space_variant = (metadata or {}).get('action_space_variant', 'txcs_duration_set') + max_duration = (metadata or {}).get('max_duration', MAX_MACRO_DURATION) + macro_types = decode_action_type(raw.T, action_space_variant, max_duration).T # (steps, agents) + + action_names = {Actions.TX.value: 'TX', Actions.CS.value: 'CS', Actions.IDLE.value: 'IDLE'} + ch_names = {1: 'SUCCESS', 0: 'IDLE', -1: 'COLLISION'} + + print(f"\n── Epoch {epoch}, first {max_steps} slots ──\n") + col_hdr = "raw/macro/slot/rew/pkt" if has_pkts else "raw/macro/slot/rew" + header = f"{'Slot':>5} {'Channel':>10} " + " ".join(f"A{i}({col_hdr})" for i in range(n_agents)) + print(header) + print("-" * len(header)) + + for t in range(min(max_steps, n_steps_total)): + ch_str = ch_names.get(int(ch[t]), str(ch[t])) + agent_cols = [] + for i in range(n_agents): + macro_str = action_names.get(int(macro_types[t, i]), str(macro_types[t, i])) + slot_str = action_names.get(int(slot_raw[t, i]), '?') if has_slot else '?' + pkt_part = f"/{int(pkts_raw[t,i])}" if has_pkts else "" + agent_cols.append(f"{int(raw[t,i]):>3}/{macro_str:<4}/{slot_str:<4}/{rewards[t,i]:>+6.2f}{pkt_part}") + print(f"{t:>5} {ch_str:>10} " + " ".join(agent_cols)) + + print(f"\n── Observations (last obs window, epoch {epoch}, first {max_steps} steps) ──\n") + obs_names = [name for name, _ in sorted(vars(ObsIdx).items(), key=lambda x: x[1]) + if not name.startswith('_')] + for t in range(min(max_steps, n_steps_total)): + print(f" Step {t}:") + for i in range(n_agents): + label = f'DRL {i}' if i < n_drl else f'Legacy {i-n_drl}' + vals = obs[t, i, -1] # last window entry + obs_str = " ".join(f"{obs_names[j]}={vals[j]:.2f}" for j in range(len(obs_names))) + print(f" [{label}] {obs_str}") + + print(f"\n── Buffer states (epoch {epoch}, first {max_steps} steps) ──\n") + for t in range(min(max_steps, n_steps_total)): + bufs = " ".join(f"A{i}:{list(buf[t,i])}" for i in range(n_agents)) + print(f" Step {t}: {bufs}") + + print("\n── Summary statistics ──\n") + has_pkts_hist = hasattr(history, 'successful_packets') and history.successful_packets is not None + for i in range(n_agents): + label = f'DRL {i}' if i < n_drl else f'Legacy {i-n_drl}' + all_types = decode_action_type( + np.array(history.actions[:, :, i]).flatten(), + action_space_variant, max_duration, + ) + tx_pct = 100 * np.mean(all_types == Actions.TX.value) + cs_pct = 100 * np.mean(all_types == Actions.CS.value) + total_r = float(np.array(history.rewards[:, :, i]).sum()) + mean_r = float(np.array(history.rewards[:, :, i]).mean()) + pkt_str = '' + if has_pkts_hist: + total_pkts = int(np.array(history.successful_packets[:, :, i]).sum()) + pkt_str = f' packets={total_pkts}' + print(f" {label:12s} TX={tx_pct:.1f}% CS={cs_pct:.1f}% " + f"total_reward={total_r:.1f} mean_reward={mean_r:.4f}{pkt_str}") + + ch_all = np.array(history.channel_state).flatten() + print(f"\n Channel: success={100*np.mean(ch_all==1):.1f}% " + f"idle={100*np.mean(ch_all==0):.1f}% " + f"collision={100*np.mean(ch_all==-1):.1f}%") + + +# ── CLI ─────────────────────────────────────────────────────────────────────── + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description='Analyze simulation history file') + parser.add_argument('filename', help='Path to .pkl.lz4 history file') + parser.add_argument('--epoch', type=int, default=0, help='Epoch index to visualize') + parser.add_argument('--steps', type=int, default=300, help='Max slots to show') + parser.add_argument('--grid', type=int, default=10, help='Vertical grid every N slots (0=off)') + parser.add_argument('--inspect', action='store_true', help='Print raw data instead of plotting') + parser.add_argument('--inspect_steps', type=int, default=20, help='Slots to print in inspect mode') + args = parser.parse_args() + + n, n_drl, seed, _ = parse_history_filename(args.filename) + + with lz4.frame.open(args.filename, 'rb') as f: + _, history, metadata = unpack_history(cloudpickle.load(f)) + + if args.inspect: + inspect(history, metadata, n, n_drl, seed, epoch=args.epoch, max_steps=args.inspect_steps) + else: + plot_combined(history, metadata, n, n_drl, seed, epoch=args.epoch, max_steps=args.steps, + grid_every=args.grid if args.grid > 0 else 10**9) From e6cc9fa5af65456a7d19829044ac4c35c2043a7d Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Wed, 6 May 2026 16:48:46 +0200 Subject: [PATCH 17/24] Fix rewards --- ltc/run.py | 15 ++++-- ltc/sim/constants.py | 4 +- ltc/sim/process_output.py | 92 ++++++++++++++++++++------------- ltc/utils/structs.py | 1 + test/sim/test_action_space.py | 43 ++++++++------- test/sim/test_process_output.py | 61 +++++++++++++--------- 6 files changed, 131 insertions(+), 85 deletions(-) diff --git a/ltc/run.py b/ltc/run.py index 33a7c35..01304aa 100644 --- a/ltc/run.py +++ b/ltc/run.py @@ -374,6 +374,7 @@ def rl_step_fn(c, step): packet_collided = tx_packet_mask & sub_collision_flag k_tx_accum = c.macro.tx_success_accum + successful_packets.astype(jnp.float32) + k_tx_planned = c.macro.tx_planned_accum + tx_packet_mask.astype(jnp.float32) k_coll_ltc = c.macro.k_coll_ltc + jnp.where(packet_collided & ~sub_collision_other_flag, 1.0, 0.0) k_coll_coex = c.macro.k_coll_coex + jnp.where(packet_collided & sub_collision_other_flag, 1.0, 0.0) @@ -381,12 +382,13 @@ def rl_step_fn(c, step): enable_cs_tx_same_type=obs_enable_cs_tx_same_type, enable_tx_collision_other=obs_enable_tx_collision_other, ) - obs, rewards_new, powers = process_output( + obs, rewards, powers = process_output( c.buffer_states, buffer_states, buffer_birth_steps, c.power_states, channel_state, c.obs, slot_actions, c.terminals, reward_key, step, obs_features, obs_config, roll_obs=done_now, done_now=done_now, k_tx=k_tx_accum, + k_tx_planned=k_tx_planned, k_coll_ltc=k_coll_ltc, k_coll_coex=k_coll_coex, header_collision=header_collision, @@ -395,12 +397,13 @@ def rl_step_fn(c, step): ack_collision_other=ack_collision_other_now, initial_ret_c=initial_ret_c, tx_started_empty=tx_started_empty, + tx_packet_mask=tx_packet_mask, + tx_success_mask=tx_success_mask, ) - rewards = jnp.where(done_now, rewards_new, c.rewards) - # Reset per-macro accumulators when action completes k_tx_accum = jnp.where(done_now, 0.0, k_tx_accum) + k_tx_planned = jnp.where(done_now, 0.0, k_tx_planned) k_coll_ltc = jnp.where(done_now, 0.0, k_coll_ltc) k_coll_coex = jnp.where(done_now, 0.0, k_coll_coex) macro_remaining = jnp.where(slot_executing, macro_remaining - 1, macro_remaining) @@ -419,6 +422,7 @@ def rl_step_fn(c, step): remaining=macro_remaining, action_types=macro_action_types, tx_success_accum=k_tx_accum, + tx_planned_accum=k_tx_planned, k_coll_ltc=k_coll_ltc, k_coll_coex=k_coll_coex, sub_collision_flag=sub_collision_flag, @@ -455,8 +459,8 @@ def rl_step_fn(c, step): parser = argparse.ArgumentParser(description="Run the RL network simulation with configurable parameters.") parser.add_argument('--n', type=int, default=5, help='Initial number of agents in the simulation.') parser.add_argument('--n_final', type=int, help='Final number of agents in the simulation.') - parser.add_argument('--n_epochs', type=int, default=50, help='Number of training epochs to run.') - parser.add_argument('--n_steps', type=int, default=2000, help='Number of steps per epoch.') + parser.add_argument('--n_epochs', type=int, default=100, help='Number of training epochs to run.') + parser.add_argument('--n_steps', type=int, default=3000, help='Number of steps per epoch.') parser.add_argument('--window_size', type=int, default=1, help='Size of the observation window for each agent.') parser.add_argument('--queue_size', type=int, default=16, help='Fixed packet-queue length for each station.') parser.add_argument('--obs_stats_window', type=int, default=1000, help='Window length for rolling observation statistics.') @@ -590,6 +594,7 @@ def rl_step_fn(c, step): remaining=jnp.zeros(n, dtype=jnp.int32), action_types=jnp.full(n, Actions.IDLE.value, dtype=jnp.int32), tx_success_accum=jnp.zeros(n, dtype=jnp.float32), + tx_planned_accum=jnp.zeros(n, dtype=jnp.float32), k_coll_ltc=jnp.zeros(n, dtype=jnp.float32), k_coll_coex=jnp.zeros(n, dtype=jnp.float32), sub_collision_flag=jnp.zeros(n, dtype=bool), diff --git a/ltc/sim/constants.py b/ltc/sim/constants.py index c0bec8a..732bdb7 100644 --- a/ltc/sim/constants.py +++ b/ltc/sim/constants.py @@ -33,8 +33,8 @@ class Actions(Enum): NO_TX_PENALTY = -1.0 COLLISION_PENALTY = -1.0 -TX_SIZE_THRESHOLD = 5 -TX_SIZE_PENALTY_WINDOW = 5 +TX_SIZE_THRESHOLD = 3 +TX_SIZE_PENALTY_WINDOW = 2 TX_SIZE_PENALTY = -1.0 TX_LTC_COLLISION_PENALTY = -1.0 TX_COEX_COLLISION_PENALTY = -5.0 diff --git a/ltc/sim/process_output.py b/ltc/sim/process_output.py index c769d37..48aa994 100644 --- a/ltc/sim/process_output.py +++ b/ltc/sim/process_output.py @@ -41,40 +41,23 @@ def no_transmission_long(args): return reward, ret_c, no_tx -def transmission(args): - _, _, _, channel_state, _, _ = args - return jax.lax.cond(channel_state == 1, reset_counters, transmission_with_collision, args) - - -def transmission_with_collision(args): - _, _, ret_c, _, _, _ = args - return jax.lax.cond(ret_c < MAX_RETRANSMISSION, retransmission, reset_counters, args) - - def reset_counters(_): ret_c = 0 no_tx = 0 return 0.0, ret_c, no_tx -def retransmission(args): - _, _, ret_c, _, _, _ = args - ret_c = ret_c + 1 - no_tx = 0 - return 0.0, ret_c, no_tx - - def tx_macro_reward( - k_tx, k_coll_ltc, k_coll_coex, tx_started_empty, initial_ret_c, + k_tx, k_tx_planned, k_coll_ltc, k_coll_coex, tx_started_empty, initial_ret_c, header_collision, ack_collision, header_collision_other, ack_collision_other, ): hdr_ack = header_collision | ack_collision k_tx_eff = jnp.where(hdr_ack, 0.0, k_tx) - # When header/ACK collides, attribute all k_tx frames as collisions of the same type + # Header/ACK collision makes all frames unreadable: all planned sub-windows become collisions collision_with_other = (header_collision & header_collision_other) | (ack_collision & ack_collision_other) - k_coll_ltc_eff = k_coll_ltc + jnp.where(hdr_ack & ~collision_with_other, k_tx, 0.0) - k_coll_coex_eff = k_coll_coex + jnp.where(hdr_ack & collision_with_other, k_tx, 0.0) + k_coll_ltc_eff = jnp.where(hdr_ack, jnp.where(~collision_with_other, k_tx_planned, 0.0), k_coll_ltc) + k_coll_coex_eff = jnp.where(hdr_ack, jnp.where(collision_with_other, k_tx_planned, 0.0), k_coll_coex) k_coll = k_coll_ltc_eff + k_coll_coex_eff any_collision = k_coll > 0 @@ -86,24 +69,24 @@ def tx_macro_reward( ) size_penalty = jnp.where( - k_tx_eff >= TX_SIZE_THRESHOLD, - TX_SIZE_PENALTY * jnp.minimum(1.0, (k_tx_eff - TX_SIZE_THRESHOLD + 1.0) / TX_SIZE_PENALTY_WINDOW), + k_tx_planned >= TX_SIZE_THRESHOLD, + TX_SIZE_PENALTY * jnp.minimum(1.0, (k_tx_planned - TX_SIZE_THRESHOLD + 1.0) / TX_SIZE_PENALTY_WINDOW), 0.0, ) + # Collision penalty scales with absolute count; max-retry penalty is additive collision_reward = jnp.where( any_collision, - jnp.where( - initial_ret_c >= MAX_RETRANSMISSION, - TX_MAX_RETRANSMISSION_PENALTY, - TX_LTC_COLLISION_PENALTY * k_coll_ltc_eff + TX_COEX_COLLISION_PENALTY * k_coll_coex_eff, - ), + TX_LTC_COLLISION_PENALTY * k_coll_ltc_eff + + TX_COEX_COLLISION_PENALTY * k_coll_coex_eff + + jnp.where(initial_ret_c >= MAX_RETRANSMISSION, TX_MAX_RETRANSMISSION_PENALTY, 0.0), 0.0, ) + # Empty-buffer penalty scales with planned duration empty_penalty = jnp.where( jnp.logical_and(tx_started_empty, ~any_collision), - TX_EMPTY_BUFFER_PENALTY, + TX_EMPTY_BUFFER_PENALTY * k_tx_planned, 0.0, ) @@ -172,6 +155,7 @@ def process_output_i( roll_obs, done_now, k_tx, + k_tx_planned, k_coll_ltc, k_coll_coex, header_collision, @@ -180,17 +164,42 @@ def process_output_i( ack_collision_other, initial_ret_c, tx_started_empty, + tx_packet_mask, + tx_success_mask, ): - ret_c = obs[-1, ObsIdx.STATUS_RETRY_COUNTER].astype(jnp.int32) - no_tx = obs[-1, ObsIdx.STATUS_NO_TX_COUNTER].astype(jnp.int32) - args = (action, buffer_state, ret_c, channel_state, no_tx, key) + old_ret_c = obs[-1, ObsIdx.STATUS_RETRY_COUNTER].astype(jnp.int32) + no_tx = obs[-1, ObsIdx.STATUS_NO_TX_COUNTER].astype(jnp.int32) + + is_tx = action == Actions.TX.value + + # CS/IDLE: get slot reward, no_tx update, and ret_c (resets if buffer empty) + cs_args = (action, buffer_state, old_ret_c, channel_state, no_tx, key) + cs_reward, cs_ret_c, cs_no_tx = no_transmission(cs_args) + + # TX: no_tx always resets (agent is transmitting); slot reward is 0 (macro reward below) + no_tx = jnp.where(is_tx, jnp.zeros_like(no_tx), cs_no_tx) + slot_reward = jnp.where(is_tx, 0.0, cs_reward) + + # ret_c for TX: update once per sub-window completion (tx_packet_mask), not every slot + tx_sub_success = is_tx & tx_packet_mask & tx_success_mask + tx_sub_coll_retry = is_tx & tx_packet_mask & ~tx_success_mask & (old_ret_c < MAX_RETRANSMISSION) + tx_sub_coll_reset = is_tx & tx_packet_mask & ~tx_success_mask & (old_ret_c >= MAX_RETRANSMISSION) + + ret_c = jnp.where(is_tx, old_ret_c, cs_ret_c) # CS uses cs logic; TX holds by default + ret_c = jnp.where(tx_sub_success, 0, ret_c) + ret_c = jnp.where(tx_sub_coll_retry, old_ret_c + 1, ret_c) + ret_c = jnp.where(tx_sub_coll_reset, 0, ret_c) - # Per-slot call: updates ret_c/no_tx for observation; reward only used for CS/IDLE - slot_reward, ret_c, no_tx = jax.lax.cond(action == Actions.TX.value, transmission, no_transmission, args) + if obs_config.enable_tx_collision_other: + k_coll_ltc_r, k_coll_coex_r = k_coll_ltc, k_coll_coex + hdr_other_r, ack_other_r = header_collision_other, ack_collision_other + else: + k_coll_ltc_r, k_coll_coex_r = k_coll_ltc + k_coll_coex, jnp.zeros_like(k_coll_coex) + hdr_other_r, ack_other_r = jnp.zeros_like(header_collision_other), jnp.zeros_like(ack_collision_other) macro_reward = tx_macro_reward( - k_tx, k_coll_ltc, k_coll_coex, tx_started_empty, initial_ret_c, - header_collision, ack_collision, header_collision_other, ack_collision_other, + k_tx, k_tx_planned, k_coll_ltc_r, k_coll_coex_r, tx_started_empty, initial_ret_c, + header_collision, ack_collision, hdr_other_r, ack_other_r, ) reward = jnp.where( action == Actions.TX.value, @@ -241,6 +250,7 @@ def process_output( roll_obs=True, done_now=False, k_tx=None, + k_tx_planned=None, k_coll_ltc=None, k_coll_coex=None, header_collision=None, @@ -249,11 +259,14 @@ def process_output( ack_collision_other=None, initial_ret_c=None, tx_started_empty=None, + tx_packet_mask=None, + tx_success_mask=None, ): n = actions.shape[0] roll_obs = jnp.broadcast_to(jnp.asarray(roll_obs), actions.shape) done_now = jnp.broadcast_to(jnp.asarray(done_now), actions.shape) k_tx = jnp.broadcast_to(jnp.asarray(k_tx if k_tx is not None else 0.0, dtype=jnp.float32), (n,)) + k_tx_planned = jnp.broadcast_to(jnp.asarray(k_tx_planned if k_tx_planned is not None else 0.0, dtype=jnp.float32), (n,)) k_coll_ltc = jnp.broadcast_to(jnp.asarray(k_coll_ltc if k_coll_ltc is not None else 0.0, dtype=jnp.float32), (n,)) k_coll_coex = jnp.broadcast_to(jnp.asarray(k_coll_coex if k_coll_coex is not None else 0.0, dtype=jnp.float32), (n,)) header_collision = jnp.broadcast_to(jnp.asarray(header_collision if header_collision is not None else False), (n,)) @@ -262,8 +275,10 @@ def process_output( ack_collision_other = jnp.broadcast_to(jnp.asarray(ack_collision_other if ack_collision_other is not None else False), (n,)) initial_ret_c = jnp.broadcast_to(jnp.asarray(initial_ret_c if initial_ret_c is not None else 0, dtype=jnp.int32), (n,)) tx_started_empty = jnp.broadcast_to(jnp.asarray(tx_started_empty if tx_started_empty is not None else False), (n,)) + tx_packet_mask = jnp.broadcast_to(jnp.asarray(tx_packet_mask if tx_packet_mask is not None else False), (n,)) + tx_success_mask = jnp.broadcast_to(jnp.asarray(tx_success_mask if tx_success_mask is not None else False), (n,)) - return jax.vmap(process_output_i, in_axes=(0, 0, 0, 0, None, 0, 0, 0, None, None, 0, None, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0))( + return jax.vmap(process_output_i, in_axes=(0, 0, 0, 0, None, 0, 0, 0, None, None, 0, None, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0))( buffer_states, new_buffer_states, new_buffer_birth_states, @@ -279,6 +294,7 @@ def process_output( roll_obs, done_now, k_tx, + k_tx_planned, k_coll_ltc, k_coll_coex, header_collision, @@ -287,4 +303,6 @@ def process_output( ack_collision_other, initial_ret_c, tx_started_empty, + tx_packet_mask, + tx_success_mask, ) diff --git a/ltc/utils/structs.py b/ltc/utils/structs.py index a285515..71d0059 100644 --- a/ltc/utils/structs.py +++ b/ltc/utils/structs.py @@ -10,6 +10,7 @@ class MacroActionState: remaining: jax.Array action_types: jax.Array tx_success_accum: jax.Array # k_tx: successful subframes + tx_planned_accum: jax.Array # k_tx_planned: planned subframes (= TX macro duration) k_coll_ltc: jax.Array # collided subframes vs LTC k_coll_coex: jax.Array # collided subframes vs coex sub_collision_flag: jax.Array # True if any collision in current TX_SLOTS sub-window diff --git a/test/sim/test_action_space.py b/test/sim/test_action_space.py index 2afacdf..141b2fe 100644 --- a/test/sim/test_action_space.py +++ b/test/sim/test_action_space.py @@ -35,7 +35,8 @@ def test_legacy_decode(self): def _macro_reward(self, **kwargs): defaults = dict( - k_tx=jnp.array(0.0), k_coll_ltc=jnp.array(0.0), k_coll_coex=jnp.array(0.0), + k_tx=jnp.array(0.0), k_tx_planned=jnp.array(0.0), + k_coll_ltc=jnp.array(0.0), k_coll_coex=jnp.array(0.0), tx_started_empty=jnp.array(False), initial_ret_c=jnp.array(0), header_collision=jnp.array(False), ack_collision=jnp.array(False), header_collision_other=jnp.array(False), ack_collision_other=jnp.array(False), @@ -44,40 +45,46 @@ def _macro_reward(self, **kwargs): return tx_macro_reward(**defaults) def test_tx_macro_reward_branches(self): - # success: k_tx=1, no collision - r = self._macro_reward(k_tx=jnp.array(1.0)) + # success: 1 sub-window, no collision + r = self._macro_reward(k_tx=jnp.array(1.0), k_tx_planned=jnp.array(1.0)) expected = TX_REWARD * (1.0 * TX_SLOTS - 2.0) / TX_SLOTS self.assertAlmostEqual(float(r), expected, places=6) - # coex collision, ret_c < MAX - r = self._macro_reward(k_coll_coex=jnp.array(1.0)) + # coex collision, ret_c < MAX: 1 collided sub-window → full coex penalty + r = self._macro_reward(k_coll_coex=jnp.array(1.0), k_tx_planned=jnp.array(1.0)) self.assertAlmostEqual(float(r), TX_COEX_COLLISION_PENALTY * 1.0, places=6) - # LTC collision, max retransmission - r = self._macro_reward(k_coll_ltc=jnp.array(1.0), initial_ret_c=jnp.array(MAX_RETRANSMISSION)) - self.assertAlmostEqual(float(r), TX_MAX_RETRANSMISSION_PENALTY, places=6) + # LTC collision + max retransmission: collision penalty + max-retry penalty (additive) + r = self._macro_reward( + k_coll_ltc=jnp.array(1.0), k_tx_planned=jnp.array(1.0), + initial_ret_c=jnp.array(MAX_RETRANSMISSION), + ) + expected = TX_LTC_COLLISION_PENALTY * 1.0 + TX_MAX_RETRANSMISSION_PENALTY + self.assertAlmostEqual(float(r), expected, places=6) - # empty buffer, no collision - r = self._macro_reward(tx_started_empty=jnp.array(True)) - self.assertAlmostEqual(float(r), TX_EMPTY_BUFFER_PENALTY, places=6) + # empty buffer, no collision: penalty scales with planned duration + r = self._macro_reward(tx_started_empty=jnp.array(True), k_tx_planned=jnp.array(2.0)) + self.assertAlmostEqual(float(r), TX_EMPTY_BUFFER_PENALTY * 2.0, places=6) - # header collision (LTC): k_tx=3 successful frames all become LTC collisions + # header collision (LTC): all 3 planned sub-windows → 3 LTC collisions + # size penalty applies: k_tx_planned=3, ramp=(3-2)/2=0.5 r = self._macro_reward( - k_tx=jnp.array(3.0), + k_tx=jnp.array(3.0), k_tx_planned=jnp.array(3.0), header_collision=jnp.array(True), header_collision_other=jnp.array(False), ) - self.assertAlmostEqual(float(r), TX_LTC_COLLISION_PENALTY * 3.0, places=6) + expected = TX_LTC_COLLISION_PENALTY * 3.0 + TX_SIZE_PENALTY * 0.5 + self.assertAlmostEqual(float(r), expected, places=6) - # header collision (coex): k_tx=2 successful frames become coex collisions + # header collision (coex): all 2 planned sub-windows → 2 coex collisions r = self._macro_reward( - k_tx=jnp.array(2.0), + k_tx=jnp.array(2.0), k_tx_planned=jnp.array(2.0), header_collision=jnp.array(True), header_collision_other=jnp.array(True), ) self.assertAlmostEqual(float(r), TX_COEX_COLLISION_PENALTY * 2.0, places=6) - # ACK collision (LTC) with mixed mid-frame: k_tx=1 successful also becomes LTC collision + # ACK collision (LTC), 2 planned sub-windows: all flagged → 2 LTC collisions r = self._macro_reward( - k_tx=jnp.array(1.0), k_coll_ltc=jnp.array(1.0), + k_tx=jnp.array(1.0), k_coll_ltc=jnp.array(1.0), k_tx_planned=jnp.array(2.0), ack_collision=jnp.array(True), ack_collision_other=jnp.array(False), ) self.assertAlmostEqual(float(r), TX_LTC_COLLISION_PENALTY * 2.0, places=6) diff --git a/test/sim/test_process_output.py b/test/sim/test_process_output.py index 1673d97..bad46bc 100644 --- a/test/sim/test_process_output.py +++ b/test/sim/test_process_output.py @@ -18,35 +18,50 @@ def test_no_transmission(self): self.assertEqual(reward, 0.0) self.assertEqual(r, 5) - def test_transmission_with_collision_retransmission(self): - args = (Actions.TX.value, jnp.array([5, EMPTY, EMPTY]), 2, -1, 1, KEY) - reward, r, no_tx = transmission_with_collision(args) - self.assertEqual(reward, 0.0) + def _make_tx_process_output(self, ret_c_init, channel_state, tx_packet_mask, tx_success_mask): + """Helper: one TX agent, one step, returns obs ret_c.""" + buf = jnp.array([[5, EMPTY, EMPTY]]) + obs = jnp.zeros((1, 3, OBS_SIZE), dtype=jnp.float32) + obs = obs.at[0, -1, ObsIdx.STATUS_RETRY_COUNTER].set(ret_c_init) + obs_features = ObsFeatureInputs( + traffic_mean_arrival_rate=jnp.array([1.0]), + channel_occupancy_pct_window=jnp.array([0.0]), + channel_collisions_pct_window=jnp.array([0.0]), + back_pct=jnp.array([0.0]), + unique_ltc_tx_window=jnp.array([0.0]), + cs_tx_same_type_now=jnp.array([0.0]), + tx_collision_other_now=jnp.array([0.0]), + ) + obs_config = ObsFeatureConfig(enable_cs_tx_same_type=False, enable_tx_collision_other=False) + result_obs, _, _ = process_output( + buffer_states=buf, new_buffer_states=buf, new_buffer_birth_states=jnp.array([[-1,-1,-1]]), + power_states=jnp.array([100]), channel_state=channel_state, + obs=obs, actions=jnp.array([Actions.TX.value]), terminals=jnp.array([False]), + key=KEY, current_step=0, obs_features=obs_features, obs_config=obs_config, + tx_packet_mask=jnp.array([tx_packet_mask]), + tx_success_mask=jnp.array([tx_success_mask]), + ) + return int(result_obs[0, -1, ObsIdx.STATUS_RETRY_COUNTER]) + + def test_tx_sub_window_collision_increments_ret_c(self): + r = self._make_tx_process_output(ret_c_init=2, channel_state=-1, tx_packet_mask=True, tx_success_mask=False) self.assertEqual(r, 3) - def test_transmission_with_collision_max_retransmission(self): - args = (Actions.TX.value, jnp.array([5, EMPTY, EMPTY]), MAX_RETRANSMISSION, -1, 1, KEY) - reward, r, no_tx = transmission_with_collision(args) - self.assertEqual(reward, 0.0) + def test_tx_sub_window_collision_at_max_resets_ret_c(self): + r = self._make_tx_process_output(ret_c_init=MAX_RETRANSMISSION, channel_state=-1, tx_packet_mask=True, tx_success_mask=False) self.assertEqual(r, 0) - def test_transmission_collision_path(self): - args = (Actions.TX.value, jnp.array([5, EMPTY, EMPTY]), 7, -1, 1, KEY) - reward, r, no_tx = transmission(args) - self.assertEqual(reward, 0.0) + def test_tx_sub_window_collision_below_max(self): + r = self._make_tx_process_output(ret_c_init=7, channel_state=-1, tx_packet_mask=True, tx_success_mask=False) self.assertEqual(r, 8) - def test_transmission_successful_path(self): - args = (Actions.TX.value, jnp.array([5, EMPTY, EMPTY]), 1, 1, 1, KEY) - reward, r, no_tx = transmission(args) - self.assertEqual(reward, 0.0) + def test_tx_sub_window_success_resets_ret_c(self): + r = self._make_tx_process_output(ret_c_init=3, channel_state=1, tx_packet_mask=True, tx_success_mask=True) self.assertEqual(r, 0) - def test_transmission_empty_buffer(self): - args = (Actions.TX.value, jnp.array([EMPTY, EMPTY, EMPTY]), 1, 1, 1, KEY) - reward, r, no_tx = transmission(args) - self.assertEqual(reward, 0.0) - self.assertEqual(r, 0) + def test_tx_mid_macro_ret_c_unchanged(self): + r = self._make_tx_process_output(ret_c_init=5, channel_state=-1, tx_packet_mask=False, tx_success_mask=False) + self.assertEqual(r, 5) def test_process_rl_output(self): buffer_states = jnp.array([ @@ -73,7 +88,7 @@ def test_process_rl_output(self): channel_state = -1 obs = jnp.zeros((4, 3, OBS_SIZE), dtype=jnp.float32) obs = obs.at[:, -1, ObsIdx.STATUS_RETRY_COUNTER].set(8) - expected_R_0 = -1.0 + expected_R_0 = TX_COEX_COLLISION_PENALTY * 1.0 + TX_MAX_RETRANSMISSION_PENALTY # -5.0 + -1.0 expected_power = power_states + jnp.array([-TX_CONSUMPTION, -CS_CONSUMPTION, -TX_CONSUMPTION, -TX_CONSUMPTION]) obs_features = ObsFeatureInputs( @@ -90,7 +105,7 @@ def test_process_rl_output(self): enable_tx_collision_other=True, ) - # Agent 0: TX, collision, ret_c=8 >= MAX_RETRANSMISSION, coex collision → TX_MAX_RETRANSMISSION_PENALTY + # Agent 0: TX, collision, ret_c=8 >= MAX_RETRANSMISSION, coex collision → coex penalty + max-retry penalty done_now = jnp.array([True, True, True, True]) k_tx = jnp.zeros(4, dtype=jnp.float32) k_coll_ltc = jnp.zeros(4, dtype=jnp.float32) From ed705ec48ea213e8982f1545494087e557ce577c Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Wed, 6 May 2026 16:50:45 +0200 Subject: [PATCH 18/24] gitignore plots --- .gitignore | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index 8269ccd..290b7ae 100644 --- a/.gitignore +++ b/.gitignore @@ -170,4 +170,5 @@ cython_debug/ # PyPI configuration file .pypirc -scripts/ \ No newline at end of file +scripts/ +plots/ \ No newline at end of file From b2b69b2183215964f5ce749036cebe0aa2456273 Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Wed, 6 May 2026 19:24:12 +0200 Subject: [PATCH 19/24] Fix plot --- ltc/utils/plot_actions.py | 32 +++++++++++++++++++++++++++----- 1 file changed, 27 insertions(+), 5 deletions(-) diff --git a/ltc/utils/plot_actions.py b/ltc/utils/plot_actions.py index 1c6e4a4..91634c0 100644 --- a/ltc/utils/plot_actions.py +++ b/ltc/utils/plot_actions.py @@ -41,13 +41,36 @@ def decode_action_type(raw_action, action_space_variant, max_duration): return raw_action +def resolve_max_duration(metadata, raw_actions): + """ + Resolve max_duration for plotting. + + Older history files may not store max_duration in metadata. In that case, + infer it from observed raw actions (best-effort) instead of falling back to + current code constant, which can misdecode old runs. + """ + md = metadata or {} + md_val = md.get('max_duration') + if md_val is not None: + try: + md_int = int(md_val) + if md_int > 0: + return md_int + except (TypeError, ValueError): + pass + + # txcs_duration_set has action range [0, 2*max_duration-1]. + # Infer max_duration from observed maximum action index. + raw_max = int(np.max(raw_actions)) + return max(1, (raw_max + 2) // 2) + + # ── 1. Actions ──────────────────────────────────────────────────────────────── def plot_action_slots(history, metadata, n, n_drl, seed, epoch=0, max_steps=300): action_space_variant = (metadata or {}).get('action_space_variant', 'txcs_duration_set') - max_duration = (metadata or {}).get('max_duration', MAX_MACRO_DURATION) - raw = np.array(history.actions[epoch, :max_steps]) # (steps, n_agents) + max_duration = resolve_max_duration(metadata, raw) n_steps, n_agents = raw.shape has_slot = hasattr(history, 'slot_actions') and history.slot_actions is not None if has_slot: @@ -159,9 +182,8 @@ def plot_reward_slots(history, n, n_drl, seed, epoch=0, max_steps=300): def plot_combined(history, metadata, n, n_drl, seed, epoch=0, max_steps=300, grid_every=10): action_space_variant = (metadata or {}).get('action_space_variant', 'txcs_duration_set') - max_duration = (metadata or {}).get('max_duration', MAX_MACRO_DURATION) - raw = np.array(history.actions[epoch, :max_steps]) # (steps, n_agents) + max_duration = resolve_max_duration(metadata, raw) has_slot = hasattr(history, 'slot_actions') and history.slot_actions is not None slot_act = np.array(history.slot_actions[epoch, :max_steps]) if has_slot else None ch = np.array(history.channel_state[epoch, :max_steps]) # (steps,) @@ -306,7 +328,7 @@ def inspect(history, metadata, n, n_drl, seed, epoch=0, max_steps=20): slot_raw = np.array(history.slot_actions[epoch, :max_steps]) if has_slot else None action_space_variant = (metadata or {}).get('action_space_variant', 'txcs_duration_set') - max_duration = (metadata or {}).get('max_duration', MAX_MACRO_DURATION) + max_duration = resolve_max_duration(metadata, raw) macro_types = decode_action_type(raw.T, action_space_variant, max_duration).T # (steps, agents) action_names = {Actions.TX.value: 'TX', Actions.CS.value: 'CS', Actions.IDLE.value: 'IDLE'} From 175bf76211aa443779bdeceeb936ec611e6b4af2 Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Wed, 6 May 2026 22:41:00 +0200 Subject: [PATCH 20/24] Fix Cantor pairing --- ltc/agents/dcf.py | 8 ++++++-- ltc/run.py | 13 +++++++------ ltc/sim/process_output.py | 11 ++++++----- ltc/sim/sim.py | 19 +++++++++++++------ 4 files changed, 32 insertions(+), 19 deletions(-) diff --git a/ltc/agents/dcf.py b/ltc/agents/dcf.py index 0bb48d2..96ee116 100644 --- a/ltc/agents/dcf.py +++ b/ltc/agents/dcf.py @@ -48,6 +48,10 @@ def double_cw(): channel_busy = env_state[-1][ObsIdx.CHANNEL_LAST_CS_BUSY] ret_c = env_state[-1][ObsIdx.STATUS_RETRY_COUNTER] + is_tx = action == Actions.TX.value + is_success = jax.lax.bitwise_and(is_tx, jax.lax.bitwise_or(reward > 0, ret_c == 0)) + is_failure = jax.lax.bitwise_and(is_tx, reward < 0) + return jax.lax.cond( buffer_count == 0, reset, @@ -58,10 +62,10 @@ def double_cw(): jax.lax.bitwise_and(state.backoff > 0, channel_busy != 0), freeze, lambda: jax.lax.cond( - jax.lax.bitwise_or(reward > 0, ret_c == 0), + is_success, reset, lambda: jax.lax.cond( - reward < 0, + is_failure, double_cw, freeze ) diff --git a/ltc/run.py b/ltc/run.py index 01304aa..6ef09b2 100644 --- a/ltc/run.py +++ b/ltc/run.py @@ -1,6 +1,7 @@ import os os.environ['XLA_PYTHON_CLIENT_PREALLOCATE'] = 'false' os.environ['XLA_FLAGS'] = '--xla_gpu_enable_triton_gemm=false' +os.environ['JAX_ENABLE_X64'] = 'true' import argparse from functools import partial @@ -331,7 +332,7 @@ def rl_step_fn(c, step): decoded_actions, n_drl ) - slot_executing = macro_remaining > 0 + slot_executing = jnp.logical_and(active, macro_remaining > 0) slot_actions = jnp.where(slot_executing, macro_action_types, Actions.IDLE.value) tx_mask = slot_actions == Actions.TX.value @@ -363,7 +364,7 @@ def rl_step_fn(c, step): planned_packets, successful_packets, ) = simulate(c.buffer_states, c.obs_tracker.buffer_birth_steps, new_frames, slot_actions, c.packet_seqs, step, - tx_packet_mask=tx_packet_mask, tx_success_mask=tx_success_mask) + tx_packet_mask=tx_packet_mask, tx_success_mask=tx_success_mask, active=active) (arrival_hist, planned_tx_hist, success_tx_hist, channel_busy_hist, collision_hist, tx_hist), obs_features = \ compute_obs_features( @@ -579,14 +580,14 @@ def rl_step_fn(c, step): drl_states=drl_states, legacy_states=legacy_states, traffic_states=traffic_states, - packet_seqs=jnp.zeros(n, dtype=jnp.int32), - buffer_states=jnp.full((n, args.queue_size), EMPTY_PACKET_ID, dtype=jnp.int32), + packet_seqs=jnp.zeros(n, dtype=jnp.int64), + buffer_states=jnp.full((n, args.queue_size), EMPTY_PACKET_ID, dtype=jnp.int64), tx_hist=jnp.full((args.obs_stats_window, n), -1, dtype=jnp.int32), - power_states=jnp.full(n, INITIAL_CAPACITY, dtype=int), + power_states=jnp.full(n, INITIAL_CAPACITY, dtype=jnp.int32), channel_state=0, key=key, obs=jnp.zeros((n, args.window_size, OBS_SIZE), dtype=jnp.float32), - actions=jnp.zeros(n, dtype=int), + actions=jnp.zeros(n, dtype=jnp.int32), rewards=jnp.zeros(n), terminals=jnp.full(n, False, dtype=bool), active=jnp.ones(n, dtype=bool).at[n_init:].set(False), diff --git a/ltc/sim/process_output.py b/ltc/sim/process_output.py index 48aa994..5e91d62 100644 --- a/ltc/sim/process_output.py +++ b/ltc/sim/process_output.py @@ -20,7 +20,7 @@ def idle_empty_buffer(_): def idle_full_buffer(args): _, _, _, _, no_tx, key = args rnd_factor = jax.random.normal(key) * SAFE_IDLE_PERIOD_STD - rnd_factor = jnp.round(rnd_factor).astype(int) + rnd_factor = jnp.round(rnd_factor).astype(jnp.int32) safe_period = SAFE_IDLE_PERIOD + rnd_factor args += (safe_period,) return jax.lax.cond(no_tx < safe_period, no_transmission_short, no_transmission_long, args) @@ -167,8 +167,8 @@ def process_output_i( tx_packet_mask, tx_success_mask, ): - old_ret_c = obs[-1, ObsIdx.STATUS_RETRY_COUNTER].astype(jnp.int32) - no_tx = obs[-1, ObsIdx.STATUS_NO_TX_COUNTER].astype(jnp.int32) + old_ret_c = obs[-1, ObsIdx.STATUS_RETRY_COUNTER].astype(jnp.int64) + no_tx = obs[-1, ObsIdx.STATUS_NO_TX_COUNTER].astype(jnp.int64) is_tx = action == Actions.TX.value @@ -263,6 +263,7 @@ def process_output( tx_success_mask=None, ): n = actions.shape[0] + keys = jax.random.split(key, n) roll_obs = jnp.broadcast_to(jnp.asarray(roll_obs), actions.shape) done_now = jnp.broadcast_to(jnp.asarray(done_now), actions.shape) k_tx = jnp.broadcast_to(jnp.asarray(k_tx if k_tx is not None else 0.0, dtype=jnp.float32), (n,)) @@ -278,7 +279,7 @@ def process_output( tx_packet_mask = jnp.broadcast_to(jnp.asarray(tx_packet_mask if tx_packet_mask is not None else False), (n,)) tx_success_mask = jnp.broadcast_to(jnp.asarray(tx_success_mask if tx_success_mask is not None else False), (n,)) - return jax.vmap(process_output_i, in_axes=(0, 0, 0, 0, None, 0, 0, 0, None, None, 0, None, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0))( + return jax.vmap(process_output_i, in_axes=(0, 0, 0, 0, None, 0, 0, 0, 0, None, 0, None, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0))( buffer_states, new_buffer_states, new_buffer_birth_states, @@ -287,7 +288,7 @@ def process_output( obs, actions, terminals, - key, + keys, current_step, obs_features, obs_config, diff --git a/ltc/sim/sim.py b/ltc/sim/sim.py index e276bd0..47d044f 100644 --- a/ltc/sim/sim.py +++ b/ltc/sim/sim.py @@ -59,8 +59,8 @@ def enqueue_generated_packets(buffer_state, buffer_birth_state, new_frames, stat new_rel = idx - kept_old take_new = jnp.logical_and(idx >= kept_old, new_rel < kept_new) - new_local_ids = new_start + jnp.maximum(new_rel, 0) - new_packets = cantor_pairing(station_id.astype(jnp.int32), new_local_ids) + new_local_ids = new_start + jnp.maximum(new_rel, 0).astype(jnp.int64) + new_packets = cantor_pairing(station_id.astype(jnp.int64), new_local_ids) new_births = jnp.full((q,), current_step, dtype=jnp.int32) next_state = jnp.where(take_old, old_packets, jnp.where(take_new, new_packets, EMPTY_PACKET_ID)) @@ -70,18 +70,25 @@ def enqueue_generated_packets(buffer_state, buffer_birth_state, new_frames, stat def add_new_frames(buffer_states, buffer_birth_steps, new_frames, packet_seqs, current_step): - station_ids = jnp.arange(buffer_states.shape[0], dtype=jnp.int32) + station_ids = jnp.arange(buffer_states.shape[0], dtype=jnp.int64) return jax.vmap(enqueue_generated_packets, in_axes=(0, 0, 0, 0, 0, None))( buffer_states, buffer_birth_steps, new_frames, station_ids, packet_seqs, current_step ) def simulate(buffer_states, buffer_birth_steps, new_frames, actions, packet_seqs, current_step, - tx_packet_mask, tx_success_mask): + tx_packet_mask, tx_success_mask, active=None): + if active is None: + active = jnp.ones_like(actions, dtype=bool) + + active = jnp.asarray(active, dtype=bool) + actions = jnp.where(active, actions, Actions.IDLE.value) + new_frames = jnp.where(active, new_frames, 0) + channel_state = channel_state_selector(actions) pre_nonempty = jnp.any(buffer_states != EMPTY_PACKET_ID, axis=1) - packet_tx = tx_success_mask & pre_nonempty - planned_packets = tx_packet_mask.astype(jnp.int32) + packet_tx = tx_success_mask & pre_nonempty & active + planned_packets = (tx_packet_mask & active).astype(jnp.int32) successful_packets = packet_tx.astype(jnp.int32) queue_size = buffer_states.shape[1] head_mask = jnp.arange(queue_size) == 0 From eaff4e37bd65fd2e7984625ddee21acc7aea7fbf Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Wed, 6 May 2026 23:20:00 +0200 Subject: [PATCH 21/24] Normalize observations and refactor --- ltc/run.py | 169 ++++++++++++++++++++------------------ ltc/sim/__init__.py | 2 +- ltc/sim/constants.py | 10 ++- ltc/sim/process_output.py | 9 ++ 4 files changed, 109 insertions(+), 81 deletions(-) diff --git a/ltc/run.py b/ltc/run.py index 6ef09b2..1b05540 100644 --- a/ltc/run.py +++ b/ltc/run.py @@ -4,6 +4,7 @@ os.environ['JAX_ENABLE_X64'] = 'true' import argparse +from dataclasses import dataclass from functools import partial import cloudpickle @@ -14,13 +15,35 @@ from tqdm import trange from ltc.agents import BayesianDDQN, DCF, QNetwork, StochasticVariationalNetwork -from ltc.sim import InitialStateConf, cox_traffic, process_output, simulate, channel_state_selector +from ltc.sim import InitialStateConf, cox_traffic, normalize_obs, process_output, simulate, channel_state_selector from ltc.sim.constants import * from ltc.utils.history import build_history_filename, ensure_clean_git_worktree, get_short_commit_hash from ltc.utils.structs import * from ltc.utils.plots import plot_all, plot_first +@dataclass(frozen=True) +class StationScheduleConfig: + one_shot_step: int | None = None + one_shot_target: int | None = None + station_change_interval: int | None = None + station_change_delta: int = 0 + station_change_start_step: int = 0 + station_change_stop_step: int | None = None + station_change_target: int | None = None + + +@dataclass(frozen=True) +class RLStepConfig: + n_bins: int = 50 + schedule: StationScheduleConfig = StationScheduleConfig() + obs_enable_cs_tx_same_type: bool = False + obs_enable_tx_collision_other: bool = False + action_space_variant: str = 'txcs_duration_set' + max_duration: int = 5 + legacy_tx_duration: int = 5 + + def init_agents(agent, key, n): keys = jax.random.split(key, n) states = jax.vmap(agent.init)(keys) @@ -52,40 +75,30 @@ def init_traffic(traffic, key, n): return states, step_fn -def schedule_active_stations( - active, - step, - *, - one_shot_step=None, - one_shot_target=None, - station_change_interval=None, - station_change_delta=0, - station_change_start_step=0, - station_change_stop_step=None, - station_change_target=None, -): +def schedule_active_stations(active, step, schedule): n = active.shape[0] next_active = active - if one_shot_step is not None and one_shot_target is not None: - target = jnp.asarray(one_shot_target, dtype=jnp.int32) + if schedule.one_shot_step is not None and schedule.one_shot_target is not None: + target = jnp.asarray(schedule.one_shot_target, dtype=jnp.int32) one_shot_active = jnp.arange(n) < target - next_active = jnp.where(step == one_shot_step, one_shot_active, next_active) + next_active = jnp.where(step == schedule.one_shot_step, one_shot_active, next_active) - if station_change_interval is not None and station_change_delta != 0: - should_change = step >= station_change_start_step - if station_change_stop_step is not None: - should_change = jnp.logical_and(should_change, step <= station_change_stop_step) + if schedule.station_change_interval is not None and schedule.station_change_delta != 0: + should_change = step >= schedule.station_change_start_step + if schedule.station_change_stop_step is not None: + should_change = jnp.logical_and(should_change, step <= schedule.station_change_stop_step) should_change = jnp.logical_and( - should_change, (step - station_change_start_step) % station_change_interval == 0 + should_change, + (step - schedule.station_change_start_step) % schedule.station_change_interval == 0, ) current_count = jnp.sum(next_active.astype(jnp.int32)) - updated_count = jnp.clip(current_count + station_change_delta, 0, n) - if station_change_target is not None: - target = jnp.asarray(station_change_target, dtype=jnp.int32) + updated_count = jnp.clip(current_count + schedule.station_change_delta, 0, n) + if schedule.station_change_target is not None: + target = jnp.asarray(schedule.station_change_target, dtype=jnp.int32) updated_count = jnp.where( - station_change_delta > 0, + schedule.station_change_delta > 0, jnp.minimum(updated_count, target), jnp.maximum(updated_count, target), ) @@ -197,6 +210,14 @@ def update_packet_tracking( ) +def scale_tx_durations(action_types, durations): + return jnp.where( + action_types == Actions.TX.value, + durations * TX_SLOTS, + durations, + ) + + def compute_obs_features( obs_tracker, tx_hist, n_drl, new_frames, planned_packets, successful_packets, channel_state, tx_mask, @@ -247,11 +268,7 @@ def update_macro_state_from_ready_agents( macro_action_types = macro_action_types.at[:n_drl].set( jnp.where(ready_drl, decoded_actions.drl_action_types, macro_state.action_types[:n_drl]) ) - drl_scaled_durations = jnp.where( - decoded_actions.drl_action_types == Actions.TX.value, - decoded_actions.drl_durations * TX_SLOTS, - decoded_actions.drl_durations, - ) + drl_scaled_durations = scale_tx_durations(decoded_actions.drl_action_types, decoded_actions.drl_durations) macro_remaining = macro_remaining.at[:n_drl].set( jnp.where(ready_drl, drl_scaled_durations, macro_state.remaining[:n_drl]) ) @@ -262,10 +279,8 @@ def update_macro_state_from_ready_agents( macro_action_types = macro_action_types.at[n_drl:].set( jnp.where(ready_legacy, decoded_actions.legacy_action_types, macro_state.action_types[n_drl:]) ) - legacy_scaled_durations = jnp.where( - decoded_actions.legacy_action_types == Actions.TX.value, - decoded_actions.legacy_durations * TX_SLOTS, - decoded_actions.legacy_durations, + legacy_scaled_durations = scale_tx_durations( + decoded_actions.legacy_action_types, decoded_actions.legacy_durations ) macro_remaining = macro_remaining.at[n_drl:].set( jnp.where(ready_legacy, legacy_scaled_durations, macro_state.remaining[n_drl:]) @@ -274,25 +289,25 @@ def update_macro_state_from_ready_agents( return macro_action_types, macro_remaining, staged_tx -def rl_step( - drl_step, legacy_step, traffic_step, n, n_drl, n_bins=50, - one_shot_step=None, one_shot_target=None, station_change_interval=None, station_change_delta=0, - station_change_start_step=0, station_change_stop_step=None, station_change_target=None, - obs_enable_cs_tx_same_type=False, obs_enable_tx_collision_other=False, - action_space_variant='txcs_duration_set', max_duration=5, legacy_tx_duration=5, -): +def build_traffic(traffic_type, f3dB, loc, scale): + if traffic_type == 'custom': + params = {'f3dB': f3dB, 'loc': loc, 'scale': scale} + else: + presets = { + 'constant': {'f3dB': 1.0, 'loc': -1.0, 'scale': 0.0}, + 'saturated': {'f3dB': 1.0, 'loc': 5.0, 'scale': 0.0}, + 'bursty': {'f3dB': 0.1, 'loc': -5.0, 'scale': 5.0}, + } + params = presets.get(traffic_type) + if params is None: + raise ValueError(f'Unknown traffic type: {traffic_type}') + + return cox_traffic(initial_state=InitialStateConf.ZERO, **params) + + +def rl_step(drl_step, legacy_step, traffic_step, n, n_drl, config=RLStepConfig()): def rl_step_fn(c, step): - active = schedule_active_stations( - c.active, - step, - one_shot_step=one_shot_step, - one_shot_target=one_shot_target, - station_change_interval=station_change_interval, - station_change_delta=station_change_delta, - station_change_start_step=station_change_start_step, - station_change_stop_step=station_change_stop_step, - station_change_target=station_change_target, - ) + active = schedule_active_stations(c.active, step, config.schedule) key, drl_keys, legacy_keys, traffic_key, reward_key = jax.random.split(c.key, 5) drl_keys = jax.random.split(drl_keys, n_drl) @@ -303,11 +318,12 @@ def rl_step_fn(c, step): drl_ready = jnp.logical_and(active[:n_drl], ready[:n_drl]) legacy_ready = jnp.logical_and(active[n_drl:], ready[n_drl:]) + obs_norm = normalize_obs(c.obs, queue_size=c.buffer_states.shape[1]) drl_states, drl_actions_raw = drl_step( - c.drl_states, drl_keys, c.obs[:n_drl], c.actions[:n_drl], c.rewards[:n_drl], c.terminals[:n_drl], drl_ready + c.drl_states, drl_keys, obs_norm[:n_drl], c.actions[:n_drl], c.rewards[:n_drl], c.terminals[:n_drl], drl_ready ) legacy_states, legacy_actions_raw = legacy_step( - c.legacy_states, legacy_keys, c.obs[n_drl:], c.actions[n_drl:], c.rewards[n_drl:], c.terminals[n_drl:], legacy_ready + c.legacy_states, legacy_keys, obs_norm[n_drl:], c.actions[n_drl:], c.rewards[n_drl:], c.terminals[n_drl:], legacy_ready ) raw_actions = c.actions @@ -315,9 +331,11 @@ def rl_step_fn(c, step): raw_actions = raw_actions.at[n_drl:].set(jnp.where(legacy_ready, legacy_actions_raw, c.actions[n_drl:])) drl_action_types, drl_durations, drl_staged_next = decode_drl_actions( - raw_actions[:n_drl], c.obs_tracker.staged_tx[:n_drl], action_space_variant, max_duration + raw_actions[:n_drl], c.obs_tracker.staged_tx[:n_drl], config.action_space_variant, config.max_duration + ) + legacy_action_types, legacy_durations = decode_legacy_actions( + raw_actions[n_drl:], config.legacy_tx_duration ) - legacy_action_types, legacy_durations = decode_legacy_actions(raw_actions[n_drl:], legacy_tx_duration) decoded_actions = ActionDecodingResults( drl_action_types=drl_action_types, @@ -380,8 +398,8 @@ def rl_step_fn(c, step): k_coll_coex = c.macro.k_coll_coex + jnp.where(packet_collided & sub_collision_other_flag, 1.0, 0.0) obs_config = ObsFeatureConfig( - enable_cs_tx_same_type=obs_enable_cs_tx_same_type, - enable_tx_collision_other=obs_enable_tx_collision_other, + enable_cs_tx_same_type=config.obs_enable_cs_tx_same_type, + enable_tx_collision_other=config.obs_enable_tx_collision_other, ) obs, rewards, powers = process_output( c.buffer_states, buffer_states, buffer_birth_steps, c.power_states, channel_state, c.obs, slot_actions, @@ -415,7 +433,7 @@ def rl_step_fn(c, step): flat_params, _ = jax.tree.flatten(params) flat_params = jax.tree.map(lambda x: x.reshape(n_drl, -1), flat_params) flat_params = jnp.hstack(flat_params) - hist, bin_edges = jax.vmap(jnp.histogram, in_axes=(0, None))(flat_params, n_bins) + hist, bin_edges = jax.vmap(jnp.histogram, in_axes=(0, None))(flat_params, config.n_bins) else: hist, bin_edges = None, None @@ -548,33 +566,26 @@ def rl_step_fn(c, step): key, init_key = jax.random.split(key) - if args.traffic_type == 'constant': - traffic = cox_traffic(f3dB=1.0, loc=-1.0, scale=0.0, initial_state=InitialStateConf.ZERO) - elif args.traffic_type == 'saturated': - traffic = cox_traffic(f3dB=1.0, loc=5.0, scale=0.0, initial_state=InitialStateConf.ZERO) - elif args.traffic_type == 'bursty': - traffic = cox_traffic(f3dB=0.1, loc=-5.0, scale=5.0, initial_state=InitialStateConf.ZERO) - elif args.traffic_type == 'custom': - traffic = cox_traffic(f3dB=args.f3dB, loc=args.loc, scale=args.scale, initial_state=InitialStateConf.ZERO) - else: - raise ValueError(f'Unknown traffic type: {args.traffic_type}') + traffic = build_traffic(args.traffic_type, args.f3dB, args.loc, args.scale) traffic_states, traffic_step = init_traffic(traffic, init_key, n) - rl_step_fn = jax.jit(rl_step( - drl_step, legacy_step, traffic_step, n, n, - one_shot_step=one_shot_step, - one_shot_target=one_shot_target, - station_change_interval=args.station_change_interval, - station_change_delta=args.station_change_delta, - station_change_start_step=station_change_start_step, - station_change_stop_step=args.station_change_stop_step, - station_change_target=station_change_target, + step_config = RLStepConfig( + schedule=StationScheduleConfig( + one_shot_step=one_shot_step, + one_shot_target=one_shot_target, + station_change_interval=args.station_change_interval, + station_change_delta=args.station_change_delta, + station_change_start_step=station_change_start_step, + station_change_stop_step=args.station_change_stop_step, + station_change_target=station_change_target, + ), obs_enable_cs_tx_same_type=args.obs_enable_cs_tx_same_type, obs_enable_tx_collision_other=args.obs_enable_tx_collision_other, action_space_variant=args.action_space_variant, max_duration=MAX_MACRO_DURATION, legacy_tx_duration=LEGACY_TX_DURATION, - )) + ) + rl_step_fn = jax.jit(rl_step(drl_step, legacy_step, traffic_step, n, n, config=step_config)) carry = Carry( drl_states=drl_states, diff --git a/ltc/sim/__init__.py b/ltc/sim/__init__.py index 1b9ba0d..f428fa0 100644 --- a/ltc/sim/__init__.py +++ b/ltc/sim/__init__.py @@ -1,3 +1,3 @@ from ltc.sim.traffic import InitialStateConf, ModelState, TrafficModel, cox_traffic -from ltc.sim.process_output import process_output +from ltc.sim.process_output import process_output, normalize_obs from ltc.sim.sim import simulate, channel_state_selector diff --git a/ltc/sim/constants.py b/ltc/sim/constants.py index 732bdb7..2eeeb95 100644 --- a/ltc/sim/constants.py +++ b/ltc/sim/constants.py @@ -55,7 +55,9 @@ class Actions(Enum): """ TAU = 5.484 * 1e-3 -# Queue representation +""" +Queue representation +""" EMPTY_PACKET_ID = -1 """ @@ -64,6 +66,12 @@ class Actions(Enum): MAX_MACRO_DURATION = 5 LEGACY_TX_DURATION = 5 +""" +Observation normalization constants +""" +OBS_UNIQUE_LTC_TX_NORM = 10 +OBS_TRAFFIC_ARRIVAL_NORM = 10 + class ObsIdx: BUFFER_OCCUPANCY_PCT = 0 diff --git a/ltc/sim/process_output.py b/ltc/sim/process_output.py index 5e91d62..0da22d0 100644 --- a/ltc/sim/process_output.py +++ b/ltc/sim/process_output.py @@ -139,6 +139,15 @@ def build_observation_entry( ], dtype=jnp.float32) +def normalize_obs(obs, queue_size): + obs = obs.at[..., ObsIdx.BUFFER_PACKET_COUNT].divide(queue_size) + obs = obs.at[..., ObsIdx.TRAFFIC_MEAN_ARRIVAL_RATE].divide(OBS_TRAFFIC_ARRIVAL_NORM) + obs = obs.at[..., ObsIdx.STATUS_RETRY_COUNTER].divide(MAX_RETRANSMISSION) + obs = obs.at[..., ObsIdx.STATUS_NO_TX_COUNTER].divide(SAFE_IDLE_PERIOD + PENALIZED_IDLE_PERIOD) + obs = obs.at[..., ObsIdx.STATUS_UNIQUE_LTC_TX_WINDOW].divide(OBS_UNIQUE_LTC_TX_NORM) + return obs + + def process_output_i( buffer_state, new_buffer_state, From 0076b973ecbabdc0099921b280240e3a55f9e64f Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Wed, 6 May 2026 23:24:35 +0200 Subject: [PATCH 22/24] Fix tests --- test/sim/test_process_output.py | 2 +- test/sim/test_traffic.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/test/sim/test_process_output.py b/test/sim/test_process_output.py index bad46bc..f77015a 100644 --- a/test/sim/test_process_output.py +++ b/test/sim/test_process_output.py @@ -13,7 +13,7 @@ class ProcessOutputTestCase(unittest.TestCase): def test_no_transmission(self): - args = (Actions.IDLE.value, jnp.array([1, EMPTY, EMPTY]), 5, 0, 1, KEY) + args = (Actions.IDLE.value, jnp.array([1, EMPTY, EMPTY]), jnp.array(5, dtype=jnp.int64), 0, jnp.array(1, dtype=jnp.int64), KEY) reward, r, no_tx = no_transmission(args) self.assertEqual(reward, 0.0) self.assertEqual(r, 5) diff --git a/test/sim/test_traffic.py b/test/sim/test_traffic.py index e0e68b5..13f71f9 100644 --- a/test/sim/test_traffic.py +++ b/test/sim/test_traffic.py @@ -61,7 +61,7 @@ def test_stats(self): estimated_acf=np.corrcoef(yjax[:-1]-yjax.mean(), yjax[1:]-yjax.mean())[1,0] - self.assertAlmostEqual(yjax.mean(), stats.mean, places=3) + self.assertAlmostEqual(yjax.mean(), stats.mean, places=2) self.assertAlmostEqual(yjax.var(), stats.variance, places=1) self.assertAlmostEqual(estimated_acf, stats.acf_lag1, places=2) ... From 316930dd26e660e85366e34bc932790f17a88cff Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Fri, 15 May 2026 18:01:44 +0200 Subject: [PATCH 23/24] Fix code for tx length equal to 1 --- ltc/run.py | 2 +- ltc/sim/constants.py | 14 ++++++++------ ltc/sim/process_output.py | 18 ++++++++++++++---- 3 files changed, 23 insertions(+), 11 deletions(-) diff --git a/ltc/run.py b/ltc/run.py index 1b05540..4a114ed 100644 --- a/ltc/run.py +++ b/ltc/run.py @@ -183,7 +183,7 @@ def update_packet_tracking( tx_executing_pre & (channel_state_prelim == -1) & (tx_collision_other_now > 0.5), True, sub_collision_other_flag, ) - tx_packet_mask = (macro_remaining % TX_SLOTS == 1) & slot_executing & tx_executing_pre + tx_packet_mask = ((macro_remaining - 1) % TX_SLOTS == 0) & slot_executing & tx_executing_pre tx_success_mask = tx_packet_mask & ~sub_collision_flag collision_with_other = tx_collision_other_now > 0.5 diff --git a/ltc/sim/constants.py b/ltc/sim/constants.py index 2eeeb95..a4affa3 100644 --- a/ltc/sim/constants.py +++ b/ltc/sim/constants.py @@ -83,11 +83,13 @@ class ObsIdx: CHANNEL_LAST_TX_COLLISION_OTHER = 6 CHANNEL_OCCUPANCY_PCT_WINDOW = 7 CHANNEL_COLLISIONS_PCT_WINDOW = 8 - ACTION_LAST = 9 - ACTION_BACK_PCT = 10 - STATUS_RETRY_COUNTER = 11 - STATUS_NO_TX_COUNTER = 12 - STATUS_UNIQUE_LTC_TX_WINDOW = 13 + ACTION_TX = 9 + ACTION_CS = 10 + ACTION_IDLE = 11 + ACTION_BACK_PCT = 12 + STATUS_RETRY_COUNTER = 13 + STATUS_NO_TX_COUNTER = 14 + STATUS_UNIQUE_LTC_TX_WINDOW = 15 -OBS_SIZE = 14 +OBS_SIZE = 16 diff --git a/ltc/sim/process_output.py b/ltc/sim/process_output.py index 0da22d0..360e532 100644 --- a/ltc/sim/process_output.py +++ b/ltc/sim/process_output.py @@ -62,9 +62,10 @@ def tx_macro_reward( k_coll = k_coll_ltc_eff + k_coll_coex_eff any_collision = k_coll > 0 + net_data_slots = jnp.maximum(k_tx_eff * TX_SLOTS - 2.0, k_tx_eff) success_reward = jnp.where( k_tx_eff > 0, - TX_REWARD * (k_tx_eff * TX_SLOTS - 2.0) / TX_SLOTS, + TX_REWARD * net_data_slots / TX_SLOTS, 0.0, ) @@ -97,8 +98,7 @@ def oldest_packet_age_norm(buffer_state, buffer_birth_state, current_step): valid = buffer_state != EMPTY_PACKET_ID oldest_birth = jnp.min(jnp.where(valid, buffer_birth_state, current_step)) age = jnp.where(jnp.any(valid), current_step - oldest_birth, -1) - queue_size = jnp.maximum(buffer_state.shape[0], 1) - return jnp.where(age >= 0, age.astype(jnp.float32) / queue_size, -1.0) + return jnp.where(age >= 0, age.astype(jnp.float32), -1.0) def build_observation_entry( @@ -121,6 +121,10 @@ def build_observation_entry( -1.0, ) + action_tx = jnp.where(action == Actions.TX.value, 1.0, 0.0) + action_cs = jnp.where(action == Actions.CS.value, 1.0, 0.0) + action_idle = jnp.where(action == Actions.IDLE.value, 1.0, 0.0) + return jnp.array([ buffer_occupancy_pct, buffer_count.astype(jnp.float32), @@ -131,7 +135,9 @@ def build_observation_entry( tx_collision_other.astype(jnp.float32), obs_features.channel_occupancy_pct_window, obs_features.channel_collisions_pct_window, - action.astype(jnp.float32), + action_tx, + action_cs, + action_idle, obs_features.back_pct, ret_c.astype(jnp.float32), no_tx.astype(jnp.float32), @@ -141,6 +147,10 @@ def build_observation_entry( def normalize_obs(obs, queue_size): obs = obs.at[..., ObsIdx.BUFFER_PACKET_COUNT].divide(queue_size) + raw_age = obs[..., ObsIdx.BUFFER_OLDEST_AGE_NORM] + obs = obs.at[..., ObsIdx.BUFFER_OLDEST_AGE_NORM].set( + jnp.where(raw_age >= 0, jnp.minimum(raw_age / SAFE_IDLE_PERIOD, 1.0), raw_age) + ) obs = obs.at[..., ObsIdx.TRAFFIC_MEAN_ARRIVAL_RATE].divide(OBS_TRAFFIC_ARRIVAL_NORM) obs = obs.at[..., ObsIdx.STATUS_RETRY_COUNTER].divide(MAX_RETRANSMISSION) obs = obs.at[..., ObsIdx.STATUS_NO_TX_COUNTER].divide(SAFE_IDLE_PERIOD + PENALIZED_IDLE_PERIOD) From 3df66b3d876b40adc9a2cd0271363e7b955459e0 Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Fri, 15 May 2026 18:02:00 +0200 Subject: [PATCH 24/24] Fix fairness --- ltc/run.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ltc/run.py b/ltc/run.py index 4a114ed..da15dd2 100644 --- a/ltc/run.py +++ b/ltc/run.py @@ -553,7 +553,7 @@ def rl_step_fn(c, step): discount=0.95, epsilon=1.0, epsilon_decay=0.999, - epsilon_min=0.0, + epsilon_min=0.05, tau=0.05 ) key, init_key = jax.random.split(key)