From ef938964d475180b20ba066814d2da62a1c83cf0 Mon Sep 17 00:00:00 2001 From: kaszczech Date: Tue, 30 Sep 2025 08:55:14 +0200 Subject: [PATCH 01/17] Implementing an agent from the literature --- ltc/sim/constants.py | 9 +---- ltc/sim/process_output.py | 79 +++++---------------------------------- 2 files changed, 10 insertions(+), 78 deletions(-) diff --git a/ltc/sim/constants.py b/ltc/sim/constants.py index 864db46..d588386 100644 --- a/ltc/sim/constants.py +++ b/ltc/sim/constants.py @@ -6,8 +6,7 @@ class Actions(Enum): Actions for the agent """ TX = 0 - CS = 1 - IDLE = 2 + IDLE = 1 """ @@ -21,19 +20,13 @@ class Actions(Enum): Rewards and penalties """ TX_REWARD = 1.0 -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 """ Power consumption """ INITIAL_CAPACITY = 10000 TX_CONSUMPTION = 5 -CS_CONSUMPTION = 1 IDLE_CONSUMPTION = 0 """ diff --git a/ltc/sim/process_output.py b/ltc/sim/process_output.py index 9b4d8da..f3c4190 100644 --- a/ltc/sim/process_output.py +++ b/ltc/sim/process_output.py @@ -5,83 +5,27 @@ def no_transmission(args): - action, buffer_state, _, _, no_tx = args - return jax.lax.cond( - action == Actions.IDLE.value, - lambda: jax.lax.cond(buffer_state == 0, idle_empty_buffer, idle_full_buffer, args), - lambda: jax.lax.cond(no_tx < SAFE_IDLE_PERIOD, no_transmission_short, no_transmission_long, args), - ) - - -def idle_empty_buffer(_): - reward = EMPTY_BUFFER_REWARD + _, _, ret_c, _, _ = args + reward = NO_TX_REWARD ret_c = 0 no_tx = 0 return reward, ret_c, no_tx - -def idle_full_buffer(args): - _, _, ret_c, _, no_tx = args - reward = NO_TX_REWARD - no_tx = no_tx + 1 - return reward, ret_c, no_tx - - -def no_transmission_short(args): - _, buffer_state, ret_c, _, no_tx = args - reward = NO_TX_REWARD - no_tx = jnp.where(buffer_state == 0, 0, no_tx + 1) - return reward, ret_c, no_tx - - -def no_transmission_long(args): - _, _, ret_c, _, no_tx = args - scale = jax.lax.min(1., (no_tx - SAFE_IDLE_PERIOD + 1) / PENALIZED_IDLE_PERIOD) - reward = scale * NO_TX_PENALTY - no_tx = no_tx + 1 - return reward, ret_c, no_tx - - 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, successful_transmission, 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) - - -def max_retransmission_collision(_): - reward = MAX_RETRANSMISSION_PENALTY - ret_c = 0 - no_tx = 0 - return reward, 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 - - -def transmission_without_collision(args): - _, buffer_state, _, _, _ = args - return jax.lax.cond(buffer_state > 0, successful_transmission, empty_buffer_transmission, args) - - -def empty_buffer_transmission(_): - reward = EMPTY_TX_PENALTY + reward = NO_TX_REWARD ret_c = 0 no_tx = 0 return reward, ret_c, no_tx - def successful_transmission(args): _, _, ret_c, _, _ = args - reward = TX_REWARD / (ret_c + 1) + reward = TX_REWARD ret_c = 0 no_tx = 0 return reward, ret_c, no_tx @@ -94,19 +38,14 @@ def process_output_i(buffer_state, new_buffer_state, power_state, channel_state, reward, ret_c, no_tx = jax.lax.cond(action == Actions.TX.value, transmission, no_transmission, args) reward = jnp.where(terminal, 0., reward) - channel_state = jnp.where(action == Actions.CS.value, channel_state, -1) + # channel_state = jnp.where(action == Actions.CS.value, channel_state, -1) + power = jnp.where( action == Actions.TX.value, power_state - TX_CONSUMPTION, - jnp.where( - action == Actions.CS.value, power_state - CS_CONSUMPTION, - jnp.where( - action == Actions.IDLE.value, power_state - IDLE_CONSUMPTION, - power_state - ) - ) + power_state ) - obs_t = jnp.array([new_buffer_state, channel_state, ret_c, no_tx, power]) + obs_t = jnp.array([channel_state, action]) obs = jnp.roll(obs, -1, axis=0) obs = obs.at[-1].set(obs_t) From 5f2c8e94759b05e1732cd6817ee84cb4719700e0 Mon Sep 17 00:00:00 2001 From: kaszczech Date: Tue, 30 Sep 2025 15:32:23 +0200 Subject: [PATCH 02/17] Implementing an agent from the literature --- ltc/run.py | 15 +++++++++------ ltc/sim/constants.py | 9 ++++++++- ltc/sim/process_output.py | 3 ++- 3 files changed, 19 insertions(+), 8 deletions(-) diff --git a/ltc/run.py b/ltc/run.py index e3ec9a8..23e3538 100644 --- a/ltc/run.py +++ b/ltc/run.py @@ -57,8 +57,9 @@ def rl_step_fn(c, _): traffic_keys = jax.random.split(traffic_key, n) drl_states, drl_actions = drl_step(c.drl_states, drl_keys, c.obs[:n_drl], c.actions[:n_drl], c.rewards[:n_drl], c.terminals[:n_drl]) - dcf_states, dcf_actions = dcf_step(c.dcf_states, dcf_keys, c.obs[n_drl:], c.actions[n_drl:], c.rewards[n_drl:], c.terminals[n_drl:]) - actions = jnp.concatenate([drl_actions, dcf_actions]) + # dcf_states, dcf_actions = dcf_step(c.dcf_states, dcf_keys, c.obs[n_drl:], c.actions[n_drl:], c.rewards[n_drl:], c.terminals[n_drl:]) + # actions = jnp.concatenate([drl_actions, dcf_actions]) + actions = drl_actions traffic_states, new_frames = traffic_step(c.traffic_states, traffic_keys) buffer_states, channel_state = simulate(c.buffer_states, new_frames, actions) @@ -85,7 +86,7 @@ def rl_step_fn(c, _): if __name__ == '__main__': - n, n_drl = 10, 5 + n, n_drl = 10, 10 n_epochs, n_steps = 50, 10000 window_size = 5 seed = 42 @@ -95,7 +96,7 @@ def rl_step_fn(c, _): buffer_states = jnp.zeros(n, dtype=int) power_states = jnp.full(n, INITIAL_CAPACITY, dtype=int) channel_state = 0 - obs = jnp.zeros((n, window_size, 5), dtype=int).at[:, -1].set(INITIAL_CAPACITY) + obs = jnp.zeros((n, window_size, 2), dtype=int).at[:, -1].set(INITIAL_CAPACITY) rewards = jnp.zeros(n) terminals = jnp.full(n, False, dtype=bool) @@ -117,9 +118,11 @@ def rl_step_fn(c, _): drl_states, drl_step = init_agents(drl, init_key, n_drl) drl_states = drl_states.replace(prev_env_state=drl_states.prev_env_state.astype(int)) - dcf = DCF() key, init_key = jax.random.split(key) - dcf_states, dcf_step = init_agents(dcf, init_key, n - n_drl) + # dcf = DCF() + # dcf_states, dcf_step = init_agents(dcf, init_key, n - n_drl) + dcf_states = None + dcf_step = None key, init_key = jax.random.split(key) traffic = cox_traffic(f3dB=0.1, loc=-5.0, scale=5.0, initial_state=InitialStateConf.ZERO) diff --git a/ltc/sim/constants.py b/ltc/sim/constants.py index d588386..864db46 100644 --- a/ltc/sim/constants.py +++ b/ltc/sim/constants.py @@ -6,7 +6,8 @@ class Actions(Enum): Actions for the agent """ TX = 0 - IDLE = 1 + CS = 1 + IDLE = 2 """ @@ -20,13 +21,19 @@ class Actions(Enum): Rewards and penalties """ TX_REWARD = 1.0 +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 """ Power consumption """ INITIAL_CAPACITY = 10000 TX_CONSUMPTION = 5 +CS_CONSUMPTION = 1 IDLE_CONSUMPTION = 0 """ diff --git a/ltc/sim/process_output.py b/ltc/sim/process_output.py index f3c4190..9733259 100644 --- a/ltc/sim/process_output.py +++ b/ltc/sim/process_output.py @@ -32,7 +32,8 @@ def successful_transmission(args): def process_output_i(buffer_state, new_buffer_state, power_state, channel_state, obs, action, terminal): - _, _, ret_c, no_tx, _ = obs[-1] + ret_c = 0 + no_tx = 0 args = (action, buffer_state, ret_c, channel_state, no_tx) reward, ret_c, no_tx = jax.lax.cond(action == Actions.TX.value, transmission, no_transmission, args) From db9fda8fd5084d71fa8266eaf07f1b4ee454660b Mon Sep 17 00:00:00 2001 From: kaszczech Date: Wed, 1 Oct 2025 09:52:41 +0200 Subject: [PATCH 03/17] Resolve comments --- ltc/sim/process_output.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/ltc/sim/process_output.py b/ltc/sim/process_output.py index 9733259..874dc0b 100644 --- a/ltc/sim/process_output.py +++ b/ltc/sim/process_output.py @@ -43,7 +43,13 @@ def process_output_i(buffer_state, new_buffer_state, power_state, channel_state, power = jnp.where( action == Actions.TX.value, power_state - TX_CONSUMPTION, - power_state + jnp.where( + action == Actions.CS.value, power_state - CS_CONSUMPTION, + jnp.where( + action == Actions.IDLE.value, power_state - IDLE_CONSUMPTION, + power_state + ) + ) ) obs_t = jnp.array([channel_state, action]) From 04f2d05145545e59b91eff05cc49c996add7563a Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Wed, 1 Oct 2025 17:10:10 +0200 Subject: [PATCH 04/17] Fix initial observations --- ltc/run.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ltc/run.py b/ltc/run.py index 23e3538..8ad67ea 100644 --- a/ltc/run.py +++ b/ltc/run.py @@ -96,7 +96,7 @@ def rl_step_fn(c, _): buffer_states = jnp.zeros(n, dtype=int) power_states = jnp.full(n, INITIAL_CAPACITY, dtype=int) channel_state = 0 - obs = jnp.zeros((n, window_size, 2), dtype=int).at[:, -1].set(INITIAL_CAPACITY) + obs = jnp.zeros((n, window_size, 2), dtype=int) rewards = jnp.zeros(n) terminals = jnp.full(n, False, dtype=bool) From a70a4c75c7b96fcf963b4f25467b970e95738009 Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Wed, 1 Oct 2025 17:20:35 +0200 Subject: [PATCH 05/17] Fix GPU installation --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index b2703de..4144975 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,7 +27,7 @@ dependencies = [ [project.optional-dependencies] gpu = [ - "jax[cuda12_local]", + "jax[cuda12-local]", ] notebook = [ "jupyterlab~=4.3.5", From 6138f6c924191a2d75873d07a013186f571f4530 Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Thu, 2 Oct 2025 13:46:21 +0200 Subject: [PATCH 06/17] Fix DLMA agent --- ltc/agents/__init__.py | 3 +- ltc/agents/dlma.py | 32 +++++++++++++++ ltc/run.py | 15 ++++--- ltc/utils/plots.py | 89 ++++++++++++++++++++++++++---------------- 4 files changed, 96 insertions(+), 43 deletions(-) create mode 100644 ltc/agents/dlma.py diff --git a/ltc/agents/__init__.py b/ltc/agents/__init__.py index 34a1957..db94046 100644 --- a/ltc/agents/__init__.py +++ b/ltc/agents/__init__.py @@ -1,3 +1,4 @@ from ltc.agents.dcf import DCF -from ltc.agents.deep import BayesianDDQN +from ltc.agents.deep import BayesianDDQN, DDQN +from ltc.agents.dlma import DLMANetwork from ltc.agents.drl import QNetwork, QNetworkDropout, StochasticVariationalNetwork diff --git a/ltc/agents/dlma.py b/ltc/agents/dlma.py new file mode 100644 index 0000000..922da6b --- /dev/null +++ b/ltc/agents/dlma.py @@ -0,0 +1,32 @@ +from functools import partial + +import flax.linen as nn + +from ltc.agents.drl import add_batch_dim + + +class DLMANetwork(nn.Module): + num_actions: int = 2 + + @nn.compact + def __call__(self, s): + dense = partial(nn.Dense, kernel_init=nn.initializers.he_normal()) + + s = add_batch_dim(s) + b, *_ = s.shape + x = s.reshape(b, -1) + + for _ in range(2): + x = dense(64)(x) + x = nn.relu(x) + + for _ in range(2): + x_res = x + x = dense(64)(x) + x = nn.relu(x) + x = dense(64)(x) + x = nn.relu(x) + x = x + x_res + + x = dense(self.num_actions)(x) + return x diff --git a/ltc/run.py b/ltc/run.py index 8ad67ea..ab0448e 100644 --- a/ltc/run.py +++ b/ltc/run.py @@ -11,7 +11,7 @@ import optax from tqdm import trange -from ltc.agents import BayesianDDQN, DCF, QNetwork, QNetworkDropout, StochasticVariationalNetwork +from ltc.agents import BayesianDDQN, DCF, DDQN, DLMANetwork, QNetwork, QNetworkDropout, StochasticVariationalNetwork from ltc.sim import InitialStateConf, cox_traffic, process_output, simulate from ltc.sim.constants import INITIAL_CAPACITY, Actions from ltc.utils.scan_states import Carry, Output @@ -53,7 +53,7 @@ def rl_step(drl_step, dcf_step, traffic_step, n, n_drl, n_bins=50): def rl_step_fn(c, _): key, drl_keys, dcf_keys, traffic_key = jax.random.split(c.key, 4) drl_keys = jax.random.split(drl_keys, n_drl) - dcf_keys = jax.random.split(dcf_keys, n - n_drl) + # dcf_keys = jax.random.split(dcf_keys, n - n_drl) traffic_keys = jax.random.split(traffic_key, n) drl_states, drl_actions = drl_step(c.drl_states, drl_keys, c.obs[:n_drl], c.actions[:n_drl], c.rewards[:n_drl], c.terminals[:n_drl]) @@ -88,7 +88,7 @@ def rl_step_fn(c, _): if __name__ == '__main__': n, n_drl = 10, 10 n_epochs, n_steps = 50, 10000 - window_size = 5 + window_size = 20 seed = 42 key = jax.random.key(seed) @@ -100,10 +100,10 @@ def rl_step_fn(c, _): rewards = jnp.zeros(n) terminals = jnp.full(n, False, dtype=bool) - drl = BayesianDDQN( - q_network=StochasticVariationalNetwork(QNetworkDropout()), + drl = DDQN( + q_network=DLMANetwork(), obs_space_shape=obs.shape[1:], - act_space_size=3, + act_space_size=2, optimizer=optax.adam(1e-4), experience_replay_buffer_size=10000, experience_replay_batch_size=128, @@ -121,8 +121,7 @@ def rl_step_fn(c, _): key, init_key = jax.random.split(key) # dcf = DCF() # dcf_states, dcf_step = init_agents(dcf, init_key, n - n_drl) - dcf_states = None - dcf_step = None + dcf_states, dcf_step = None, None key, init_key = jax.random.split(key) traffic = cox_traffic(f3dB=0.1, loc=-5.0, scale=5.0, initial_state=InitialStateConf.ZERO) diff --git a/ltc/utils/plots.py b/ltc/utils/plots.py index 9ec9034..5a3a00c 100644 --- a/ltc/utils/plots.py +++ b/ltc/utils/plots.py @@ -33,6 +33,8 @@ 'ytick.major.width': 0.5, } +SAVE_FORMAT = 'pdf' + class PlotType(Enum): ALL = 'all' @@ -67,8 +69,7 @@ def plot_powers(power_states, n, n_drl, seed, name): plt.legend() plt.grid() plt.tight_layout() - plt.savefig(f'{name.value}_power_{n}_{n_drl}_{seed}.pdf', bbox_inches='tight') - plt.savefig(f'{name.value}_power_{n}_{n_drl}_{seed}.png', bbox_inches='tight', dpi=300) + plt.savefig(f'{name.value}_power_{n}_{n_drl}_{seed}.{SAVE_FORMAT}', bbox_inches='tight', dpi=300) plt.clf() @@ -104,8 +105,7 @@ def plot_power_per_tx(actions, channel_states, power_states, n, n_drl, seed, nam plt.legend() plt.grid() plt.tight_layout() - plt.savefig(f'{name.value}_power_per_succ_tx_{n}_{n_drl}_{seed}.pdf', bbox_inches='tight') - plt.savefig(f'{name.value}_power_per_succ_tx_{n}_{n_drl}_{seed}.png', bbox_inches='tight', dpi=300) + plt.savefig(f'{name.value}_power_per_succ_tx_{n}_{n_drl}_{seed}.{SAVE_FORMAT}', bbox_inches='tight', dpi=300) plt.clf() @@ -135,8 +135,7 @@ def plot_rewards(rewards, n, n_drl, seed, name): plt.legend() plt.grid() plt.tight_layout() - plt.savefig(f'{name.value}_rewards_{n}_{n_drl}_{seed}.pdf', bbox_inches='tight') - plt.savefig(f'{name.value}_rewards_{n}_{n_drl}_{seed}.png', bbox_inches='tight', dpi=300) + plt.savefig(f'{name.value}_rewards_{n}_{n_drl}_{seed}.{SAVE_FORMAT}', bbox_inches='tight', dpi=300) plt.clf() @@ -167,8 +166,7 @@ def plot_cumulative_rewards(rewards, n, n_drl, seed, name): plt.legend() plt.grid() plt.tight_layout() - plt.savefig(f'{name.value}_cum_rewards_{n}_{n_drl}_{seed}.pdf', bbox_inches='tight') - plt.savefig(f'{name.value}_cum_rewards_{n}_{n_drl}_{seed}.png', bbox_inches='tight', dpi=300) + plt.savefig(f'{name.value}_cum_rewards_{n}_{n_drl}_{seed}.{SAVE_FORMAT}', bbox_inches='tight', dpi=300) plt.clf() @@ -201,8 +199,41 @@ def plot_successful_transmissions(actions, channel_states, n, n_drl, seed, name) plt.legend() plt.grid() plt.tight_layout() - plt.savefig(f'{name.value}_tx_{n}_{n_drl}_{seed}.pdf', bbox_inches='tight') - plt.savefig(f'{name.value}_tx_{n}_{n_drl}_{seed}.png', bbox_inches='tight', dpi=300) + plt.savefig(f'{name.value}_tx_{n}_{n_drl}_{seed}.{SAVE_FORMAT}', bbox_inches='tight', dpi=300) + plt.clf() + + +def plot_cum_successful_transmissions(actions, channel_states, n, n_drl, seed, name): + plt.rcParams.update(PLOT_PARAMS) + + n_epochs, n_steps = actions.shape[:2] + if name == PlotType.ALL: + xs = np.arange(n_epochs) + 1 + else: + xs = np.linspace(0, n_epochs * n_steps, n_epochs) + 1 + + actions = actions.at[*np.where(channel_states != 1), :].set(-1) + actions = actions == Actions.TX.value + actions = actions.sum(axis=1) + actions = actions.cumsum(axis=0) + + for i in range(n_drl): + plt.plot(xs, actions[:, i], color='red') + + for i in range(n_drl, n): + plt.plot(xs, actions[:, i], color='blue', linestyle='--') + + plt.plot([], color='blue', linestyle='--', label='DCF') + plt.plot([], color='red', label='DRL') + + plt.xlabel('Epoch' if name == PlotType.ALL else 'Step') + plt.ylabel('Successful transmissions') + plt.xlim(1, n_epochs if name == PlotType.ALL else n_epochs * n_steps) + plt.ylim(bottom=0) + plt.legend() + plt.grid() + plt.tight_layout() + plt.savefig(f'{name.value}_tx_cum_{n}_{n_drl}_{seed}.{SAVE_FORMAT}', bbox_inches='tight', dpi=300) plt.clf() @@ -252,8 +283,7 @@ def plot_actions(actions, terminals, n, n_drl, seed, name): plt.legend(ncol=2) plt.grid() plt.tight_layout() - plt.savefig(f'{name.value}_action_{n}_{n_drl}_{seed}.pdf', bbox_inches='tight') - plt.savefig(f'{name.value}_action_{n}_{n_drl}_{seed}.png', bbox_inches='tight', dpi=300) + plt.savefig(f'{name.value}_action_{n}_{n_drl}_{seed}.{SAVE_FORMAT}', bbox_inches='tight', dpi=300) plt.clf() @@ -286,8 +316,7 @@ def plot_buffer_states(buffer_states, terminals, n, n_drl, seed, name): plt.legend() plt.grid() plt.tight_layout() - plt.savefig(f'{name.value}_buffer_{n}_{n_drl}_{seed}.pdf', bbox_inches='tight') - plt.savefig(f'{name.value}_buffer_{n}_{n_drl}_{seed}.png', bbox_inches='tight', dpi=300) + plt.savefig(f'{name.value}_buffer_{n}_{n_drl}_{seed}.{SAVE_FORMAT}', bbox_inches='tight', dpi=300) plt.clf() @@ -320,8 +349,7 @@ def plot_channel_states(channel_states, n, n_drl, seed, name): plt.legend() plt.grid() plt.tight_layout() - plt.savefig(f'{name.value}_channel_{n}_{n_drl}_{seed}.pdf', bbox_inches='tight') - plt.savefig(f'{name.value}_channel_{n}_{n_drl}_{seed}.png', bbox_inches='tight', dpi=300) + plt.savefig(f'{name.value}_channel_{n}_{n_drl}_{seed}.{SAVE_FORMAT}', bbox_inches='tight', dpi=300) plt.clf() @@ -360,8 +388,7 @@ def plot_channel_states_fill(channel_states, n, n_drl, seed, name): plt.legend() plt.grid() plt.tight_layout() - plt.savefig(f'{name.value}_channel_fill_{n}_{n_drl}_{seed}.pdf', bbox_inches='tight') - plt.savefig(f'{name.value}_channel_fill_{n}_{n_drl}_{seed}.png', bbox_inches='tight', dpi=300) + plt.savefig(f'{name.value}_channel_fill_{n}_{n_drl}_{seed}.{SAVE_FORMAT}', bbox_inches='tight', dpi=300) plt.clf() @@ -393,8 +420,7 @@ def plot_throughput(actions, channel_states, terminals, n, n_drl, seed, name): plt.legend() plt.grid() plt.tight_layout() - plt.savefig(f'{name.value}_throughput_{n}_{n_drl}_{seed}.pdf', bbox_inches='tight') - plt.savefig(f'{name.value}_throughput_{n}_{n_drl}_{seed}.png', bbox_inches='tight', dpi=300) + plt.savefig(f'{name.value}_throughput_{n}_{n_drl}_{seed}.{SAVE_FORMAT}', bbox_inches='tight', dpi=300) plt.clf() @@ -428,8 +454,7 @@ def plot_throughput_fill(actions, channel_states, terminals, n, n_drl, seed, nam plt.legend() plt.grid() plt.tight_layout() - plt.savefig(f'{name.value}_throughput_fill_{n}_{n_drl}_{seed}.pdf', bbox_inches='tight') - plt.savefig(f'{name.value}_throughput_fill_{n}_{n_drl}_{seed}.png', bbox_inches='tight', dpi=300) + plt.savefig(f'{name.value}_throughput_fill_{n}_{n_drl}_{seed}.{SAVE_FORMAT}', bbox_inches='tight', dpi=300) plt.clf() @@ -461,8 +486,7 @@ def plot_throughput_fill_nn(actions, channel_states, terminals, n, n_drl, seed, plt.legend() plt.grid() plt.tight_layout() - plt.savefig(f'{name.value}_throughput_fill_nn_{n}_{n_drl}_{seed}.pdf', bbox_inches='tight') - plt.savefig(f'{name.value}_throughput_fill_nn_{n}_{n_drl}_{seed}.png', bbox_inches='tight', dpi=300) + plt.savefig(f'{name.value}_throughput_fill_nn_{n}_{n_drl}_{seed}.{SAVE_FORMAT}', bbox_inches='tight', dpi=300) plt.clf() @@ -494,8 +518,7 @@ def plot_channel_access_delay(buffer_states, new_frames, terminals, n, n_drl, se plt.legend() plt.grid() plt.tight_layout() - plt.savefig(f'{name.value}_ch_acc_delay_{n}_{n_drl}_{seed}.pdf', bbox_inches='tight') - plt.savefig(f'{name.value}_ch_acc_delay_{n}_{n_drl}_{seed}.png', bbox_inches='tight', dpi=300) + plt.savefig(f'{name.value}_ch_acc_delay_{n}_{n_drl}_{seed}.{SAVE_FORMAT}', bbox_inches='tight', dpi=300) plt.clf() @@ -534,8 +557,7 @@ def plot_xnor(actions, channel_states, buffer_states, terminals, n, n_drl, seed, plt.legend() plt.grid() plt.tight_layout() - plt.savefig(f'{name.value}_xnor_{n}_{n_drl}_{seed}.pdf', bbox_inches='tight') - plt.savefig(f'{name.value}_xnor_{n}_{n_drl}_{seed}.png', bbox_inches='tight', dpi=300) + plt.savefig(f'{name.value}_xnor_{n}_{n_drl}_{seed}.{SAVE_FORMAT}', bbox_inches='tight', dpi=300) plt.clf() @@ -570,8 +592,7 @@ def plot_weights(histogram, bin_edges, n, n_drl, seed, name): plt.yticks(np.linspace(0, (n_epochs - 1) * offset, 11), labels=ys) plt.grid(alpha=0.3) plt.tight_layout() - plt.savefig(f'{name.value}_hist_id{i}_{n}_{n_drl}_{seed}.pdf', bbox_inches='tight') - plt.savefig(f'{name.value}_hist_id{i}_{n}_{n_drl}_{seed}.png', bbox_inches='tight', dpi=300) + plt.savefig(f'{name.value}_hist_id{i}_{n}_{n_drl}_{seed}.{SAVE_FORMAT}', bbox_inches='tight', dpi=300) plt.clf() @@ -588,6 +609,7 @@ def plot_all(filename): plot_rewards(history.rewards, n, n_drl, seed, PlotType.ALL) plot_cumulative_rewards(history.rewards, n, n_drl, seed, PlotType.ALL) plot_successful_transmissions(history.actions, history.channel_state, n, n_drl, seed, PlotType.ALL) + plot_cum_successful_transmissions(history.actions, history.channel_state, n, n_drl, seed, PlotType.ALL) plot_actions(history.actions, history.terminals, n, n_drl, seed, PlotType.ALL) plot_buffer_states(history.buffer_states, history.terminals, n, n_drl, seed, PlotType.ALL) plot_channel_states(history.channel_state, n, n_drl, seed, PlotType.ALL) @@ -616,6 +638,7 @@ def plot_first(filename, n_epochs=10, aggregation=500): plot_rewards(history.rewards, n, n_drl, seed, PlotType.FIRST) plot_cumulative_rewards(history.rewards, n, n_drl, seed, PlotType.FIRST) plot_successful_transmissions(history.actions, history.channel_state, n, n_drl, seed, PlotType.FIRST) + plot_cum_successful_transmissions(history.actions, history.channel_state, n, n_drl, seed, PlotType.FIRST) plot_actions(history.actions, history.terminals, n, n_drl, seed, PlotType.FIRST) plot_buffer_states(history.buffer_states, history.terminals, n, n_drl, seed, PlotType.FIRST) plot_channel_states(history.channel_state, n, n_drl, seed, PlotType.FIRST) @@ -648,8 +671,7 @@ def plot_dcf_collision_probabilities(actions, rewards, seed, name): plt.legend() plt.grid() plt.tight_layout() - plt.savefig(f'{name.value}_collision_{seed}.pdf', bbox_inches='tight') - plt.savefig(f'{name.value}_collision_{seed}.png', bbox_inches='tight', dpi=300) + plt.savefig(f'{name.value}_collision_{seed}.{SAVE_FORMAT}', bbox_inches='tight', dpi=300) plt.clf() @@ -668,8 +690,7 @@ def plot_cw_values(backoff, actions, n, n_drl, seed, name): plt.xticks(cw, fontsize=7) plt.grid(axis='y') plt.tight_layout() - plt.savefig(f'{name.value}_cw_{n}_{n_drl}_{seed}.pdf', bbox_inches='tight') - plt.savefig(f'{name.value}_cw_{n}_{n_drl}_{seed}.png', bbox_inches='tight', dpi=300) + plt.savefig(f'{name.value}_cw_{n}_{n_drl}_{seed}.{SAVE_FORMAT}', bbox_inches='tight', dpi=300) plt.clf() From 108ecb23d6465888923311b1bc80c4e1994039e4 Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Fri, 3 Oct 2025 10:13:45 +0200 Subject: [PATCH 07/17] Add all baselines from DLMA --- ltc/agents/__init__.py | 4 +++ ltc/agents/dlma.py | 2 ++ ltc/agents/eb_aloha.py | 53 ++++++++++++++++++++++++++++ ltc/agents/fw_aloha.py | 43 +++++++++++++++++++++++ ltc/agents/q_aloha.py | 42 ++++++++++++++++++++++ ltc/agents/tdma.py | 44 +++++++++++++++++++++++ ltc/run.py | 74 ++++++++++++++++++++++++--------------- ltc/sim/constants.py | 2 +- ltc/sim/process_output.py | 2 +- ltc/utils/plots.py | 28 +++++++-------- ltc/utils/scan_states.py | 4 +-- pyproject.toml | 4 +-- 12 files changed, 254 insertions(+), 48 deletions(-) create mode 100644 ltc/agents/eb_aloha.py create mode 100644 ltc/agents/fw_aloha.py create mode 100644 ltc/agents/q_aloha.py create mode 100644 ltc/agents/tdma.py diff --git a/ltc/agents/__init__.py b/ltc/agents/__init__.py index db94046..c6f764a 100644 --- a/ltc/agents/__init__.py +++ b/ltc/agents/__init__.py @@ -2,3 +2,7 @@ from ltc.agents.deep import BayesianDDQN, DDQN from ltc.agents.dlma import DLMANetwork from ltc.agents.drl import QNetwork, QNetworkDropout, StochasticVariationalNetwork +from ltc.agents.eb_aloha import EBALOHA +from ltc.agents.fw_aloha import FWALOHA +from ltc.agents.q_aloha import QALOHA +from ltc.agents.tdma import TDMA diff --git a/ltc/agents/dlma.py b/ltc/agents/dlma.py index 922da6b..1c017d8 100644 --- a/ltc/agents/dlma.py +++ b/ltc/agents/dlma.py @@ -12,7 +12,9 @@ class DLMANetwork(nn.Module): def __call__(self, s): dense = partial(nn.Dense, kernel_init=nn.initializers.he_normal()) + s = s[..., :-1] # remove buffer state s = add_batch_dim(s) + b, *_ = s.shape x = s.reshape(b, -1) diff --git a/ltc/agents/eb_aloha.py b/ltc/agents/eb_aloha.py new file mode 100644 index 0000000..08ce77c --- /dev/null +++ b/ltc/agents/eb_aloha.py @@ -0,0 +1,53 @@ +from functools import partial + +import jax +import jax.numpy as jnp +from chex import dataclass +from reinforced_lib.agents import BaseAgent, AgentState + +from ltc.sim.constants import Actions + + +@dataclass +class EBALOHAState(AgentState): + backoff: int + ret_c: int + + +class EBALOHA(BaseAgent): + def __init__(self, window_size: int, max_backoff: int): + self.window_size = window_size + self.max_backoff = max_backoff + self.init = jax.jit(partial(self.init, window_size=window_size)) + self.update = jax.jit(partial(self.update, window_size=window_size, max_backoff=max_backoff)) + self.sample = jax.jit(self.sample) + + @staticmethod + def init(key, window_size: int): + ret_c = 0 + backoff = jax.random.randint(key, shape=(), minval=0, maxval=window_size) + return EBALOHAState(backoff=backoff, ret_c=ret_c) + + @staticmethod + def update(state, key, env_state, action, reward, terminal, window_size: int, max_backoff: int): + *_, buffer = env_state[-1] + + collision = (action == Actions.TX.value) & (reward == 0) & (buffer > 0) + new_ret_c = jnp.where(collision, jnp.minimum(state.ret_c + 1, max_backoff), 0) + window = window_size * (2 ** new_ret_c) + + reset_value = jax.random.randint(key, shape=(), minval=0, maxval=window) + new_backoff = state.backoff - 1 + new_backoff = jnp.where(new_backoff < 0, reset_value, new_backoff) + + return EBALOHAState(backoff=new_backoff, ret_c=new_ret_c) + + @staticmethod + def sample(state, key, env_state): + *_, buffer = env_state[-1] + + return jnp.where( + (buffer > 0) & (state.backoff == 0), + Actions.TX.value, + Actions.IDLE.value + ) diff --git a/ltc/agents/fw_aloha.py b/ltc/agents/fw_aloha.py new file mode 100644 index 0000000..fa338d8 --- /dev/null +++ b/ltc/agents/fw_aloha.py @@ -0,0 +1,43 @@ +from functools import partial + +import jax +import jax.numpy as jnp +from chex import dataclass +from reinforced_lib.agents import BaseAgent, AgentState + +from ltc.sim.constants import Actions + + +@dataclass +class FWALOHAState(AgentState): + backoff: int + + +class FWALOHA(BaseAgent): + def __init__(self, window_size: int): + self.window_size = window_size + self.init = jax.jit(partial(self.init, window_size=window_size)) + self.update = jax.jit(partial(self.update, window_size=window_size)) + self.sample = jax.jit(self.sample) + + @staticmethod + def init(key, window_size: int): + backoff = jax.random.randint(key, shape=(), minval=0, maxval=window_size) + return FWALOHAState(backoff=backoff) + + @staticmethod + def update(state, key, env_state, action, reward, terminal, window_size: int): + reset_value = jax.random.randint(key, shape=(), minval=0, maxval=window_size) + new_backoff = state.backoff - 1 + new_backoff = jnp.where(new_backoff < 0, reset_value, new_backoff) + return FWALOHAState(backoff=new_backoff) + + @staticmethod + def sample(state, key, env_state): + *_, buffer = env_state[-1] + + return jnp.where( + (buffer > 0) & (state.backoff == 0), + Actions.TX.value, + Actions.IDLE.value + ) diff --git a/ltc/agents/q_aloha.py b/ltc/agents/q_aloha.py new file mode 100644 index 0000000..1c422b5 --- /dev/null +++ b/ltc/agents/q_aloha.py @@ -0,0 +1,42 @@ +from functools import partial + +import jax +import jax.numpy as jnp +from chex import dataclass +from reinforced_lib.agents import BaseAgent, AgentState + +from ltc.sim.constants import Actions + + +@dataclass +class QALOHAState(AgentState): + pass + + +class QALOHA(BaseAgent): + def __init__(self, q): + self.init = jax.jit(self.init) + self.update = jax.jit(self.update) + self.sample = jax.jit(partial(self.sample, q=q)) + + @staticmethod + def init(key): + return QALOHAState() + + @staticmethod + def update(state, key, env_state, action, reward, terminal): + return state + + @staticmethod + def sample(state, key, env_state, q): + *_, buffer = env_state[-1] + + return jnp.where( + buffer == 0, + Actions.IDLE.value, + jnp.where( + jax.random.uniform(key) < q, + Actions.TX.value, + Actions.IDLE.value + ) + ) diff --git a/ltc/agents/tdma.py b/ltc/agents/tdma.py new file mode 100644 index 0000000..d30493d --- /dev/null +++ b/ltc/agents/tdma.py @@ -0,0 +1,44 @@ +from functools import partial + +import jax +import jax.numpy as jnp +from chex import dataclass +from reinforced_lib.agents import BaseAgent, AgentState + +from ltc.sim.constants import Actions + + +@dataclass +class TDMAState(AgentState): + counter: int + tdma_slots: jnp.ndarray + + +class TDMA(BaseAgent): + def __init__(self, state_size: int, assigned_slots: int): + self.state_size = state_size + self.assigned_slots = assigned_slots + self.init = jax.jit(partial(self.init, state_size=state_size, assigned_slots=assigned_slots)) + self.update = jax.jit(self.update) + self.sample = jax.jit(self.sample) + + @staticmethod + def init(key, state_size: int, assigned_slots: int): + tdma_slots = jnp.ones(state_size, dtype=int).at[:assigned_slots].set(1) + tdma_slots = jax.random.permutation(key, tdma_slots) + return TDMAState(counter=0, tdma_slots=tdma_slots) + + @staticmethod + def update(state, key, env_state, action, reward, terminal): + new_count = (state.counter + 1) % state.tdma_slots.shape[0] + return TDMAState(counter=new_count, tdma_slots=state.tdma_slots) + + @staticmethod + def sample(state, key, env_state): + *_, buffer = env_state[-1] + + return jnp.where( + (buffer > 0) & (state.tdma_slots[state.counter] == 1), + Actions.TX.value, + Actions.IDLE.value + ) diff --git a/ltc/run.py b/ltc/run.py index ab0448e..7070b2c 100644 --- a/ltc/run.py +++ b/ltc/run.py @@ -11,7 +11,7 @@ import optax from tqdm import trange -from ltc.agents import BayesianDDQN, DCF, DDQN, DLMANetwork, QNetwork, QNetworkDropout, StochasticVariationalNetwork +from ltc.agents import BayesianDDQN, DDQN, DLMANetwork, QNetwork, QNetworkDropout, StochasticVariationalNetwork, DCF, QALOHA, EBALOHA, FWALOHA, TDMA from ltc.sim import InitialStateConf, cox_traffic, process_output, simulate from ltc.sim.constants import INITIAL_CAPACITY, Actions from ltc.utils.scan_states import Carry, Output @@ -49,17 +49,16 @@ def init_traffic(traffic, key, n): return states, step_fn -def rl_step(drl_step, dcf_step, traffic_step, n, n_drl, n_bins=50): +def rl_step(drl_step, legacy_step, traffic_step, n, n_drl, n_bins=50): def rl_step_fn(c, _): - key, drl_keys, dcf_keys, traffic_key = jax.random.split(c.key, 4) + key, drl_keys, legacy_keys, traffic_key = jax.random.split(c.key, 4) drl_keys = jax.random.split(drl_keys, n_drl) - # dcf_keys = jax.random.split(dcf_keys, n - n_drl) + legacy_keys = jax.random.split(legacy_keys, n - n_drl) traffic_keys = jax.random.split(traffic_key, n) drl_states, drl_actions = drl_step(c.drl_states, drl_keys, c.obs[:n_drl], c.actions[:n_drl], c.rewards[:n_drl], c.terminals[:n_drl]) - # dcf_states, dcf_actions = dcf_step(c.dcf_states, dcf_keys, c.obs[n_drl:], c.actions[n_drl:], c.rewards[n_drl:], c.terminals[n_drl:]) - # actions = jnp.concatenate([drl_actions, dcf_actions]) - actions = drl_actions + legacy_states, legacy_actions = legacy_step(c.legacy_states, legacy_keys, c.obs[n_drl:], c.actions[n_drl:], c.rewards[n_drl:], c.terminals[n_drl:]) + 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) @@ -73,11 +72,11 @@ def rl_step_fn(c, _): hist, bin_edges = jax.vmap(jnp.histogram, in_axes=(0, None))(flat_params, n_bins) c = Carry( - drl_states, dcf_states, traffic_states, buffer_states, powers, + drl_states, legacy_states, traffic_states, buffer_states, powers, channel_state, key, obs, actions, rewards, terminals ) o = Output( - dcf_states, obs, actions, rewards, terminals, buffer_states, powers, + legacy_states, obs, actions, rewards, terminals, buffer_states, powers, (new_frames > 0).astype(int), channel_state, hist, bin_edges ) return c, o @@ -86,17 +85,19 @@ def rl_step_fn(c, _): if __name__ == '__main__': - n, n_drl = 10, 10 - n_epochs, n_steps = 50, 10000 + n, n_drl = 2, 1 + n_epochs, n_steps = 1, 500000 window_size = 20 seed = 42 + traffic_type = 'saturated' # 'saturated', 'bursty' + legacy_type = 'q-aloha' # 'q-aloha', 'eb-aloha', 'fw-aloha', 'tdma' key = jax.random.key(seed) actions = jnp.zeros(n, dtype=int) buffer_states = jnp.zeros(n, dtype=int) power_states = jnp.full(n, INITIAL_CAPACITY, dtype=int) channel_state = 0 - obs = jnp.zeros((n, window_size, 2), dtype=int) + obs = jnp.zeros((n, window_size, 3), dtype=int) rewards = jnp.zeros(n) terminals = jnp.full(n, False, dtype=bool) @@ -104,32 +105,49 @@ def rl_step_fn(c, _): q_network=DLMANetwork(), obs_space_shape=obs.shape[1:], act_space_size=2, - optimizer=optax.adam(1e-4), - experience_replay_buffer_size=10000, - experience_replay_batch_size=128, - experience_replay_steps=5, - discount=1.0, - epsilon=1.0, - epsilon_decay=0.999, - epsilon_min=0.001, - tau=0.01 + optimizer=optax.rmsprop(1e-2), + experience_replay_buffer_size=500, + experience_replay_batch_size=32, + experience_replay_steps=1, + discount=0.9, + epsilon=0.1, + epsilon_decay=0.9999, + epsilon_min=0.005, + tau=0.02 ) key, init_key = jax.random.split(key) drl_states, drl_step = init_agents(drl, init_key, n_drl) drl_states = drl_states.replace(prev_env_state=drl_states.prev_env_state.astype(int)) key, init_key = jax.random.split(key) - # dcf = DCF() - # dcf_states, dcf_step = init_agents(dcf, init_key, n - n_drl) - dcf_states, dcf_step = None, None + + if legacy_type == 'q-aloha': + legacy = QALOHA(q=0.2) + elif legacy_type == 'eb-aloha': + legacy = EBALOHA(window_size=4, max_backoff=3) + elif legacy_type == 'fw-aloha': + legacy = FWALOHA(window_size=8) + elif legacy_type == 'tdma': + legacy = TDMA(state_size=10, assigned_slots=5) + else: + raise ValueError(f'Unknown legacy type: {legacy_type}') + + legacy_states, legacy_step = init_agents(legacy, init_key, n - n_drl) key, init_key = jax.random.split(key) - traffic = cox_traffic(f3dB=0.1, loc=-5.0, scale=5.0, initial_state=InitialStateConf.ZERO) + + if traffic_type == 'saturated': + traffic = cox_traffic(f3dB=1.0, loc=0.0, scale=0.0, initial_state=InitialStateConf.ZERO) + elif traffic_type == 'bursty': + traffic = cox_traffic(f3dB=0.1, loc=-5.0, scale=5.0, initial_state=InitialStateConf.ZERO) + else: + raise ValueError(f'Unknown traffic type: {traffic_type}') + traffic_states, traffic_step = init_traffic(traffic, init_key, n) - rl_step_fn = jax.jit(rl_step(drl_step, dcf_step, traffic_step, n, n_drl)) + rl_step_fn = jax.jit(rl_step(drl_step, legacy_step, traffic_step, n, n_drl)) init_carry = Carry( - drl_states, dcf_states, traffic_states, buffer_states, power_states, + drl_states, legacy_states, traffic_states, buffer_states, power_states, channel_state, key, obs, actions, rewards, terminals ) all_outputs = [] @@ -149,5 +167,5 @@ def rl_step_fn(c, _): with lz4.frame.open(filename, 'wb') as f: cloudpickle.dump((init_carry.drl_states, all_outputs), f) - plot_all(filename) + # plot_all(filename) plot_first(filename) diff --git a/ltc/sim/constants.py b/ltc/sim/constants.py index 864db46..fa16e1d 100644 --- a/ltc/sim/constants.py +++ b/ltc/sim/constants.py @@ -31,7 +31,7 @@ class Actions(Enum): """ Power consumption """ -INITIAL_CAPACITY = 10000 +INITIAL_CAPACITY = int(1e9) TX_CONSUMPTION = 5 CS_CONSUMPTION = 1 IDLE_CONSUMPTION = 0 diff --git a/ltc/sim/process_output.py b/ltc/sim/process_output.py index 874dc0b..2e65c68 100644 --- a/ltc/sim/process_output.py +++ b/ltc/sim/process_output.py @@ -52,7 +52,7 @@ def process_output_i(buffer_state, new_buffer_state, power_state, channel_state, ) ) - obs_t = jnp.array([channel_state, action]) + obs_t = jnp.array([channel_state, action, new_buffer_state]) obs = jnp.roll(obs, -1, axis=0) obs = obs.at[-1].set(obs_t) diff --git a/ltc/utils/plots.py b/ltc/utils/plots.py index 5a3a00c..7de3822 100644 --- a/ltc/utils/plots.py +++ b/ltc/utils/plots.py @@ -58,7 +58,7 @@ def plot_powers(power_states, n, n_drl, seed, name): for i in range(n_drl, n): plt.plot(xs, consumed_power[:, i], color='blue', linestyle='--') - plt.plot([], color='blue', linestyle='--', label='DCF') + plt.plot([], color='blue', linestyle='--', label='Legacy') plt.plot([], color='red', label='DRL') plt.xlabel('Epoch' if name == PlotType.ALL else 'Step') @@ -95,7 +95,7 @@ def plot_power_per_tx(actions, channel_states, power_states, n, n_drl, seed, nam for i in range(n_drl, n): plt.plot(xs, power_per_tx[:, i], color='blue', linestyle='--') - plt.plot([], color='blue', linestyle='--', label='DCF') + plt.plot([], color='blue', linestyle='--', label='Legacy') plt.plot([], color='red', label='DRL') plt.xlabel('Epoch' if name == PlotType.ALL else 'Step') @@ -126,7 +126,7 @@ def plot_rewards(rewards, n, n_drl, seed, name): for i in range(n_drl, n): plt.plot(xs, rewards[:, i], color='blue', linestyle='--') - plt.plot([], color='blue', linestyle='--', label='DCF') + plt.plot([], color='blue', linestyle='--', label='Legacy') plt.plot([], color='red', label='DRL') plt.xlabel('Epoch' if name == PlotType.ALL else 'Step') @@ -157,7 +157,7 @@ def plot_cumulative_rewards(rewards, n, n_drl, seed, name): for i in range(n_drl, n): plt.plot(xs, cum_rewards[:, i], color='blue', linestyle='--') - plt.plot([], color='blue', linestyle='--', label='DCF') + plt.plot([], color='blue', linestyle='--', label='Legacy') plt.plot([], color='red', label='DRL') plt.xlabel('Epoch' if name == PlotType.ALL else 'Step') @@ -189,7 +189,7 @@ def plot_successful_transmissions(actions, channel_states, n, n_drl, seed, name) for i in range(n_drl, n): plt.plot(xs, actions[:, i], color='blue', linestyle='--') - plt.plot([], color='blue', linestyle='--', label='DCF') + plt.plot([], color='blue', linestyle='--', label='Legacy') plt.plot([], color='red', label='DRL') plt.xlabel('Epoch' if name == PlotType.ALL else 'Step') @@ -223,7 +223,7 @@ def plot_cum_successful_transmissions(actions, channel_states, n, n_drl, seed, n for i in range(n_drl, n): plt.plot(xs, actions[:, i], color='blue', linestyle='--') - plt.plot([], color='blue', linestyle='--', label='DCF') + plt.plot([], color='blue', linestyle='--', label='Legacy') plt.plot([], color='red', label='DRL') plt.xlabel('Epoch' if name == PlotType.ALL else 'Step') @@ -273,7 +273,7 @@ def plot_actions(actions, terminals, n, n_drl, seed, name): plt.plot([], color='green', label='CS') plt.plot([], color='blue', label='Idle') plt.plot([], color='k', linestyle='-', label='DRL') - plt.plot([], color='k', linestyle='--', label='DCF', alpha=0.5) + plt.plot([], color='k', linestyle='--', label='Legacy', alpha=0.5) plt.xlabel('Epoch' if name == PlotType.ALL else 'Step') plt.ylabel('Action') @@ -305,7 +305,7 @@ def plot_buffer_states(buffer_states, terminals, n, n_drl, seed, name): for i in range(n_drl, n): plt.plot(xs, buffer_states[:, i], color='blue', linestyle='--') - plt.plot([], color='blue', linestyle='--', label='DCF') + plt.plot([], color='blue', linestyle='--', label='Legacy') plt.plot([], color='red', label='DRL') plt.xlabel('Epoch' if name == PlotType.ALL else 'Step') @@ -410,7 +410,7 @@ def plot_throughput(actions, channel_states, terminals, n, n_drl, seed, name): for i in range(n_drl, n): plt.plot(xs, throughput[:, i], color='blue', linestyle='--') - plt.plot([], color='blue', linestyle='--', label='DCF') + plt.plot([], color='blue', linestyle='--', label='Legacy') plt.plot([], color='red', label='DRL') plt.xlabel('Epoch' if name == PlotType.ALL else 'Step') @@ -444,7 +444,7 @@ def plot_throughput_fill(actions, channel_states, terminals, n, n_drl, seed, nam plt.plot(xs, throughput[:, i], color='gray') plt.fill_between(xs, np.zeros_like(xs), throughput[:, n_drl - 1], color='red', alpha=0.3, label='DRL', linewidth=0) - plt.fill_between(xs, throughput[:, n_drl - 1], throughput[:, -1], color='blue', alpha=0.3, label='DCF', linewidth=0) + plt.fill_between(xs, throughput[:, n_drl - 1], throughput[:, -1], color='blue', alpha=0.3, label='Legacy', linewidth=0) plt.xlabel('Epoch' if name == PlotType.ALL else 'Step') plt.ylabel('Throughput') @@ -477,7 +477,7 @@ def plot_throughput_fill_nn(actions, channel_states, terminals, n, n_drl, seed, plt.plot(xs, throughput[:, i], color='gray') plt.fill_between(xs, np.zeros_like(xs), throughput[:, n_drl - 1], color='red', alpha=0.3, label='DRL', linewidth=0) - plt.fill_between(xs, throughput[:, n_drl - 1], throughput[:, -1], color='blue', alpha=0.3, label='DCF', linewidth=0) + plt.fill_between(xs, throughput[:, n_drl - 1], throughput[:, -1], color='blue', alpha=0.3, label='Legacy', linewidth=0) plt.xlabel('Epoch' if name == PlotType.ALL else 'Step') plt.ylabel('Throughput') @@ -508,7 +508,7 @@ def plot_channel_access_delay(buffer_states, new_frames, terminals, n, n_drl, se for i in range(n_drl, n): plt.plot(xs, cad[:, i], color='blue', linestyle='--') - plt.plot([], color='blue', linestyle='--', label='DCF') + plt.plot([], color='blue', linestyle='--', label='Legacy') plt.plot([], color='red', label='DRL') plt.xlabel('Epoch' if name == PlotType.ALL else 'Step') @@ -547,7 +547,7 @@ def plot_xnor(actions, channel_states, buffer_states, terminals, n, n_drl, seed, for i in range(n_drl, n): plt.plot(xs, xnor[:, i], color='blue', linestyle='--') - plt.plot([], color='blue', linestyle='--', label='DCF') + plt.plot([], color='blue', linestyle='--', label='Legacy') plt.plot([], color='red', label='DRL') plt.xlabel('Epoch' if name == PlotType.ALL else 'Step') @@ -622,7 +622,7 @@ def plot_all(filename): plot_weights(history.weights_histogram, history.weights_bin_edges, n, n_drl, seed, PlotType.ALL) -def plot_first(filename, n_epochs=10, aggregation=500): +def plot_first(filename, n_epochs=10, aggregation=1000): with lz4.frame.open(filename, 'rb') as f: _, history = cloudpickle.load(f) history = jax.tree.map(lambda x: x[:n_epochs], asdict(history)) diff --git a/ltc/utils/scan_states.py b/ltc/utils/scan_states.py index 6342ae8..f048a3e 100644 --- a/ltc/utils/scan_states.py +++ b/ltc/utils/scan_states.py @@ -10,7 +10,7 @@ @dataclass class Carry: drl_states: AgentState - dcf_states: AgentState + legacy_states: AgentState traffic_states: ModelState buffer_states: jax.Array power_states: jax.Array @@ -25,7 +25,7 @@ class Carry: @jax.tree_util.register_dataclass @dataclass class Output: - dcf_states: AgentState + legacy_states: AgentState observations: jax.Array actions: jax.Array rewards: jax.Array diff --git a/pyproject.toml b/pyproject.toml index 4144975..93dfc63 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,7 +21,7 @@ authors = [ ] dependencies = [ - "reinforced-lib==1.2.2", + "reinforced-lib==2.0.0", "tqdm~=4.67.1", ] @@ -30,7 +30,7 @@ gpu = [ "jax[cuda12-local]", ] notebook = [ - "jupyterlab~=4.3.5", + "jupyterlab~=4.4.9", "seaborn~=0.13.2", ] From 903b1fd8466b0c45c8840b3b47580ef7c1f744bc Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Fri, 3 Oct 2025 12:48:00 +0200 Subject: [PATCH 08/17] Add tqdm for jax scan --- ltc/agents/tdma.py | 2 +- ltc/run.py | 22 +++++++++++++--------- pyproject.toml | 1 + 3 files changed, 15 insertions(+), 10 deletions(-) diff --git a/ltc/agents/tdma.py b/ltc/agents/tdma.py index d30493d..02c0f02 100644 --- a/ltc/agents/tdma.py +++ b/ltc/agents/tdma.py @@ -24,7 +24,7 @@ def __init__(self, state_size: int, assigned_slots: int): @staticmethod def init(key, state_size: int, assigned_slots: int): - tdma_slots = jnp.ones(state_size, dtype=int).at[:assigned_slots].set(1) + tdma_slots = jnp.zeros(state_size, dtype=int).at[:assigned_slots].set(1) tdma_slots = jax.random.permutation(key, tdma_slots) return TDMAState(counter=0, tdma_slots=tdma_slots) diff --git a/ltc/run.py b/ltc/run.py index 7070b2c..ba4945b 100644 --- a/ltc/run.py +++ b/ltc/run.py @@ -9,7 +9,7 @@ import jax.numpy as jnp import lz4.frame import optax -from tqdm import trange +from jax_tqdm import scan_tqdm from ltc.agents import BayesianDDQN, DDQN, DLMANetwork, QNetwork, QNetworkDropout, StochasticVariationalNetwork, DCF, QALOHA, EBALOHA, FWALOHA, TDMA from ltc.sim import InitialStateConf, cox_traffic, process_output, simulate @@ -89,7 +89,7 @@ def rl_step_fn(c, _): n_epochs, n_steps = 1, 500000 window_size = 20 seed = 42 - traffic_type = 'saturated' # 'saturated', 'bursty' + traffic_type = 'saturated' # 'constant', 'saturated', 'bursty' legacy_type = 'q-aloha' # 'q-aloha', 'eb-aloha', 'fw-aloha', 'tdma' key = jax.random.key(seed) @@ -122,11 +122,11 @@ def rl_step_fn(c, _): key, init_key = jax.random.split(key) if legacy_type == 'q-aloha': - legacy = QALOHA(q=0.2) + legacy = QALOHA(q=0.5) elif legacy_type == 'eb-aloha': - legacy = EBALOHA(window_size=4, max_backoff=3) + legacy = EBALOHA(window_size=4, max_backoff=2) elif legacy_type == 'fw-aloha': - legacy = FWALOHA(window_size=8) + legacy = FWALOHA(window_size=4) elif legacy_type == 'tdma': legacy = TDMA(state_size=10, assigned_slots=5) else: @@ -136,8 +136,10 @@ def rl_step_fn(c, _): key, init_key = jax.random.split(key) - if traffic_type == 'saturated': - traffic = cox_traffic(f3dB=1.0, loc=0.0, scale=0.0, initial_state=InitialStateConf.ZERO) + if traffic_type == 'constant': + traffic = cox_traffic(f3dB=1.0, loc=-1.0, scale=0.0, initial_state=InitialStateConf.ZERO) + elif traffic_type == 'saturated': + traffic = cox_traffic(f3dB=1.0, loc=2.0, scale=0.0, initial_state=InitialStateConf.ZERO) elif traffic_type == 'bursty': traffic = cox_traffic(f3dB=0.1, loc=-5.0, scale=5.0, initial_state=InitialStateConf.ZERO) else: @@ -146,14 +148,16 @@ def rl_step_fn(c, _): 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_drl)) + rl_step_fn = scan_tqdm(n_steps)(rl_step_fn) init_carry = Carry( drl_states, legacy_states, traffic_states, buffer_states, power_states, channel_state, key, obs, actions, rewards, terminals ) all_outputs = [] - for _ in trange(n_epochs): - carry, output = jax.lax.scan(rl_step_fn, init_carry, length=n_steps) + for epoch in range(n_epochs): + print(f'Epoch {epoch + 1}/{n_epochs}') + carry, output = jax.lax.scan(rl_step_fn, init_carry, jnp.arange(n_steps)) init_carry = replace( init_carry, drl_states=carry.drl_states, diff --git a/pyproject.toml b/pyproject.toml index 93dfc63..5e1b048 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,6 +21,7 @@ authors = [ ] dependencies = [ + "jax-tqdm==0.4.0", "reinforced-lib==2.0.0", "tqdm~=4.67.1", ] From 26c27d77700e0c833a1cf5f26a5294b52693d01b Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Fri, 3 Oct 2025 14:46:04 +0200 Subject: [PATCH 09/17] Fix EB-ALOHA agent --- ltc/agents/eb_aloha.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/ltc/agents/eb_aloha.py b/ltc/agents/eb_aloha.py index 08ce77c..3114a0f 100644 --- a/ltc/agents/eb_aloha.py +++ b/ltc/agents/eb_aloha.py @@ -30,10 +30,14 @@ def init(key, window_size: int): @staticmethod def update(state, key, env_state, action, reward, terminal, window_size: int, max_backoff: int): - *_, buffer = env_state[-1] + collision = (action == Actions.TX.value) & (reward == 0) + success = (action == Actions.TX.value) & (reward > 0) - collision = (action == Actions.TX.value) & (reward == 0) & (buffer > 0) - new_ret_c = jnp.where(collision, jnp.minimum(state.ret_c + 1, max_backoff), 0) + new_ret_c = jnp.where( + collision, + jnp.minimum(state.ret_c + 1, max_backoff), + jnp.where(success, 0, state.ret_c) + ) window = window_size * (2 ** new_ret_c) reset_value = jax.random.randint(key, shape=(), minval=0, maxval=window) From 3c3753b9893760fb7b6ae3f45ed6d10725755e10 Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Mon, 6 Oct 2025 11:23:20 +0200 Subject: [PATCH 10/17] Fix DLMA observation and reward --- ltc/agents/dlma.py | 57 ++++++++++++++++++++++++++++++++++++++- ltc/run.py | 6 ++--- ltc/sim/process_output.py | 35 +----------------------- 3 files changed, 60 insertions(+), 38 deletions(-) diff --git a/ltc/agents/dlma.py b/ltc/agents/dlma.py index 1c017d8..8ddc4ed 100644 --- a/ltc/agents/dlma.py +++ b/ltc/agents/dlma.py @@ -1,10 +1,65 @@ from functools import partial +import jax +import jax.numpy as jnp import flax.linen as nn from ltc.agents.drl import add_batch_dim +@jax.vmap +@jax.vmap +def prepare_observation(obs): + def prepare_action(obs): + return jax.nn.one_hot(obs[1], 2) + + def prepare_state(obs): + """ + **WARNING**: This function assumes that there are only two nodes in the network (1 DLMA and 1 legacy)! + + 1. ch = -1, act = 0 (collision) + 2. ch = 1, act = 0 (DLMA transmits) + 3. ch = 1, act = 1 (legacy transmits) + 4. ch = 0, act = 1 (idle) + + s = max(ch + 2 * act, 0) + + 1. s = max(-1 + 0, 0) = 0 (collision) + 2. s = max(1 + 0, 0) = 1 (DLMA transmits) + 3. s = max(1 + 2, 0) = 3 (legacy transmits) + 4. s = max(0 + 2, 0) = 2 (idle) + """ + + state = obs[:-1] # remove buffer state + state = state[0] + 2 * state[1] + state = jnp.maximum(state, 0) + state = jax.nn.one_hot(state, 4) + return state + + def prepare_rewards(obs): + """ + **WARNING**: This function assumes that there are only two nodes in the network (1 DLMA and 1 legacy)! + + 1. ch = -1, act = 0 (collision) -> DLMA reward = 0, legacy reward = 0 + 2. ch = 1, act = 0 (DLMA transmits) -> DLMA reward = 1, legacy reward = 0 + 3. ch = 1, act = 1 (legacy transmits) -> DLMA reward = 0, legacy reward = 1 + 4. ch = 0, act = 1 (idle) -> DLMA reward = 0, legacy reward = 0 + """ + + state = obs[:-1] # remove buffer state + no_tx = jnp.maximum(state[0], 0) + dlma_tx = no_tx * (1 - state[1]) + legacy_tx = no_tx * state[1] + reward = jnp.array([dlma_tx, legacy_tx]) + return reward + + action = prepare_action(obs) + state = prepare_state(obs) + reward = prepare_rewards(obs) + obs = jnp.concatenate([action, state, reward], axis=0) + return obs + + class DLMANetwork(nn.Module): num_actions: int = 2 @@ -12,8 +67,8 @@ class DLMANetwork(nn.Module): def __call__(self, s): dense = partial(nn.Dense, kernel_init=nn.initializers.he_normal()) - s = s[..., :-1] # remove buffer state s = add_batch_dim(s) + s = prepare_observation(s) b, *_ = s.shape x = s.reshape(b, -1) diff --git a/ltc/run.py b/ltc/run.py index ba4945b..85fdd41 100644 --- a/ltc/run.py +++ b/ltc/run.py @@ -86,11 +86,11 @@ def rl_step_fn(c, _): if __name__ == '__main__': n, n_drl = 2, 1 - n_epochs, n_steps = 1, 500000 + n_epochs, n_steps = 1, 100000 window_size = 20 seed = 42 traffic_type = 'saturated' # 'constant', 'saturated', 'bursty' - legacy_type = 'q-aloha' # 'q-aloha', 'eb-aloha', 'fw-aloha', 'tdma' + legacy_type = 'tdma' # 'q-aloha', 'eb-aloha', 'fw-aloha', 'tdma' key = jax.random.key(seed) actions = jnp.zeros(n, dtype=int) @@ -139,7 +139,7 @@ def rl_step_fn(c, _): if traffic_type == 'constant': traffic = cox_traffic(f3dB=1.0, loc=-1.0, scale=0.0, initial_state=InitialStateConf.ZERO) elif traffic_type == 'saturated': - traffic = cox_traffic(f3dB=1.0, loc=2.0, scale=0.0, initial_state=InitialStateConf.ZERO) + traffic = cox_traffic(f3dB=1.0, loc=5.0, scale=0.0, initial_state=InitialStateConf.ZERO) elif traffic_type == 'bursty': traffic = cox_traffic(f3dB=0.1, loc=-5.0, scale=5.0, initial_state=InitialStateConf.ZERO) else: diff --git a/ltc/sim/process_output.py b/ltc/sim/process_output.py index 2e65c68..b232e6e 100644 --- a/ltc/sim/process_output.py +++ b/ltc/sim/process_output.py @@ -4,43 +4,10 @@ from ltc.sim.constants import * -def no_transmission(args): - _, _, ret_c, _, _ = args - reward = NO_TX_REWARD - ret_c = 0 - no_tx = 0 - return reward, ret_c, no_tx - -def transmission(args): - _, _, _, channel_state, _ = args - return jax.lax.cond(channel_state == 1, successful_transmission, transmission_with_collision, args) - - -def transmission_with_collision(args): - _, _, ret_c, _, _ = args - reward = NO_TX_REWARD - ret_c = 0 - no_tx = 0 - return reward, ret_c, no_tx - -def successful_transmission(args): - _, _, ret_c, _, _ = args - reward = TX_REWARD - ret_c = 0 - no_tx = 0 - return reward, ret_c, no_tx - - def process_output_i(buffer_state, new_buffer_state, power_state, channel_state, obs, action, terminal): - ret_c = 0 - no_tx = 0 - args = (action, buffer_state, ret_c, channel_state, no_tx) - - reward, ret_c, no_tx = jax.lax.cond(action == Actions.TX.value, transmission, no_transmission, args) + reward = channel_state == 1 reward = jnp.where(terminal, 0., reward) - # channel_state = jnp.where(action == Actions.CS.value, channel_state, -1) - power = jnp.where( action == Actions.TX.value, power_state - TX_CONSUMPTION, jnp.where( From fd44a79e7431214b4197579fdea4c0d546149f96 Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Mon, 19 Jan 2026 14:29:34 +0100 Subject: [PATCH 11/17] Add dynamic number of stations --- ltc/run.py | 79 ++++++++++++++++++++++++++-------------- ltc/utils/scan_states.py | 2 + 2 files changed, 54 insertions(+), 27 deletions(-) diff --git a/ltc/run.py b/ltc/run.py index 85fdd41..132ce5d 100644 --- a/ltc/run.py +++ b/ltc/run.py @@ -4,6 +4,7 @@ from dataclasses import replace from functools import partial +import argparse import cloudpickle import jax import jax.numpy as jnp @@ -25,7 +26,7 @@ def init_agents(agent, key, n): return states, step_fn -def agent_step(agent, state, key, obs, action, reward, terminal): +def agent_step(agent, state, key, obs, action, reward, terminal, active): update_key, sample_key = jax.random.split(key) def power_on(state, update_key, sample_key, obs, action, reward, terminal): @@ -37,7 +38,7 @@ def power_off(state, update_key, sample_key, obs, action, reward, terminal): return state, Actions.IDLE.value return jax.lax.cond( - terminal, power_off, power_on, + jnp.logical_or(~active, terminal), power_off, power_on, state, update_key, sample_key, obs, action, reward, terminal ) @@ -49,15 +50,24 @@ def init_traffic(traffic, key, n): return states, step_fn -def rl_step(drl_step, legacy_step, traffic_step, n, n_drl, n_bins=50): - def rl_step_fn(c, _): +def rl_step(drl_step, legacy_step, traffic_step, n, n_drl, n_bins=50, n_switch=None, n_final=None): + def rl_step_fn(c, step): + if n_switch is not None: + active = jnp.where(step == n_switch, jnp.ones_like(c.active).at[n_final:].set(False), c.active) + else: + active = c.active + key, drl_keys, legacy_keys, traffic_key = jax.random.split(c.key, 4) 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) - drl_states, drl_actions = drl_step(c.drl_states, drl_keys, c.obs[:n_drl], c.actions[:n_drl], c.rewards[:n_drl], c.terminals[:n_drl]) - legacy_states, legacy_actions = legacy_step(c.legacy_states, legacy_keys, c.obs[n_drl:], c.actions[n_drl:], c.rewards[n_drl:], c.terminals[n_drl:]) + drl_states, drl_actions = 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( + 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) @@ -73,24 +83,44 @@ def rl_step_fn(c, _): c = Carry( drl_states, legacy_states, traffic_states, buffer_states, powers, - channel_state, key, obs, actions, rewards, terminals + 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, hist, bin_edges + (new_frames > 0).astype(int), channel_state, active, hist, bin_edges ) return c, o return rl_step_fn +def setup_args(): + 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_drl', type=int, default=5, help='Number of DRL agents in the simulation.') + parser.add_argument('--n_final', type=int, help='Final number of agents in the simulation.') + parser.add_argument('--n_steps', type=int, default=2000, help='Number of steps.') + parser.add_argument('--window_size', type=int, default=1, help='Size of the observation window for each agent.') + 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('--traffic_type', type=str, default='saturated', choices=['constant', 'saturated', 'bursty'],help="Traffic model to use: 'constant', 'saturated', or 'bursty'.") + 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'.") + args = parser.parse_args() + return args + + if __name__ == '__main__': - n, n_drl = 2, 1 - n_epochs, n_steps = 1, 100000 - window_size = 20 - seed = 42 - traffic_type = 'saturated' # 'constant', 'saturated', 'bursty' - legacy_type = 'tdma' # 'q-aloha', 'eb-aloha', 'fw-aloha', 'tdma' + args = setup_args() + + n_drl = args.n_drl + n_init = args.n + n_final = args.n_final if args.n_final is not None else n_init + n = max(n_init, n_final) + n_steps = args.n_steps + window_size = args.window_size + seed = args.seed + traffic_type = args.traffic_type + legacy_type = args.legacy_type key = jax.random.key(seed) actions = jnp.zeros(n, dtype=int) @@ -100,6 +130,7 @@ def rl_step_fn(c, _): obs = jnp.zeros((n, window_size, 3), dtype=int) rewards = jnp.zeros(n) terminals = jnp.full(n, False, dtype=bool) + active = jnp.ones(n, dtype=bool).at[n_init:].set(False) drl = DDQN( q_network=DLMANetwork(), @@ -147,7 +178,7 @@ def rl_step_fn(c, _): 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_drl)) + rl_step_fn = jax.jit(rl_step(drl_step, legacy_step, traffic_step, n, n_drl, n_switch=n_steps // 2, n_final=n_final)) rl_step_fn = scan_tqdm(n_steps)(rl_step_fn) init_carry = Carry( drl_states, legacy_states, traffic_states, buffer_states, power_states, @@ -155,21 +186,15 @@ def rl_step_fn(c, _): ) all_outputs = [] - for epoch in range(n_epochs): - print(f'Epoch {epoch + 1}/{n_epochs}') - carry, output = jax.lax.scan(rl_step_fn, init_carry, jnp.arange(n_steps)) - init_carry = replace( - init_carry, - drl_states=carry.drl_states, - key=carry.key - ) - all_outputs.append(output) + carry, output = jax.lax.scan(rl_step_fn, init_carry, jnp.arange(n_steps)) + all_outputs.append(output) all_outputs = jax.tree.map(lambda *x: jnp.stack(x), *all_outputs) - filename = f'history_{n}_{n_drl}_{seed}.pkl.lz4' + filename = f'history_{n_init}_{n_final}_{seed}.pkl.lz4' with lz4.frame.open(filename, 'wb') as f: cloudpickle.dump((init_carry.drl_states, all_outputs), f) - # plot_all(filename) - plot_first(filename) + if args.save_plots: + plot_all(filename) + plot_first(filename) diff --git a/ltc/utils/scan_states.py b/ltc/utils/scan_states.py index f048a3e..51bdf93 100644 --- a/ltc/utils/scan_states.py +++ b/ltc/utils/scan_states.py @@ -20,6 +20,7 @@ class Carry: actions: jax.Array rewards: jax.Array terminals: jax.Array + active: jax.Array @jax.tree_util.register_dataclass @@ -34,5 +35,6 @@ class Output: power_states: jax.Array new_frames: jax.Array channel_state: int + active: jax.Array weights_histogram: jax.Array weights_bin_edges: jax.Array From d1296c0c83b521e0580a1a06afa2c01103b8861b Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Mon, 19 Jan 2026 14:39:36 +0100 Subject: [PATCH 12/17] Fix init carry --- ltc/run.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ltc/run.py b/ltc/run.py index 132ce5d..7051b3f 100644 --- a/ltc/run.py +++ b/ltc/run.py @@ -182,7 +182,7 @@ def setup_args(): rl_step_fn = scan_tqdm(n_steps)(rl_step_fn) init_carry = Carry( drl_states, legacy_states, traffic_states, buffer_states, power_states, - channel_state, key, obs, actions, rewards, terminals + channel_state, key, obs, actions, rewards, terminals, active ) all_outputs = [] From 9a3860a8f27d8407906a6fdc57c6dfc7187dc116 Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Wed, 21 Jan 2026 10:34:15 +0100 Subject: [PATCH 13/17] Fix run script with n_drl=0 --- ltc/run.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/ltc/run.py b/ltc/run.py index 7051b3f..9d14884 100644 --- a/ltc/run.py +++ b/ltc/run.py @@ -75,11 +75,14 @@ def rl_step_fn(c, step): obs, rewards, powers = process_output(c.buffer_states, buffer_states, c.power_states, channel_state, c.obs, actions, c.terminals) terminals = jnp.logical_or(c.terminals, powers < 0) - params = c.drl_states.params['model'] if 'model' in c.drl_states.params else c.drl_states.params - 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) + if n_drl > 0: + params = c.drl_states.params['model'] if 'model' in c.drl_states.params else c.drl_states.params + 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) + else: + hist, bin_edges = None, None c = Carry( drl_states, legacy_states, traffic_states, buffer_states, powers, From 7a9727c19e2bd5dfa0b5fd4f90399935aa685ac6 Mon Sep 17 00:00:00 2001 From: kaszczech Date: Fri, 30 Jan 2026 09:46:23 +0100 Subject: [PATCH 14/17] Add support for traffic generator parameters --- ltc/run.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/ltc/run.py b/ltc/run.py index 9d14884..5be86b0 100644 --- a/ltc/run.py +++ b/ltc/run.py @@ -106,7 +106,10 @@ def setup_args(): parser.add_argument('--window_size', type=int, default=1, help='Size of the observation window for each agent.') 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('--traffic_type', type=str, default='saturated', choices=['constant', 'saturated', 'bursty'],help="Traffic model to use: 'constant', 'saturated', or 'bursty'.") + parser.add_argument('--loc', type=float, default=5.0, help='loc traffic generator parameter.') + parser.add_argument('--scale', type=float, default=0.0, help='scale traffic generator parameter') + parser.add_argument('--f3dB', type=float, default=1.0, help='f3dB traffic generator parameter') + parser.add_argument('--traffic_type', type=str, default='saturated', choices=['constant', 'saturated', 'bursty', 'custom'],help="Traffic model to use: 'constant', 'saturated', or 'bursty'.") 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'.") args = parser.parse_args() return args @@ -125,6 +128,10 @@ def setup_args(): traffic_type = args.traffic_type legacy_type = args.legacy_type + loc = args.loc + scale = args.scale + f3dB = args.f3dB + key = jax.random.key(seed) actions = jnp.zeros(n, dtype=int) buffer_states = jnp.zeros(n, dtype=int) @@ -176,6 +183,8 @@ def setup_args(): traffic = cox_traffic(f3dB=1.0, loc=5.0, scale=0.0, initial_state=InitialStateConf.ZERO) elif 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) else: raise ValueError(f'Unknown traffic type: {traffic_type}') From f0cb43a313e8fa89edb00da985974f70b799c19d Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Sat, 31 Jan 2026 13:20:50 +0100 Subject: [PATCH 15/17] Idle with empty buffer --- ltc/run.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/ltc/run.py b/ltc/run.py index 5be86b0..0b1ca5b 100644 --- a/ltc/run.py +++ b/ltc/run.py @@ -26,7 +26,7 @@ def init_agents(agent, key, n): return states, step_fn -def agent_step(agent, state, key, obs, action, reward, terminal, active): +def agent_step(agent, state, key, obs, action, reward, terminal, buffer_state, active): update_key, sample_key = jax.random.split(key) def power_on(state, update_key, sample_key, obs, action, reward, terminal): @@ -38,7 +38,7 @@ def power_off(state, update_key, sample_key, obs, action, reward, terminal): return state, Actions.IDLE.value return jax.lax.cond( - jnp.logical_or(~active, terminal), power_off, power_on, + jnp.logical_or(jnp.logical_or(~active, terminal), ~buffer_state), power_off, power_on, state, update_key, sample_key, obs, action, reward, terminal ) @@ -63,10 +63,12 @@ def rl_step_fn(c, step): traffic_keys = jax.random.split(traffic_key, n) drl_states, drl_actions = 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], c.buffer_states[:n_drl], active[:n_drl] ) legacy_states, legacy_actions = 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:], c.buffer_states[n_drl:], active[n_drl:] ) actions = jnp.concatenate([drl_actions, legacy_actions]) From ed9d940b8db00fca8c96e77043a968ae3303e4d9 Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Sat, 31 Jan 2026 20:06:54 +0100 Subject: [PATCH 16/17] Fix empty buffer condition --- ltc/run.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ltc/run.py b/ltc/run.py index 0b1ca5b..178d136 100644 --- a/ltc/run.py +++ b/ltc/run.py @@ -38,7 +38,7 @@ def power_off(state, update_key, sample_key, obs, action, reward, terminal): return state, Actions.IDLE.value return jax.lax.cond( - jnp.logical_or(jnp.logical_or(~active, terminal), ~buffer_state), power_off, power_on, + jnp.logical_or(jnp.logical_or(~active, terminal), buffer_state == 0), power_off, power_on, state, update_key, sample_key, obs, action, reward, terminal ) From 13c02c99deae72bb37f3f586cb611e21646fc3b9 Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Mon, 23 Feb 2026 14:44:03 +0100 Subject: [PATCH 17/17] Advanced dynamic scenario --- ltc/run.py | 132 ++++++++++++++++++++++++++++++++++++++++++++++++----- 1 file changed, 121 insertions(+), 11 deletions(-) diff --git a/ltc/run.py b/ltc/run.py index 178d136..8604858 100644 --- a/ltc/run.py +++ b/ltc/run.py @@ -1,7 +1,6 @@ import os os.environ['XLA_PYTHON_CLIENT_PREALLOCATE'] = 'false' -from dataclasses import replace from functools import partial import argparse @@ -50,13 +49,82 @@ def init_traffic(traffic, key, n): return states, step_fn -def rl_step(drl_step, legacy_step, traffic_step, n, n_drl, n_bins=50, n_switch=None, n_final=None): +def active_mask_from_count(count, n): + return jnp.arange(n) < count + + +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, +): + 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) + one_shot_active = active_mask_from_count(target, n) + next_active = jnp.where(step == 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) + should_change = jnp.logical_and( + should_change, (step - station_change_start_step) % 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.where( + station_change_delta > 0, + jnp.minimum(updated_count, target), + jnp.maximum(updated_count, target), + ) + + changed_active = active_mask_from_count(updated_count, n) + next_active = jnp.where(should_change, changed_active, next_active) + + return next_active + + +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, +): def rl_step_fn(c, step): - if n_switch is not None: - active = jnp.where(step == n_switch, jnp.ones_like(c.active).at[n_final:].set(False), c.active) - else: - active = c.active - + 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, + ) + key, drl_keys, legacy_keys, traffic_key = jax.random.split(c.key, 4) drl_keys = jax.random.split(drl_keys, n_drl) legacy_keys = jax.random.split(legacy_keys, n - n_drl) @@ -111,6 +179,10 @@ def setup_args(): parser.add_argument('--loc', type=float, default=5.0, help='loc traffic generator parameter.') parser.add_argument('--scale', type=float, default=0.0, help='scale traffic generator parameter') parser.add_argument('--f3dB', type=float, default=1.0, help='f3dB traffic generator parameter') + parser.add_argument('--station_change_interval', type=int, help='Apply periodic station-count changes every N global steps.') + parser.add_argument('--station_change_delta', type=int, default=0, help='Stations to add/remove per interval tick. Negative values remove stations.') + parser.add_argument('--station_change_start_step', type=int, help='Global step when periodic station changes start. Defaults to interval.') + 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', or 'bursty'.") 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'.") args = parser.parse_args() @@ -122,13 +194,36 @@ def setup_args(): n_drl = args.n_drl n_init = args.n - n_final = args.n_final if args.n_final is not None else n_init + has_n_final = args.n_final is not None + n_final = args.n_final if has_n_final else n_init n = max(n_init, n_final) n_steps = args.n_steps window_size = args.window_size seed = args.seed traffic_type = args.traffic_type legacy_type = args.legacy_type + 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 + + if station_change_interval is not None and 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: + raise ValueError('--station_change_start_step/--station_change_stop_step require --station_change_interval.') + else: + if 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 + + 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: + one_shot_step = n_steps // 2 + one_shot_target = n_final loc = args.loc scale = args.scale @@ -192,7 +287,22 @@ def setup_args(): 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_drl, n_switch=n_steps // 2, n_final=n_final)) + rl_step_fn = jax.jit( + rl_step( + drl_step, + legacy_step, + traffic_step, + n, + n_drl, + 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, + ) + ) rl_step_fn = scan_tqdm(n_steps)(rl_step_fn) init_carry = Carry( drl_states, legacy_states, traffic_states, buffer_states, power_states, @@ -200,14 +310,14 @@ def setup_args(): ) all_outputs = [] - carry, output = jax.lax.scan(rl_step_fn, init_carry, jnp.arange(n_steps)) + carry, output = jax.lax.scan(rl_step_fn, init_carry, jnp.arange(n_steps, dtype=jnp.int32)) all_outputs.append(output) all_outputs = jax.tree.map(lambda *x: jnp.stack(x), *all_outputs) filename = f'history_{n_init}_{n_final}_{seed}.pkl.lz4' with lz4.frame.open(filename, 'wb') as f: - cloudpickle.dump((init_carry.drl_states, all_outputs), f) + cloudpickle.dump((carry.drl_states, all_outputs), f) if args.save_plots: plot_all(filename)