From f664871b6ae53bf4ca583db9986f56031bff16b7 Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Wed, 22 Oct 2025 13:10:10 +0200 Subject: [PATCH 01/15] Implement observations --- ltc/agents/__init__.py | 2 +- ltc/agents/dcf.py | 4 ++-- ltc/agents/model.py | 2 +- ltc/run.py | 33 +++++++++++++++++++-------------- ltc/sim/constants.py | 2 +- ltc/sim/process_output.py | 20 +++++++++++++------- ltc/sim/sim.py | 5 +++-- ltc/utils/scan_states.py | 1 + pyproject.toml | 1 + 9 files changed, 42 insertions(+), 28 deletions(-) diff --git a/ltc/agents/__init__.py b/ltc/agents/__init__.py index 98d4e6f..247b9aa 100644 --- a/ltc/agents/__init__.py +++ b/ltc/agents/__init__.py @@ -1,4 +1,4 @@ from ltc.agents.dcf import DCF -from ltc.agents.ddqn import BayesianDDQN +from ltc.agents.ddqn import BayesianDDQN, DDQN from ltc.agents.model import QNetwork from ltc.agents.svi import StochasticVariationalNetwork diff --git a/ltc/agents/dcf.py b/ltc/agents/dcf.py index 7231f2d..cb40729 100644 --- a/ltc/agents/dcf.py +++ b/ltc/agents/dcf.py @@ -44,7 +44,7 @@ def double_cw(): backoff = jax.random.randint(key, (), 0, state.cw) return DCFState(cw=cw, backoff=backoff) - buffer, channel, ret_c, _, _ = env_state[-1] + _, channel, _, _, _, ret_c, buffer = env_state[-1] return jax.lax.cond( buffer == 0, @@ -70,7 +70,7 @@ def double_cw(): @staticmethod def sample(state, key, env_state): - buffer, channel, _, _, _ = env_state[-1] + _, channel, _, _, _, _, buffer = env_state[-1] return jnp.where( buffer == 0, diff --git a/ltc/agents/model.py b/ltc/agents/model.py index 12dfa03..417952f 100644 --- a/ltc/agents/model.py +++ b/ltc/agents/model.py @@ -70,7 +70,7 @@ class QNetwork(nn.Module): @nn.compact def __call__(self, s, training=True): - x = add_batch_dim(s) + x = add_batch_dim(s[..., :-2]) # remove last two features (ret_c and buffer_state) b, t, _ = x.shape pos_embed = self.param('pos_embed', nn.initializers.xavier_uniform(), (1, t + 1, self.dim), self.dtype) diff --git a/ltc/run.py b/ltc/run.py index e889747..0889313 100644 --- a/ltc/run.py +++ b/ltc/run.py @@ -11,9 +11,9 @@ 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, DCF, QNetwork, StochasticVariationalNetwork +from ltc.agents import DDQN, DCF, QNetwork 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 @@ -63,10 +63,12 @@ def rl_step_fn(c, _): 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) - obs, rewards, powers = process_output(c.buffer_states, buffer_states, c.power_states, channel_state, c.obs, actions, c.terminals) + buffer_states, channel_state, d2lt = simulate(c.buffer_states, new_frames, actions, c.d2lt) + obs, rewards, powers = process_output(c.buffer_states, buffer_states, c.power_states, c.d2lt, channel_state, c.obs, actions, c.terminals) terminals = jnp.logical_or(c.terminals, powers < 0) + global_obs = jnp.concatenate([actions, d2lt / d2lt.sum()]) + 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) @@ -77,7 +79,7 @@ def rl_step_fn(c, _): hist, bin_edges = None, None c = Carry( - drl_states, legacy_states, traffic_states, buffer_states, powers, + drl_states, legacy_states, traffic_states, buffer_states, powers, d2lt, channel_state, key, obs, actions, rewards, terminals ) o = Output( @@ -93,8 +95,8 @@ def rl_step_fn(c, _): parser = argparse.ArgumentParser(description="Run the RL network simulation with configurable parameters.") parser.add_argument('--n', type=int, default=10, help='Total number of agents in the simulation.') parser.add_argument('--n_drl', type=int, default=5, help='Number of DRL agents.') - parser.add_argument('--n_epochs', type=int, default=40, 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=2, help='Number of training epochs to run.') + parser.add_argument('--n_steps', type=int, default=200, help='Number of steps per epoch.') parser.add_argument('--window_size', type=int, default=20, 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.') @@ -110,17 +112,18 @@ def rl_step_fn(c, _): traffic_type = args.traffic_type key = jax.random.key(seed) - num_actions = len(Actions) + num_actions = 2 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, 5), dtype=int).at[:, -1].set(INITIAL_CAPACITY) + obs = jnp.zeros((n, window_size, 7), dtype=int) rewards = jnp.zeros(n) terminals = jnp.full(n, False, dtype=bool) + d2lt = jnp.zeros(n, dtype=int) - drl = BayesianDDQN( - q_network=StochasticVariationalNetwork(QNetwork(num_actions, num_layers=4, dim=64, num_heads=4)), + drl = DDQN( + q_network=QNetwork(num_actions, num_layers=2, dim=32, num_heads=2), obs_space_shape=obs.shape[1:], act_space_size=num_actions, optimizer=optax.adam(1e-4), @@ -155,14 +158,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, + drl_states, legacy_states, traffic_states, buffer_states, power_states, d2lt, 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/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 9b4d8da..8e50c0b 100644 --- a/ltc/sim/process_output.py +++ b/ltc/sim/process_output.py @@ -87,14 +87,18 @@ 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): - _, _, ret_c, no_tx, _ = obs[-1] +def process_output_i(buffer_state, new_buffer_state, power_state, d2lt, idx, channel_state, obs, action, terminal): + _, _, no_tx, _, _, ret_c, _ = obs[-1] args = (action, buffer_state, ret_c, channel_state, no_tx) + d2lt_i = d2lt[idx] + d2lt_mi = d2lt.at[idx].set(jnp.inf).min() + d2lt_i, d2lt_mi = d2lt_i / (d2lt_i + d2lt_mi), d2lt_mi / (d2lt_i + d2lt_mi) + 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.TX.value, channel_state == 1, jnp.abs(channel_state)) power = jnp.where( action == Actions.TX.value, power_state - TX_CONSUMPTION, jnp.where( @@ -106,13 +110,15 @@ 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, power]) + obs_t = jnp.array([action, channel_state, 1, d2lt_i, d2lt_mi, ret_c, new_buffer_state]) 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): - channel_states = jnp.full(buffer_states.shape[0], channel_state) - return jax.vmap(process_output_i)(buffer_states, new_buffer_states, power_states, channel_states, obs, actions, terminals) +def process_output(buffer_states, new_buffer_states, power_states, d2lt, channel_state, obs, actions, terminals): + idxs = jnp.arange(buffer_states.shape[0]) + return jax.vmap(process_output_i, in_axes=(0, 0, 0, None, 0, None, 0, 0, 0))( + buffer_states, new_buffer_states, power_states, d2lt, idxs, channel_state, obs, actions, terminals + ) diff --git a/ltc/sim/sim.py b/ltc/sim/sim.py index 369c854..e920b5c 100644 --- a/ltc/sim/sim.py +++ b/ltc/sim/sim.py @@ -46,8 +46,9 @@ def add_new_frames(buffer_states, new_frames): return jnp.bitwise_or(buffer_states, new_frames) -def simulate(buffer_states, new_frames, actions): +def simulate(buffer_states, new_frames, actions, d2lt): 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 + d2lt = jnp.where((channel_state == 1) & (actions == Actions.TX.value), 0, d2lt + 1) + return buffer_states, channel_state, d2lt diff --git a/ltc/utils/scan_states.py b/ltc/utils/scan_states.py index f048a3e..539f498 100644 --- a/ltc/utils/scan_states.py +++ b/ltc/utils/scan_states.py @@ -14,6 +14,7 @@ class Carry: traffic_states: ModelState buffer_states: jax.Array power_states: jax.Array + d2lt: jax.Array channel_state: int key: jax.random.PRNGKey obs: jax.Array diff --git a/pyproject.toml b/pyproject.toml index 93dfc63..14d202d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,6 +30,7 @@ gpu = [ "jax[cuda12-local]", ] notebook = [ + "jax-tqdm==0.4.0", "jupyterlab~=4.4.9", "seaborn~=0.13.2", ] From a3a6fdbd98ed8711df4d01ff225a08db5226667f Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Wed, 22 Oct 2025 13:45:35 +0200 Subject: [PATCH 02/15] Implement actions --- ltc/run.py | 34 +++++++++++++++++++++++----------- 1 file changed, 23 insertions(+), 11 deletions(-) diff --git a/ltc/run.py b/ltc/run.py index 0889313..7f9e704 100644 --- a/ltc/run.py +++ b/ltc/run.py @@ -27,20 +27,26 @@ 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, wait): update_key, sample_key = jax.random.split(key) - def power_on(state, update_key, sample_key, obs, action, reward, terminal): - state = agent.update(state, update_key, obs, action, reward, terminal) - action = agent.sample(state, sample_key, obs) - return state, action + def power_on(state, update_key, sample_key, obs, action, reward, terminal, wait): + def update(state, update_key, sample_key, obs, action, reward, terminal): + state = agent.update(state, update_key, obs, action, reward, terminal) + action = agent.sample(state, sample_key, obs) + return state, action - def power_off(state, update_key, sample_key, obs, action, reward, terminal): + return jax.lax.cond( + wait, lambda state, *_: (state, Actions.CS.value), update, + state, update_key, sample_key, obs, action, reward, terminal + ) + + def power_off(state, update_key, sample_key, obs, action, reward, terminal, wait): return state, Actions.IDLE.value return jax.lax.cond( terminal, power_off, power_on, - state, update_key, sample_key, obs, action, reward, terminal + state, update_key, sample_key, obs, action, reward, terminal, wait ) @@ -58,8 +64,14 @@ def rl_step_fn(c, _): 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], (c.actions[:n_drl] == Actions.TX.value) | (c.channel_state != 0) + ) + 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:], jnp.zeros(n - n_drl, dtype=bool) + ) actions = jnp.concatenate([drl_actions, legacy_actions]) traffic_states, new_frames = traffic_step(c.traffic_states, traffic_keys) @@ -95,8 +107,8 @@ def rl_step_fn(c, _): parser = argparse.ArgumentParser(description="Run the RL network simulation with configurable parameters.") parser.add_argument('--n', type=int, default=10, help='Total number of agents in the simulation.') parser.add_argument('--n_drl', type=int, default=5, help='Number of DRL agents.') - parser.add_argument('--n_epochs', type=int, default=2, help='Number of training epochs to run.') - parser.add_argument('--n_steps', type=int, default=200, help='Number of steps per epoch.') + parser.add_argument('--n_epochs', type=int, default=40, 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=20, 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.') From e24e43f88877597581966567550b13b65b055f2e Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Wed, 22 Oct 2025 14:18:05 +0200 Subject: [PATCH 03/15] Implement QNetwork --- ltc/agents/model.py | 30 +++++++++++------------------- ltc/run.py | 20 ++++++++++---------- 2 files changed, 21 insertions(+), 29 deletions(-) diff --git a/ltc/agents/model.py b/ltc/agents/model.py index 417952f..d8a2dc2 100644 --- a/ltc/agents/model.py +++ b/ltc/agents/model.py @@ -1,3 +1,4 @@ +import jax import jax.numpy as jnp from flax import linen as nn @@ -61,28 +62,19 @@ def add_batch_dim(x): class QNetwork(nn.Module): num_actions: int - num_layers: int - dim: int - num_heads: int - dropout_rate: float = 0.25 - mc_dropout: bool = True - dtype: jnp.dtype = 'float32' + rnn_dim = 32 + fc_dim = 32 @nn.compact def __call__(self, s, training=True): - x = add_batch_dim(s[..., :-2]) # remove last two features (ret_c and buffer_state) - b, t, _ = x.shape + scan_gru = nn.scan(nn.GRUCell, variable_broadcast='params', split_rngs={'params': False}, in_axes=1, out_axes=1) - pos_embed = self.param('pos_embed', nn.initializers.xavier_uniform(), (1, t + 1, self.dim), self.dtype) - cls_token = self.param('cls_token', nn.initializers.xavier_uniform(), (1, 1, self.dim), self.dtype) - cls_tokens = jnp.tile(cls_token, (b, 1, 1)) + s = add_batch_dim(s[..., :-2]) # remove last two features (ret_c and buffer_state) + h = scan_gru(self.rnn_dim).initialize_carry(jax.random.PRNGKey(0), s[:, 0].shape) - x = nn.Dense(self.dim, dtype=self.dtype)(x) - x = jnp.concatenate([cls_tokens, x], axis=1) - x = x + pos_embed + _, s = scan_gru(self.rnn_dim)(h, s) + s = nn.Dense(self.fc_dim)(s[:, -1]) + s = nn.relu(s) + s = nn.Dense(self.num_actions)(s) - x = Transformer(self.num_layers, self.num_heads, 4 * self.dim, self.dropout_rate, self.dtype)(x, training=training) - x = nn.Dropout(self.dropout_rate)(x, deterministic=not training or self.mc_dropout) - x = nn.Dense(self.num_actions, dtype=self.dtype)(x[:, 0]) - - return x + return s diff --git a/ltc/run.py b/ltc/run.py index 7f9e704..8898740 100644 --- a/ltc/run.py +++ b/ltc/run.py @@ -109,7 +109,7 @@ def rl_step_fn(c, _): parser.add_argument('--n_drl', type=int, default=5, help='Number of DRL agents.') parser.add_argument('--n_epochs', type=int, default=40, 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=20, help='Size of the observation window for each agent.') + parser.add_argument('--window_size', type=int, default=10, 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'.") @@ -135,18 +135,18 @@ def rl_step_fn(c, _): d2lt = jnp.zeros(n, dtype=int) drl = DDQN( - q_network=QNetwork(num_actions, num_layers=2, dim=32, num_heads=2), + q_network=QNetwork(num_actions), obs_space_shape=obs.shape[1:], act_space_size=num_actions, - optimizer=optax.adam(1e-4), - experience_replay_buffer_size=10000, - experience_replay_batch_size=128, - experience_replay_steps=5, - discount=1.0, + optimizer=optax.rmsprop(5e-4), + experience_replay_buffer_size=500, + experience_replay_batch_size=32, + experience_replay_steps=1, + discount=0.5, epsilon=1.0, - epsilon_decay=0.999, - epsilon_min=0.001, - tau=0.01 + epsilon_decay=0.998, + epsilon_min=0.01, + tau=0.02 ) key, init_key = jax.random.split(key) drl_states, drl_step = init_agents(drl, init_key, n_drl) From 4b99183cc9afd5b387cb2846ed4b61b97378cf41 Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Wed, 22 Oct 2025 15:24:21 +0200 Subject: [PATCH 04/15] Implement MixingNetwork --- ltc/agents/model.py | 33 +++++++++++++++++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/ltc/agents/model.py b/ltc/agents/model.py index d8a2dc2..7397a30 100644 --- a/ltc/agents/model.py +++ b/ltc/agents/model.py @@ -78,3 +78,36 @@ def __call__(self, s, training=True): s = nn.Dense(self.num_actions)(s) return s + + +class MixingNetwork(nn.Module): + num_actions: int + fc_dim: int = 32 + + @nn.compact + def __call__(self, qs, g): + b, n, _ = qs.shape + _, g_feat = g.shape + + W1 = self.param('W1', nn.initializers.xavier_uniform(), (self.fc_dim * n * self.num_actions, g_feat)) + b1 = self.param('b1', nn.initializers.zeros, (self.fc_dim, g_feat)) + W2 = self.param('W2', nn.initializers.xavier_uniform(), ((n + 1) * self.fc_dim, g_feat)) + b2a = self.param('b2a', nn.initializers.zeros, (self.fc_dim, g_feat)) + b2b = self.param('b2b', nn.initializers.zeros, ((n + 1), self.fc_dim)) + + W1s = jnp.abs(g @ W1.T).reshape(b, self.fc_dim, n * self.num_actions) + b1s = g @ b1.T + W2s = jnp.abs(g @ W2.T).reshape(b, n + 1, self.fc_dim) + b2s = g @ b2a.T + b2s = nn.relu(b2s) + b2s = b2s @ b2b.T + + def fwd(q, W1, b1, W2, b2): + s = q.flatten() + s = jnp.dot(s, W1.T) + b1.T + s = nn.elu(s) + s = jnp.dot(s, W2.T) + b2.T + return s + + s = jax.vmap(fwd)(qs, W1s, b1s, W2s, b2s) + return s From 93a5720927c654f5d2d001f4842d92d9e69c1b6f Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Wed, 22 Oct 2025 16:10:25 +0200 Subject: [PATCH 05/15] Implement rewards --- ltc/run.py | 12 ++--- ltc/sim/process_output.py | 93 ++++++++++++++++----------------------- ltc/utils/scan_states.py | 2 + 3 files changed, 48 insertions(+), 59 deletions(-) diff --git a/ltc/run.py b/ltc/run.py index 8898740..03f03eb 100644 --- a/ltc/run.py +++ b/ltc/run.py @@ -76,11 +76,11 @@ def rl_step_fn(c, _): traffic_states, new_frames = traffic_step(c.traffic_states, traffic_keys) buffer_states, channel_state, d2lt = simulate(c.buffer_states, new_frames, actions, c.d2lt) - obs, rewards, powers = process_output(c.buffer_states, buffer_states, c.power_states, c.d2lt, channel_state, c.obs, actions, c.terminals) + obs, rewards, throughputs, global_obs, powers = process_output( + c.buffer_states, buffer_states, c.power_states, c.d2lt, c.throughputs, channel_state, c.obs, actions, c.terminals + ) terminals = jnp.logical_or(c.terminals, powers < 0) - global_obs = jnp.concatenate([actions, d2lt / d2lt.sum()]) - 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) @@ -91,7 +91,7 @@ def rl_step_fn(c, _): hist, bin_edges = None, None c = Carry( - drl_states, legacy_states, traffic_states, buffer_states, powers, d2lt, + drl_states, legacy_states, traffic_states, buffer_states, powers, d2lt, throughputs, global_obs, channel_state, key, obs, actions, rewards, terminals ) o = Output( @@ -133,6 +133,8 @@ def rl_step_fn(c, _): rewards = jnp.zeros(n) terminals = jnp.full(n, False, dtype=bool) d2lt = jnp.zeros(n, dtype=int) + throughputs = jnp.zeros((n, 1000), dtype=int) + global_obs = jnp.zeros(2 * n, dtype=float) drl = DDQN( q_network=QNetwork(num_actions), @@ -172,7 +174,7 @@ def rl_step_fn(c, _): 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, d2lt, + drl_states, legacy_states, traffic_states, buffer_states, power_states, d2lt, throughputs, global_obs, channel_state, key, obs, actions, rewards, terminals ) all_outputs = [] diff --git a/ltc/sim/process_output.py b/ltc/sim/process_output.py index 8e50c0b..02e221c 100644 --- a/ltc/sim/process_output.py +++ b/ltc/sim/process_output.py @@ -8,83 +8,44 @@ 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), + lambda: jax.lax.cond(buffer_state == 0, zero_counters, increment_no_tx, args), + lambda: jax.lax.cond(no_tx < SAFE_IDLE_PERIOD, no_transmission_short, increment_no_tx, args), ) -def idle_empty_buffer(_): - reward = EMPTY_BUFFER_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 + return ret_c, no_tx -def no_transmission_long(args): +def increment_no_tx(args): _, _, ret_c, _, no_tx = args - scale = jax.lax.min(1., (no_tx - SAFE_IDLE_PERIOD + 1) / PENALIZED_IDLE_PERIOD) - reward = scale * NO_TX_PENALTY no_tx = no_tx + 1 - return reward, ret_c, no_tx + return 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, zero_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) - - -def max_retransmission_collision(_): - reward = MAX_RETRANSMISSION_PENALTY - ret_c = 0 - no_tx = 0 - return reward, ret_c, no_tx + return jax.lax.cond(ret_c < MAX_RETRANSMISSION, retransmission, zero_counters, args) def retransmission(args): _, _, ret_c, _, _ = args - reward = COLLISION_PENALTY ret_c = ret_c + 1 no_tx = 0 - return reward, ret_c, no_tx + return 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 +def zero_counters(_): 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) - ret_c = 0 - no_tx = 0 - return reward, ret_c, no_tx + return ret_c, no_tx def process_output_i(buffer_state, new_buffer_state, power_state, d2lt, idx, channel_state, obs, action, terminal): @@ -95,8 +56,7 @@ def process_output_i(buffer_state, new_buffer_state, power_state, d2lt, idx, cha d2lt_mi = d2lt.at[idx].set(jnp.inf).min() d2lt_i, d2lt_mi = d2lt_i / (d2lt_i + d2lt_mi), d2lt_mi / (d2lt_i + d2lt_mi) - reward, ret_c, no_tx = jax.lax.cond(action == Actions.TX.value, transmission, no_transmission, args) - reward = jnp.where(terminal, 0., reward) + ret_c, no_tx = jax.lax.cond(action == Actions.TX.value, transmission, no_transmission, args) channel_state = jnp.where(action == Actions.TX.value, channel_state == 1, jnp.abs(channel_state)) power = jnp.where( @@ -114,11 +74,36 @@ def process_output_i(buffer_state, new_buffer_state, power_state, d2lt, idx, cha obs = jnp.roll(obs, -1, axis=0) obs = obs.at[-1].set(obs_t) - return obs, reward, power + return obs, power + + +def process_output(buffer_states, new_buffer_states, power_states, d2lt, throughputs, channel_state, obs, actions, terminals): + global_obs = jnp.concatenate([actions, d2lt / d2lt.sum()]) + + thr_t = ((actions == Actions.TX.value) & (channel_state == 1)).astype(int) + throughputs = jnp.roll(throughputs, -1, axis=1) + throughputs = throughputs.at[:, -1].set(thr_t) + priority = new_buffer_states / (throughputs.mean(axis=1) + 1e-6) + opt_action = (priority == priority.max()).astype(int) + rewards_ind = 2 * (actions == opt_action).astype(float) - 1 + rewards_ind = jnp.where(terminals, 0., rewards_ind) + + reward_tot = jnp.where( + channel_state == -1, COLLISION_PENALTY, + jnp.where( + channel_state == 0, NO_TX_REWARD, + jnp.where( + jnp.argmax(actions == Actions.TX.value) == jnp.argmax(d2lt), TX_REWARD, + d2lt[jnp.argmax(actions == Actions.TX.value)] / d2lt.sum() + ) + ) + ) + rewards = jnp.concatenate([rewards_ind, jnp.array([reward_tot])]) -def process_output(buffer_states, new_buffer_states, power_states, d2lt, channel_state, obs, actions, terminals): idxs = jnp.arange(buffer_states.shape[0]) - return jax.vmap(process_output_i, in_axes=(0, 0, 0, None, 0, None, 0, 0, 0))( + obs, powers = jax.vmap(process_output_i, in_axes=(0, 0, 0, None, 0, None, 0, 0, 0))( buffer_states, new_buffer_states, power_states, d2lt, idxs, channel_state, obs, actions, terminals ) + + return obs, rewards, throughputs, global_obs, powers diff --git a/ltc/utils/scan_states.py b/ltc/utils/scan_states.py index 539f498..22de943 100644 --- a/ltc/utils/scan_states.py +++ b/ltc/utils/scan_states.py @@ -15,6 +15,8 @@ class Carry: buffer_states: jax.Array power_states: jax.Array d2lt: jax.Array + throughputs: jax.Array + global_obs: jax.Array channel_state: int key: jax.random.PRNGKey obs: jax.Array From 686b1cc6d0040ccd0e736b66b73035704ef8b05f Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Wed, 22 Oct 2025 20:33:10 +0200 Subject: [PATCH 06/15] Implement QLBT --- ltc/agents/__init__.py | 3 +- ltc/agents/dcf.py | 6 +-- ltc/agents/model.py | 37 ++++++++++--- ltc/agents/qlbt.py | 109 ++++++++++++++++++++++++++++++++++++++ ltc/run.py | 71 ++++++++----------------- ltc/sim/process_output.py | 11 ++-- ltc/utils/scan_states.py | 1 - 7 files changed, 171 insertions(+), 67 deletions(-) create mode 100644 ltc/agents/qlbt.py diff --git a/ltc/agents/__init__.py b/ltc/agents/__init__.py index 247b9aa..b3610e9 100644 --- a/ltc/agents/__init__.py +++ b/ltc/agents/__init__.py @@ -1,4 +1,5 @@ from ltc.agents.dcf import DCF from ltc.agents.ddqn import BayesianDDQN, DDQN -from ltc.agents.model import QNetwork +from ltc.agents.model import QNetwork, QLBTNetwork from ltc.agents.svi import StochasticVariationalNetwork +from ltc.agents.qlbt import QLBT diff --git a/ltc/agents/dcf.py b/ltc/agents/dcf.py index cb40729..31e86a2 100644 --- a/ltc/agents/dcf.py +++ b/ltc/agents/dcf.py @@ -44,7 +44,7 @@ def double_cw(): backoff = jax.random.randint(key, (), 0, state.cw) return DCFState(cw=cw, backoff=backoff) - _, channel, _, _, _, ret_c, buffer = env_state[-1] + _, channel, _, _, _, _, ret_c, buffer = env_state[-1] return jax.lax.cond( buffer == 0, @@ -69,8 +69,8 @@ def double_cw(): ) @staticmethod - def sample(state, key, env_state): - _, channel, _, _, _, _, buffer = env_state[-1] + def sample(state, key, env_state, _): + _, channel, _, _, _, _, _, buffer = env_state[-1] return jnp.where( buffer == 0, diff --git a/ltc/agents/model.py b/ltc/agents/model.py index 7397a30..f94a99f 100644 --- a/ltc/agents/model.py +++ b/ltc/agents/model.py @@ -57,7 +57,7 @@ def __call__(self, x, mask=None, training=True): def add_batch_dim(x): - return x[None, ...] if x.ndim == 2 else x + return x[None, ...] if x.ndim == 3 else x class QNetwork(nn.Module): @@ -66,10 +66,8 @@ class QNetwork(nn.Module): fc_dim = 32 @nn.compact - def __call__(self, s, training=True): + def __call__(self, s): scan_gru = nn.scan(nn.GRUCell, variable_broadcast='params', split_rngs={'params': False}, in_axes=1, out_axes=1) - - s = add_batch_dim(s[..., :-2]) # remove last two features (ret_c and buffer_state) h = scan_gru(self.rnn_dim).initialize_carry(jax.random.PRNGKey(0), s[:, 0].shape) _, s = scan_gru(self.rnn_dim)(h, s) @@ -90,10 +88,10 @@ def __call__(self, qs, g): _, g_feat = g.shape W1 = self.param('W1', nn.initializers.xavier_uniform(), (self.fc_dim * n * self.num_actions, g_feat)) - b1 = self.param('b1', nn.initializers.zeros, (self.fc_dim, g_feat)) + b1 = self.param('b1', nn.zeros_init(), (self.fc_dim, g_feat)) W2 = self.param('W2', nn.initializers.xavier_uniform(), ((n + 1) * self.fc_dim, g_feat)) - b2a = self.param('b2a', nn.initializers.zeros, (self.fc_dim, g_feat)) - b2b = self.param('b2b', nn.initializers.zeros, ((n + 1), self.fc_dim)) + b2a = self.param('b2a', nn.zeros_init(), (self.fc_dim, g_feat)) + b2b = self.param('b2b', nn.zeros_init(), ((n + 1), self.fc_dim)) W1s = jnp.abs(g @ W1.T).reshape(b, self.fc_dim, n * self.num_actions) b1s = g @ b1.T @@ -111,3 +109,28 @@ def fwd(q, W1, b1, W2, b2): s = jax.vmap(fwd)(qs, W1s, b1s, W2s, b2s) return s + + +class QLBTNetwork(nn.Module): + num_actions: int + rnn_dim: int = 32 + fc_dim: int = 32 + + @nn.compact + def __call__(self, s): + BatchQNetwork = nn.vmap( + QNetwork, + in_axes=1, out_axes=1, + variable_axes={'params': 0}, + split_rngs={'params': True} + ) + + s = add_batch_dim(s) + ss = s[..., :-4] # remove auxiliary features (raw d2lt, ret_c, buffer_state, and reward) + actions, d2lt = s[..., -1, 0], s[..., -1, 5] + d2lt = d2lt / (d2lt.sum(axis=-1, keepdims=True) + 1e-6) + g = jnp.concatenate([actions, d2lt], axis=-1) + + qs = BatchQNetwork(self.num_actions)(ss) + q_tot = MixingNetwork(self.num_actions, self.fc_dim)(qs, g) + return jnp.concatenate([q_tot, qs.reshape(s.shape[0], -1)], axis=-1) diff --git a/ltc/agents/qlbt.py b/ltc/agents/qlbt.py new file mode 100644 index 0000000..4663dee --- /dev/null +++ b/ltc/agents/qlbt.py @@ -0,0 +1,109 @@ +from typing import Callable + +import flax.linen as nn +import jax +import jax.numpy as jnp +import optax +from chex import Array, PRNGKey, Scalar +from reinforced_lib.agents.deep.ddqn import DDQN, DDQNState +from reinforced_lib.utils.experience_replay import ExperienceReplay +from reinforced_lib.utils.jax_utils import forward + +from ltc.sim.constants import Actions + + +class QLBT(DDQN): + @staticmethod + def loss_fn( + params: dict, + key: PRNGKey, + state: DDQNState, + batch: tuple, + q_network: nn.Module, + discount: Scalar + ) -> tuple[Scalar, dict]: + states, actions, rewards_tot, terminals, next_states = batch + q_key, q_target_key = jax.random.split(key) + + rewards_ind = next_states[:, :, -1, -1] + b, n = rewards_ind.shape + beta = n + + q_values, net_state = forward(q_network, params, state.net_state, q_key, states) + actions = q_values[..., n + 1:].reshape(b, n, -1) + actions = jnp.argmax(actions, axis=-1) + q_tot, q_ind = q_values[..., 0][..., None], q_values[..., 1:n + 1] + next_states = next_states.at[:, :, -1, 0].set(actions) + + q_values_target, _ = forward(q_network, state.params_target, state.net_state_target, q_target_key, next_states) + q_values_target = q_values_target[..., :n + 1] + q_tot_target, q_ind_target = q_values_target[..., 0], q_values_target[..., 1:n + 1] + q_tot_target = q_tot_target[..., None] + + target_tot = rewards_tot + (1 - terminals) * discount * q_tot_target + target_ind = rewards_ind + (1 - terminals) * discount * q_ind_target + + target_tot = jax.lax.stop_gradient(target_tot) + target_ind = jax.lax.stop_gradient(target_ind) + loss = optax.l2_loss(q_tot, target_tot).mean() + beta * optax.l2_loss(q_ind, target_ind).mean() + + return loss, net_state + + @staticmethod + def update( + state: DDQNState, + key: PRNGKey, + env_state: Array, + actions: Array, + rewards: Scalar, + terminal: Array, + step_fn: Callable, + er: ExperienceReplay, + experience_replay_steps: int, + epsilon_decay: Scalar, + epsilon_min: Scalar, + tau: Scalar + ) -> DDQNState: + filled_rewards = jnp.repeat(rewards[1:].reshape(-1, 1), env_state.shape[1], axis=1)[..., None] + env_state = jnp.concatenate([env_state, filled_rewards], axis=-1) + + replay_buffer = er.append(state.replay_buffer, state.prev_env_state, 0, rewards[0], False, env_state) + batch_key, network_key = jax.random.split(key) + + loss_params = (network_key, state, er.sample(replay_buffer, batch_key)) + params, net_state, opt_state, _ = step_fn(state.params, loss_params, state.opt_state) + params_target, net_state_target = optax.incremental_update((params, net_state), (state.params_target, state.net_state_target), tau) + + return DDQNState( + params=params, + net_state=net_state, + params_target=params_target, + net_state_target=net_state_target, + opt_state=opt_state, + replay_buffer=replay_buffer, + prev_env_state=env_state, + epsilon=jax.lax.max(state.epsilon * epsilon_decay, epsilon_min) + ) + + @staticmethod + def sample( + state: DDQNState, + key: PRNGKey, + env_state: Array, + wait: Array, + q_network: nn.Module, + act_space_size: int + ) -> int: + network_key, action_key = jax.random.split(key) + dummy_state = jnp.zeros_like(env_state[..., 0])[..., None] + env_state = jnp.concatenate([env_state, dummy_state], axis=-1) + + q, _ = forward(q_network, state.params, state.net_state, network_key, env_state) + n = (q.shape[-1] - 1) // (act_space_size + 1) + q = q[0, n + 1:].reshape(n, act_space_size) + + max_q = (q == q.max(axis=-1, keepdims=True)).astype(float) + probs = (1 - state.epsilon) * max_q / jnp.sum(max_q, axis=-1, keepdims=True) + state.epsilon / q.shape[-1] + + action = jax.random.categorical(action_key, jnp.log(probs), axis=-1) + return jnp.where(wait, Actions.CS.value, action) diff --git a/ltc/run.py b/ltc/run.py index 03f03eb..570c97f 100644 --- a/ltc/run.py +++ b/ltc/run.py @@ -13,41 +13,26 @@ import optax from jax_tqdm import scan_tqdm -from ltc.agents import DDQN, DCF, QNetwork +from ltc.agents import DDQN, DCF, QNetwork, QLBT, QLBTNetwork 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 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) - step_fn = jax.vmap(partial(agent_step, agent)) +def init_agents(agent, key, n, apply_vmap): + map_fn = jax.vmap if apply_vmap else lambda f: f + keys = jax.random.split(key, n) if apply_vmap else key + states = map_fn(agent.init)(keys) + step_fn = map_fn(partial(agent_step, agent)) return states, step_fn def agent_step(agent, state, key, obs, action, reward, terminal, wait): update_key, sample_key = jax.random.split(key) - - def power_on(state, update_key, sample_key, obs, action, reward, terminal, wait): - def update(state, update_key, sample_key, obs, action, reward, terminal): - state = agent.update(state, update_key, obs, action, reward, terminal) - action = agent.sample(state, sample_key, obs) - return state, action - - return jax.lax.cond( - wait, lambda state, *_: (state, Actions.CS.value), update, - state, update_key, sample_key, obs, action, reward, terminal - ) - - def power_off(state, update_key, sample_key, obs, action, reward, terminal, wait): - return state, Actions.IDLE.value - - return jax.lax.cond( - terminal, power_off, power_on, - state, update_key, sample_key, obs, action, reward, terminal, wait - ) + state = agent.update(state, update_key, obs, action, reward, terminal) + action = agent.sample(state, sample_key, obs, wait) + return state, action def init_traffic(traffic, key, n): @@ -60,43 +45,33 @@ def init_traffic(traffic, key, n): def rl_step(drl_step, legacy_step, traffic_step, n, n_drl, n_bins=50): def rl_step_fn(c, _): 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.drl_states, drl_keys, c.obs[:n_drl], c.actions[:n_drl], c.rewards[:n_drl + 1], c.terminals[:n_drl], (c.actions[:n_drl] == Actions.TX.value) | (c.channel_state != 0) ) legacy_states, legacy_actions = legacy_step( - c.legacy_states, legacy_keys, c.obs[n_drl:], c.actions[n_drl:], c.rewards[n_drl:], + c.legacy_states, legacy_keys, c.obs[n_drl:], c.actions[n_drl:], c.rewards[n_drl + 1:], c.terminals[n_drl:], jnp.zeros(n - n_drl, dtype=bool) ) actions = jnp.concatenate([drl_actions, legacy_actions]) traffic_states, new_frames = traffic_step(c.traffic_states, traffic_keys) buffer_states, channel_state, d2lt = simulate(c.buffer_states, new_frames, actions, c.d2lt) - obs, rewards, throughputs, global_obs, powers = process_output( + obs, rewards, throughputs, powers = process_output( c.buffer_states, buffer_states, c.power_states, c.d2lt, c.throughputs, channel_state, c.obs, actions, c.terminals ) terminals = jnp.logical_or(c.terminals, powers < 0) - 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, d2lt, throughputs, global_obs, + drl_states, legacy_states, traffic_states, buffer_states, powers, d2lt, throughputs, channel_state, key, obs, actions, rewards, terminals ) 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, None, None ) return c, o @@ -129,16 +104,15 @@ 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, 7), dtype=int) - rewards = jnp.zeros(n) + obs = jnp.zeros((n, window_size, 8), dtype=float) + rewards = jnp.zeros(n + 1) terminals = jnp.full(n, False, dtype=bool) d2lt = jnp.zeros(n, dtype=int) throughputs = jnp.zeros((n, 1000), dtype=int) - global_obs = jnp.zeros(2 * n, dtype=float) - drl = DDQN( - q_network=QNetwork(num_actions), - obs_space_shape=obs.shape[1:], + drl = QLBT( + q_network=QLBTNetwork(num_actions), + obs_space_shape=(n_drl, obs.shape[1], obs.shape[2] + 1), act_space_size=num_actions, optimizer=optax.rmsprop(5e-4), experience_replay_buffer_size=500, @@ -151,12 +125,11 @@ def rl_step_fn(c, _): 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)) + drl_states, drl_step = init_agents(drl, init_key, n_drl, apply_vmap=False) dcf = DCF() key, init_key = jax.random.split(key) - legacy_states, legacy_step = init_agents(dcf, init_key, n - n_drl) + legacy_states, legacy_step = init_agents(dcf, init_key, n - n_drl, apply_vmap=True) key, init_key = jax.random.split(key) @@ -174,7 +147,7 @@ def rl_step_fn(c, _): 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, d2lt, throughputs, global_obs, + drl_states, legacy_states, traffic_states, buffer_states, power_states, d2lt, throughputs, channel_state, key, obs, actions, rewards, terminals ) all_outputs = [] diff --git a/ltc/sim/process_output.py b/ltc/sim/process_output.py index 02e221c..fc44577 100644 --- a/ltc/sim/process_output.py +++ b/ltc/sim/process_output.py @@ -49,7 +49,8 @@ def zero_counters(_): def process_output_i(buffer_state, new_buffer_state, power_state, d2lt, idx, channel_state, obs, action, terminal): - _, _, no_tx, _, _, ret_c, _ = obs[-1] + _, _, no_tx, _, _, _, ret_c, _ = obs[-1] + no_tx, ret_c = no_tx.astype(int), ret_c.astype(int) args = (action, buffer_state, ret_c, channel_state, no_tx) d2lt_i = d2lt[idx] @@ -70,7 +71,7 @@ def process_output_i(buffer_state, new_buffer_state, power_state, d2lt, idx, cha ) ) - obs_t = jnp.array([action, channel_state, 1, d2lt_i, d2lt_mi, ret_c, new_buffer_state]) + obs_t = jnp.array([action, channel_state, 1, d2lt_i, d2lt_mi, d2lt[idx], ret_c, new_buffer_state]) obs = jnp.roll(obs, -1, axis=0) obs = obs.at[-1].set(obs_t) @@ -78,8 +79,6 @@ def process_output_i(buffer_state, new_buffer_state, power_state, d2lt, idx, cha def process_output(buffer_states, new_buffer_states, power_states, d2lt, throughputs, channel_state, obs, actions, terminals): - global_obs = jnp.concatenate([actions, d2lt / d2lt.sum()]) - thr_t = ((actions == Actions.TX.value) & (channel_state == 1)).astype(int) throughputs = jnp.roll(throughputs, -1, axis=1) throughputs = throughputs.at[:, -1].set(thr_t) @@ -99,11 +98,11 @@ def process_output(buffer_states, new_buffer_states, power_states, d2lt, through ) ) ) - rewards = jnp.concatenate([rewards_ind, jnp.array([reward_tot])]) + rewards = jnp.concatenate([jnp.array([reward_tot]), rewards_ind]) idxs = jnp.arange(buffer_states.shape[0]) obs, powers = jax.vmap(process_output_i, in_axes=(0, 0, 0, None, 0, None, 0, 0, 0))( buffer_states, new_buffer_states, power_states, d2lt, idxs, channel_state, obs, actions, terminals ) - return obs, rewards, throughputs, global_obs, powers + return obs, rewards, throughputs, powers diff --git a/ltc/utils/scan_states.py b/ltc/utils/scan_states.py index 22de943..b470be5 100644 --- a/ltc/utils/scan_states.py +++ b/ltc/utils/scan_states.py @@ -16,7 +16,6 @@ class Carry: power_states: jax.Array d2lt: jax.Array throughputs: jax.Array - global_obs: jax.Array channel_state: int key: jax.random.PRNGKey obs: jax.Array From 1ec7c005977674707e8029fbde0f87de2fff00d4 Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Tue, 18 Nov 2025 13:14:42 +0100 Subject: [PATCH 07/15] Refactor QLBT --- .gitignore | 2 ++ ltc/agents/model.py | 53 --------------------------------------------- ltc/agents/qlbt.py | 5 +++-- 3 files changed, 5 insertions(+), 55 deletions(-) diff --git a/.gitignore b/.gitignore index d86e4ff..2773c9e 100644 --- a/.gitignore +++ b/.gitignore @@ -169,3 +169,5 @@ cython_debug/ # PyPI configuration file .pypirc + +.vscode/ diff --git a/ltc/agents/model.py b/ltc/agents/model.py index f94a99f..f6d129c 100644 --- a/ltc/agents/model.py +++ b/ltc/agents/model.py @@ -3,59 +3,6 @@ from flax import linen as nn -class FeedForwardBlock(nn.Module): - ff_dim: int - dropout_rate: float - dtype: jnp.dtype - - @nn.compact - def __call__(self, x, training=True): - *_, out_dim = x.shape - x = nn.Dense(self.ff_dim, dtype=self.dtype)(x) - x = nn.gelu(x) - x = nn.Dropout(self.dropout_rate)(x, deterministic=not training) - x = nn.Dense(out_dim, dtype=self.dtype)(x) - x = nn.Dropout(self.dropout_rate)(x, deterministic=not training) - return x - - -class TransformerBlock(nn.Module): - num_heads: int - ff_dim: int - dropout_rate: float - dtype: jnp.dtype - - @nn.compact - def __call__(self, x, mask, training=True): - residual = x - x = nn.LayerNorm(dtype=self.dtype)(x) - x = nn.MultiHeadDotProductAttention(self.num_heads, qkv_features=x.shape[-1], dtype=self.dtype)(x, mask=mask) - x = x + residual - - residual = x - x = nn.LayerNorm(dtype=self.dtype)(x) - x = FeedForwardBlock(self.ff_dim, self.dropout_rate, self.dtype)(x, training=training) - x = x + residual - - return x - - -class Transformer(nn.Module): - num_layers: int - num_heads: int - ff_dim: int - dropout_rate: float - dtype: jnp.dtype - - @nn.compact - def __call__(self, x, mask=None, training=True): - for _ in range(self.num_layers): - x = TransformerBlock(self.num_heads, self.ff_dim, self.dropout_rate, self.dtype)(x, mask, training=training) - - x = nn.LayerNorm(dtype=self.dtype)(x) - return x - - def add_batch_dim(x): return x[None, ...] if x.ndim == 3 else x diff --git a/ltc/agents/qlbt.py b/ltc/agents/qlbt.py index 4663dee..38f678b 100644 --- a/ltc/agents/qlbt.py +++ b/ltc/agents/qlbt.py @@ -32,13 +32,14 @@ def loss_fn( q_values, net_state = forward(q_network, params, state.net_state, q_key, states) actions = q_values[..., n + 1:].reshape(b, n, -1) actions = jnp.argmax(actions, axis=-1) - q_tot, q_ind = q_values[..., 0][..., None], q_values[..., 1:n + 1] + q_tot, q_ind = q_values[..., 0], q_values[..., 1:n + 1] + q_tot = jnp.expand_dims(q_tot, axis=-1) next_states = next_states.at[:, :, -1, 0].set(actions) q_values_target, _ = forward(q_network, state.params_target, state.net_state_target, q_target_key, next_states) q_values_target = q_values_target[..., :n + 1] q_tot_target, q_ind_target = q_values_target[..., 0], q_values_target[..., 1:n + 1] - q_tot_target = q_tot_target[..., None] + q_tot_target = jnp.expand_dims(q_tot_target, axis=-1) target_tot = rewards_tot + (1 - terminals) * discount * q_tot_target target_ind = rewards_ind + (1 - terminals) * discount * q_ind_target From dee60033ec499b6ae6616775164c6502fa5da1be Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Tue, 18 Nov 2025 14:15:16 +0100 Subject: [PATCH 08/15] Update normalization, adjust simulation parameters --- ltc/agents/model.py | 2 +- ltc/run.py | 9 ++++----- ltc/sim/process_output.py | 5 +++-- ltc/utils/plots.py | 4 +++- 4 files changed, 11 insertions(+), 9 deletions(-) diff --git a/ltc/agents/model.py b/ltc/agents/model.py index f6d129c..50ebe5d 100644 --- a/ltc/agents/model.py +++ b/ltc/agents/model.py @@ -75,7 +75,7 @@ def __call__(self, s): s = add_batch_dim(s) ss = s[..., :-4] # remove auxiliary features (raw d2lt, ret_c, buffer_state, and reward) actions, d2lt = s[..., -1, 0], s[..., -1, 5] - d2lt = d2lt / (d2lt.sum(axis=-1, keepdims=True) + 1e-6) + d2lt = d2lt / jnp.maximum(d2lt.sum(axis=-1, keepdims=True), 1) g = jnp.concatenate([actions, d2lt], axis=-1) qs = BatchQNetwork(self.num_actions)(ss) diff --git a/ltc/run.py b/ltc/run.py index 570c97f..0176238 100644 --- a/ltc/run.py +++ b/ltc/run.py @@ -80,13 +80,13 @@ def rl_step_fn(c, _): if __name__ == '__main__': parser = argparse.ArgumentParser(description="Run the RL network simulation with configurable parameters.") - parser.add_argument('--n', type=int, default=10, help='Total number of agents in the simulation.') + parser.add_argument('--n', type=int, default=5, help='Total number of agents in the simulation.') parser.add_argument('--n_drl', type=int, default=5, help='Number of DRL agents.') - parser.add_argument('--n_epochs', type=int, default=40, 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=1, help='Number of training epochs to run.') + parser.add_argument('--n_steps', type=int, default=100000, help='Number of steps per epoch.') parser.add_argument('--window_size', type=int, default=10, 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('--save-plots', action='store_false', default=True, 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'.") args = parser.parse_args() @@ -169,5 +169,4 @@ def rl_step_fn(c, _): cloudpickle.dump((init_carry.drl_states, all_outputs), f) if args.save_plots: - plot_all(filename) plot_first(filename) diff --git a/ltc/sim/process_output.py b/ltc/sim/process_output.py index fc44577..4110756 100644 --- a/ltc/sim/process_output.py +++ b/ltc/sim/process_output.py @@ -55,7 +55,8 @@ def process_output_i(buffer_state, new_buffer_state, power_state, d2lt, idx, cha d2lt_i = d2lt[idx] d2lt_mi = d2lt.at[idx].set(jnp.inf).min() - d2lt_i, d2lt_mi = d2lt_i / (d2lt_i + d2lt_mi), d2lt_mi / (d2lt_i + d2lt_mi) + denominator = jnp.maximum(d2lt_i + d2lt_mi, 1) + d2lt_i, d2lt_mi = d2lt_i / denominator, d2lt_mi / denominator ret_c, no_tx = jax.lax.cond(action == Actions.TX.value, transmission, no_transmission, args) @@ -94,7 +95,7 @@ def process_output(buffer_states, new_buffer_states, power_states, d2lt, through channel_state == 0, NO_TX_REWARD, jnp.where( jnp.argmax(actions == Actions.TX.value) == jnp.argmax(d2lt), TX_REWARD, - d2lt[jnp.argmax(actions == Actions.TX.value)] / d2lt.sum() + d2lt[jnp.argmax(actions == Actions.TX.value)] / jnp.maximum(d2lt.sum(), 1) ) ) ) diff --git a/ltc/utils/plots.py b/ltc/utils/plots.py index 99975d8..7485f05 100644 --- a/ltc/utils/plots.py +++ b/ltc/utils/plots.py @@ -631,11 +631,13 @@ 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=1000): +def plot_first(filename, n_epochs=1, aggregation=500): with lz4.frame.open(filename, 'rb') as f: _, history = cloudpickle.load(f) history = jax.tree.map(lambda x: x[:n_epochs], asdict(history)) + del history['legacy_states'] history = jax.tree.map(lambda x: x.reshape(-1, aggregation, *x.shape[2:]), history) + history['legacy_states'] = None history = Output(**history) _, n, n_drl, seed_r, *_ = filename.split('_') From e2012f768920d276f565f4b7df2055c5fc6f4231 Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Tue, 18 Nov 2025 16:50:09 +0100 Subject: [PATCH 09/15] Fix QLBT initialization --- ltc/agents/model.py | 13 ++++---- ltc/agents/qlbt.py | 17 +++++------ ltc/run.py | 6 ++-- ltc/sim/process_output.py | 64 ++++++--------------------------------- 4 files changed, 26 insertions(+), 74 deletions(-) diff --git a/ltc/agents/model.py b/ltc/agents/model.py index 50ebe5d..c2680e3 100644 --- a/ltc/agents/model.py +++ b/ltc/agents/model.py @@ -35,10 +35,10 @@ def __call__(self, qs, g): _, g_feat = g.shape W1 = self.param('W1', nn.initializers.xavier_uniform(), (self.fc_dim * n * self.num_actions, g_feat)) - b1 = self.param('b1', nn.zeros_init(), (self.fc_dim, g_feat)) + b1 = self.param('b1', nn.initializers.xavier_uniform(), (self.fc_dim, g_feat)) W2 = self.param('W2', nn.initializers.xavier_uniform(), ((n + 1) * self.fc_dim, g_feat)) - b2a = self.param('b2a', nn.zeros_init(), (self.fc_dim, g_feat)) - b2b = self.param('b2b', nn.zeros_init(), ((n + 1), self.fc_dim)) + b2a = self.param('b2a', nn.initializers.xavier_uniform(), (self.fc_dim, g_feat)) + b2b = self.param('b2b', nn.initializers.xavier_uniform(), ((n + 1), self.fc_dim)) W1s = jnp.abs(g @ W1.T).reshape(b, self.fc_dim, n * self.num_actions) b1s = g @ b1.T @@ -78,6 +78,7 @@ def __call__(self, s): d2lt = d2lt / jnp.maximum(d2lt.sum(axis=-1, keepdims=True), 1) g = jnp.concatenate([actions, d2lt], axis=-1) - qs = BatchQNetwork(self.num_actions)(ss) - q_tot = MixingNetwork(self.num_actions, self.fc_dim)(qs, g) - return jnp.concatenate([q_tot, qs.reshape(s.shape[0], -1)], axis=-1) + q_loc = BatchQNetwork(self.num_actions)(ss) + q_mix = MixingNetwork(self.num_actions, self.fc_dim)(q_loc, g) + q_loc = q_loc.reshape(s.shape[0], -1) + return jnp.concatenate([q_mix, q_loc], axis=-1) diff --git a/ltc/agents/qlbt.py b/ltc/agents/qlbt.py index 38f678b..9f869ec 100644 --- a/ltc/agents/qlbt.py +++ b/ltc/agents/qlbt.py @@ -22,24 +22,21 @@ def loss_fn( q_network: nn.Module, discount: Scalar ) -> tuple[Scalar, dict]: - states, actions, rewards_tot, terminals, next_states = batch + states, _, rewards_tot, terminals, next_states = batch q_key, q_target_key = jax.random.split(key) - rewards_ind = next_states[:, :, -1, -1] + rewards_ind = next_states[..., -1, -1] b, n = rewards_ind.shape beta = n q_values, net_state = forward(q_network, params, state.net_state, q_key, states) - actions = q_values[..., n + 1:].reshape(b, n, -1) - actions = jnp.argmax(actions, axis=-1) - q_tot, q_ind = q_values[..., 0], q_values[..., 1:n + 1] - q_tot = jnp.expand_dims(q_tot, axis=-1) - next_states = next_states.at[:, :, -1, 0].set(actions) + q_tot, q_ind = q_values[..., :1], q_values[..., 1:n + 1] + q_values = q_values[..., n + 1:].reshape(b, n, -1) + new_actions = jnp.argmax(q_values, axis=-1) + next_states = next_states.at[..., -1, 0].set(new_actions) q_values_target, _ = forward(q_network, state.params_target, state.net_state_target, q_target_key, next_states) - q_values_target = q_values_target[..., :n + 1] - q_tot_target, q_ind_target = q_values_target[..., 0], q_values_target[..., 1:n + 1] - q_tot_target = jnp.expand_dims(q_tot_target, axis=-1) + q_tot_target, q_ind_target = q_values_target[..., :1], q_values_target[..., 1:n + 1] target_tot = rewards_tot + (1 - terminals) * discount * q_tot_target target_ind = rewards_ind + (1 - terminals) * discount * q_ind_target diff --git a/ltc/run.py b/ltc/run.py index 0176238..1988e50 100644 --- a/ltc/run.py +++ b/ltc/run.py @@ -13,11 +13,11 @@ import optax from jax_tqdm import scan_tqdm -from ltc.agents import DDQN, DCF, QNetwork, QLBT, QLBTNetwork +from ltc.agents import DCF, QLBT, QLBTNetwork 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 -from ltc.utils.plots import plot_all, plot_first +from ltc.utils.plots import plot_first def init_agents(agent, key, n, apply_vmap): @@ -42,7 +42,7 @@ 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(drl_step, legacy_step, traffic_step, n, n_drl): def rl_step_fn(c, _): key, drl_keys, legacy_keys, traffic_key = jax.random.split(c.key, 4) legacy_keys = jax.random.split(legacy_keys, n - n_drl) diff --git a/ltc/sim/process_output.py b/ltc/sim/process_output.py index 4110756..c362170 100644 --- a/ltc/sim/process_output.py +++ b/ltc/sim/process_output.py @@ -4,62 +4,15 @@ from ltc.sim.constants import * -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, zero_counters, increment_no_tx, args), - lambda: jax.lax.cond(no_tx < SAFE_IDLE_PERIOD, no_transmission_short, increment_no_tx, args), - ) - - -def no_transmission_short(args): - _, buffer_state, ret_c, _, no_tx = args - no_tx = jnp.where(buffer_state == 0, 0, no_tx + 1) - return ret_c, no_tx - - -def increment_no_tx(args): - _, _, ret_c, _, no_tx = args - no_tx = no_tx + 1 - return ret_c, no_tx - - -def transmission(args): - _, _, _, channel_state, _ = args - return jax.lax.cond(channel_state == 1, zero_counters, transmission_with_collision, args) - - -def transmission_with_collision(args): - _, _, ret_c, _, _ = args - return jax.lax.cond(ret_c < MAX_RETRANSMISSION, retransmission, zero_counters, args) - - -def retransmission(args): - _, _, ret_c, _, _ = args - ret_c = ret_c + 1 - no_tx = 0 - return ret_c, no_tx - - -def zero_counters(_): - ret_c = 0 - no_tx = 0 - return ret_c, no_tx - - def process_output_i(buffer_state, new_buffer_state, power_state, d2lt, idx, channel_state, obs, action, terminal): - _, _, no_tx, _, _, _, ret_c, _ = obs[-1] - no_tx, ret_c = no_tx.astype(int), ret_c.astype(int) - args = (action, buffer_state, ret_c, channel_state, no_tx) - + *_, ret_c, _ = obs[-1] + ret_c = jnp.where(action == Actions.TX.value, jnp.where(channel_state == 1, 0, ret_c + 1), ret_c) + d2lt_i = d2lt[idx] d2lt_mi = d2lt.at[idx].set(jnp.inf).min() denominator = jnp.maximum(d2lt_i + d2lt_mi, 1) d2lt_i, d2lt_mi = d2lt_i / denominator, d2lt_mi / denominator - ret_c, no_tx = jax.lax.cond(action == Actions.TX.value, transmission, no_transmission, args) - channel_state = jnp.where(action == Actions.TX.value, channel_state == 1, jnp.abs(channel_state)) power = jnp.where( action == Actions.TX.value, power_state - TX_CONSUMPTION, @@ -80,13 +33,14 @@ def process_output_i(buffer_state, new_buffer_state, power_state, d2lt, idx, cha def process_output(buffer_states, new_buffer_states, power_states, d2lt, throughputs, channel_state, obs, actions, terminals): - thr_t = ((actions == Actions.TX.value) & (channel_state == 1)).astype(int) + tx_action = (actions == Actions.TX.value).astype(int) + thr_t = (tx_action & (channel_state == 1)).astype(int) throughputs = jnp.roll(throughputs, -1, axis=1) throughputs = throughputs.at[:, -1].set(thr_t) - priority = new_buffer_states / (throughputs.mean(axis=1) + 1e-6) + priority = new_buffer_states / jnp.maximum(throughputs.mean(axis=1), 1) opt_action = (priority == priority.max()).astype(int) - rewards_ind = 2 * (actions == opt_action).astype(float) - 1 + rewards_ind = 2 * (tx_action == opt_action).astype(float) - 1 rewards_ind = jnp.where(terminals, 0., rewards_ind) reward_tot = jnp.where( @@ -94,8 +48,8 @@ def process_output(buffer_states, new_buffer_states, power_states, d2lt, through jnp.where( channel_state == 0, NO_TX_REWARD, jnp.where( - jnp.argmax(actions == Actions.TX.value) == jnp.argmax(d2lt), TX_REWARD, - d2lt[jnp.argmax(actions == Actions.TX.value)] / jnp.maximum(d2lt.sum(), 1) + jnp.argmax(tx_action) == jnp.argmax(d2lt), TX_REWARD, + d2lt[jnp.argmax(tx_action)] / jnp.maximum(d2lt.sum(), 1) ) ) ) From aafe322611692626a62475fec8afde38cc5b16ed Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Tue, 18 Nov 2025 18:26:55 +0100 Subject: [PATCH 10/15] Implement stop gradient --- ltc/agents/model.py | 10 ++++---- ltc/agents/qlbt.py | 57 +++++++++++++++++++++++++++++++++++++++++---- 2 files changed, 58 insertions(+), 9 deletions(-) diff --git a/ltc/agents/model.py b/ltc/agents/model.py index c2680e3..33b816f 100644 --- a/ltc/agents/model.py +++ b/ltc/agents/model.py @@ -34,11 +34,11 @@ def __call__(self, qs, g): b, n, _ = qs.shape _, g_feat = g.shape - W1 = self.param('W1', nn.initializers.xavier_uniform(), (self.fc_dim * n * self.num_actions, g_feat)) - b1 = self.param('b1', nn.initializers.xavier_uniform(), (self.fc_dim, g_feat)) - W2 = self.param('W2', nn.initializers.xavier_uniform(), ((n + 1) * self.fc_dim, g_feat)) - b2a = self.param('b2a', nn.initializers.xavier_uniform(), (self.fc_dim, g_feat)) - b2b = self.param('b2b', nn.initializers.xavier_uniform(), ((n + 1), self.fc_dim)) + W1 = self.param('W1', nn.initializers.lecun_normal(), (self.fc_dim * n * self.num_actions, g_feat)) + b1 = self.param('b1', nn.initializers.lecun_normal(), (self.fc_dim, g_feat)) + W2 = self.param('W2', nn.initializers.lecun_normal(), ((n + 1) * self.fc_dim, g_feat)) + b2a = self.param('b2a', nn.initializers.lecun_normal(), (self.fc_dim, g_feat)) + b2b = self.param('b2b', nn.initializers.lecun_normal(), ((n + 1), self.fc_dim)) W1s = jnp.abs(g @ W1.T).reshape(b, self.fc_dim, n * self.num_actions) b1s = g @ b1.T diff --git a/ltc/agents/qlbt.py b/ltc/agents/qlbt.py index 9f869ec..c14b71b 100644 --- a/ltc/agents/qlbt.py +++ b/ltc/agents/qlbt.py @@ -1,3 +1,4 @@ +from functools import partial from typing import Callable import flax.linen as nn @@ -13,21 +14,34 @@ class QLBT(DDQN): + def __init__( + self, + q_network: nn.Module, + *args, + optimizer: optax.GradientTransformation, + discount: Scalar, + **kwargs + ) -> None: + super().__init__(q_network, *args, optimizer=optimizer, discount=discount, **kwargs) + QLBT.loss_fn = jax.jit(partial(QLBT.loss_fn, q_network=q_network, discount=discount)) + QLBT.gradient_step = jax.jit(partial(QLBT.gradient_step, optimizer=optimizer), static_argnames=['n']) + @staticmethod def loss_fn( params: dict, key: PRNGKey, state: DDQNState, batch: tuple, + idx: int, q_network: nn.Module, - discount: Scalar + discount: Scalar, + beta: Scalar = 1.0 ) -> tuple[Scalar, dict]: states, _, rewards_tot, terminals, next_states = batch q_key, q_target_key = jax.random.split(key) rewards_ind = next_states[..., -1, -1] b, n = rewards_ind.shape - beta = n q_values, net_state = forward(q_network, params, state.net_state, q_key, states) q_tot, q_ind = q_values[..., :1], q_values[..., 1:n + 1] @@ -43,9 +57,43 @@ def loss_fn( target_tot = jax.lax.stop_gradient(target_tot) target_ind = jax.lax.stop_gradient(target_ind) - loss = optax.l2_loss(q_tot, target_tot).mean() + beta * optax.l2_loss(q_ind, target_ind).mean() + def scan_fn(i, loss_i): + return i + 1, jnp.where(jax.lax.bitwise_or(idx == i, idx == n), loss_i, jax.lax.stop_gradient(loss_i)) + + loss_tot = optax.l2_loss(q_tot, target_tot).mean() + loss_ind = optax.l2_loss(q_ind, target_ind).mean(axis=0) + _, loss_ind = jax.lax.scan(scan_fn, 0, loss_ind) + + loss = loss_tot + beta * loss_ind.sum() return loss, net_state + + @staticmethod + def combine_grads(grads: dict, aux: any, n: int) -> tuple[dict, any]: + def scan_fn(g, i): + return g, jax.tree.map(lambda x: x[i, i], g) + + aux = jax.tree.map(lambda x: x[-1], aux) + grads_mix = jax.tree.map(lambda x: x[-1], grads['MixingNetwork_0']) + _, grads_q = jax.lax.scan(scan_fn, grads['VmapQNetwork_0'], jnp.arange(n)) + + grads = {'MixingNetwork_0': grads_mix, 'VmapQNetwork_0': grads_q} + return grads, aux + + @staticmethod + def gradient_step( + objective: any, + loss_params: tuple, + opt_state: optax.OptState, + n: int, + optimizer: optax.GradientTransformation + ) -> tuple[any, any, optax.OptState]: + vmaped_loss_fn = jax.vmap(jax.grad(QLBT.loss_fn, has_aux=True), in_axes=(None, None, None, None, 0)) + grads, aux = vmaped_loss_fn(objective, *loss_params, jnp.arange(loss_params[-1][-1].shape[1] + 1)) + grads, aux = QLBT.combine_grads(grads, aux, n) + updates, opt_state = optimizer.update(grads, opt_state, objective) + objective = optax.apply_updates(objective, updates) + return objective, aux, opt_state @staticmethod def update( @@ -64,12 +112,13 @@ def update( ) -> DDQNState: filled_rewards = jnp.repeat(rewards[1:].reshape(-1, 1), env_state.shape[1], axis=1)[..., None] env_state = jnp.concatenate([env_state, filled_rewards], axis=-1) + n = env_state.shape[0] replay_buffer = er.append(state.replay_buffer, state.prev_env_state, 0, rewards[0], False, env_state) batch_key, network_key = jax.random.split(key) loss_params = (network_key, state, er.sample(replay_buffer, batch_key)) - params, net_state, opt_state, _ = step_fn(state.params, loss_params, state.opt_state) + params, net_state, opt_state = QLBT.gradient_step(state.params, loss_params, state.opt_state, n) params_target, net_state_target = optax.incremental_update((params, net_state), (state.params_target, state.net_state_target), tau) return DDQNState( From 0a3adba4ad129977904355b80bad203146626948 Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Wed, 19 Nov 2025 11:33:33 +0100 Subject: [PATCH 11/15] Fix reward and wait --- ltc/agents/dcf.py | 2 +- ltc/agents/qlbt.py | 6 ++---- ltc/run.py | 18 ++++++++++++------ ltc/sim/process_output.py | 4 ++-- 4 files changed, 17 insertions(+), 13 deletions(-) diff --git a/ltc/agents/dcf.py b/ltc/agents/dcf.py index 31e86a2..66198d8 100644 --- a/ltc/agents/dcf.py +++ b/ltc/agents/dcf.py @@ -69,7 +69,7 @@ def double_cw(): ) @staticmethod - def sample(state, key, env_state, _): + def sample(state, key, env_state): _, channel, _, _, _, _, _, buffer = env_state[-1] return jnp.where( diff --git a/ltc/agents/qlbt.py b/ltc/agents/qlbt.py index c14b71b..c7d7f2a 100644 --- a/ltc/agents/qlbt.py +++ b/ltc/agents/qlbt.py @@ -89,7 +89,7 @@ def gradient_step( optimizer: optax.GradientTransformation ) -> tuple[any, any, optax.OptState]: vmaped_loss_fn = jax.vmap(jax.grad(QLBT.loss_fn, has_aux=True), in_axes=(None, None, None, None, 0)) - grads, aux = vmaped_loss_fn(objective, *loss_params, jnp.arange(loss_params[-1][-1].shape[1] + 1)) + grads, aux = vmaped_loss_fn(objective, *loss_params, jnp.arange(n + 1)) grads, aux = QLBT.combine_grads(grads, aux, n) updates, opt_state = optimizer.update(grads, opt_state, objective) objective = optax.apply_updates(objective, updates) @@ -137,7 +137,6 @@ def sample( state: DDQNState, key: PRNGKey, env_state: Array, - wait: Array, q_network: nn.Module, act_space_size: int ) -> int: @@ -152,5 +151,4 @@ def sample( max_q = (q == q.max(axis=-1, keepdims=True)).astype(float) probs = (1 - state.epsilon) * max_q / jnp.sum(max_q, axis=-1, keepdims=True) + state.epsilon / q.shape[-1] - action = jax.random.categorical(action_key, jnp.log(probs), axis=-1) - return jnp.where(wait, Actions.CS.value, action) + return jax.random.categorical(action_key, jnp.log(probs), axis=-1) diff --git a/ltc/run.py b/ltc/run.py index 1988e50..6e1d10e 100644 --- a/ltc/run.py +++ b/ltc/run.py @@ -29,10 +29,16 @@ def init_agents(agent, key, n, apply_vmap): def agent_step(agent, state, key, obs, action, reward, terminal, wait): - update_key, sample_key = jax.random.split(key) - state = agent.update(state, update_key, obs, action, reward, terminal) - action = agent.sample(state, sample_key, obs, wait) - return state, action + def agent_fn(state, key, obs, action, reward, terminal): + update_key, sample_key = jax.random.split(key) + state = agent.update(state, update_key, obs, action, reward, terminal) + action = agent.sample(state, sample_key, obs) + return state, action + + def wait_fn(state, key, obs, action, reward, terminal): + return state, jnp.full_like(action, Actions.CS.value) + + return jax.lax.cond(jnp.any(wait), wait_fn, agent_fn, state, key, obs, action, reward, terminal) def init_traffic(traffic, key, n): @@ -50,7 +56,7 @@ def rl_step_fn(c, _): drl_states, drl_actions = drl_step( c.drl_states, drl_keys, c.obs[:n_drl], c.actions[:n_drl], c.rewards[:n_drl + 1], - c.terminals[:n_drl], (c.actions[:n_drl] == Actions.TX.value) | (c.channel_state != 0) + c.terminals[:n_drl], c.channel_state != 0 ) legacy_states, legacy_actions = legacy_step( c.legacy_states, legacy_keys, c.obs[n_drl:], c.actions[n_drl:], c.rewards[n_drl + 1:], @@ -120,7 +126,7 @@ def rl_step_fn(c, _): experience_replay_steps=1, discount=0.5, epsilon=1.0, - epsilon_decay=0.998, + epsilon_decay=0.9997, epsilon_min=0.01, tau=0.02 ) diff --git a/ltc/sim/process_output.py b/ltc/sim/process_output.py index c362170..4902105 100644 --- a/ltc/sim/process_output.py +++ b/ltc/sim/process_output.py @@ -38,7 +38,7 @@ def process_output(buffer_states, new_buffer_states, power_states, d2lt, through throughputs = jnp.roll(throughputs, -1, axis=1) throughputs = throughputs.at[:, -1].set(thr_t) - priority = new_buffer_states / jnp.maximum(throughputs.mean(axis=1), 1) + priority = new_buffer_states / (throughputs.mean(axis=1) + 1e-6) opt_action = (priority == priority.max()).astype(int) rewards_ind = 2 * (tx_action == opt_action).astype(float) - 1 rewards_ind = jnp.where(terminals, 0., rewards_ind) @@ -48,7 +48,7 @@ def process_output(buffer_states, new_buffer_states, power_states, d2lt, through jnp.where( channel_state == 0, NO_TX_REWARD, jnp.where( - jnp.argmax(tx_action) == jnp.argmax(d2lt), TX_REWARD, + (tx_action & (d2lt == d2lt.max())).sum() > 0, TX_REWARD, d2lt[jnp.argmax(tx_action)] / jnp.maximum(d2lt.sum(), 1) ) ) From c34b738af64576f87fba3da8318771ae9ebc0b26 Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Thu, 20 Nov 2025 21:36:17 +0100 Subject: [PATCH 12/15] Attempt to fix QLBT --- ltc/agents/model.py | 18 ++++++++---------- ltc/agents/qlbt.py | 9 ++------- ltc/run.py | 6 +++--- ltc/utils/plots.py | 2 +- 4 files changed, 14 insertions(+), 21 deletions(-) diff --git a/ltc/agents/model.py b/ltc/agents/model.py index 33b816f..402f033 100644 --- a/ltc/agents/model.py +++ b/ltc/agents/model.py @@ -26,21 +26,20 @@ def __call__(self, s): class MixingNetwork(nn.Module): - num_actions: int fc_dim: int = 32 @nn.compact def __call__(self, qs, g): - b, n, _ = qs.shape + b, n = qs.shape _, g_feat = g.shape - W1 = self.param('W1', nn.initializers.lecun_normal(), (self.fc_dim * n * self.num_actions, g_feat)) + W1 = self.param('W1', nn.initializers.lecun_normal(), (self.fc_dim * n, g_feat)) b1 = self.param('b1', nn.initializers.lecun_normal(), (self.fc_dim, g_feat)) W2 = self.param('W2', nn.initializers.lecun_normal(), ((n + 1) * self.fc_dim, g_feat)) b2a = self.param('b2a', nn.initializers.lecun_normal(), (self.fc_dim, g_feat)) b2b = self.param('b2b', nn.initializers.lecun_normal(), ((n + 1), self.fc_dim)) - W1s = jnp.abs(g @ W1.T).reshape(b, self.fc_dim, n * self.num_actions) + W1s = jnp.abs(g @ W1.T).reshape(b, self.fc_dim, n) b1s = g @ b1.T W2s = jnp.abs(g @ W2.T).reshape(b, n + 1, self.fc_dim) b2s = g @ b2a.T @@ -60,8 +59,6 @@ def fwd(q, W1, b1, W2, b2): class QLBTNetwork(nn.Module): num_actions: int - rnn_dim: int = 32 - fc_dim: int = 32 @nn.compact def __call__(self, s): @@ -78,7 +75,8 @@ def __call__(self, s): d2lt = d2lt / jnp.maximum(d2lt.sum(axis=-1, keepdims=True), 1) g = jnp.concatenate([actions, d2lt], axis=-1) - q_loc = BatchQNetwork(self.num_actions)(ss) - q_mix = MixingNetwork(self.num_actions, self.fc_dim)(q_loc, g) - q_loc = q_loc.reshape(s.shape[0], -1) - return jnp.concatenate([q_mix, q_loc], axis=-1) + q_idn = BatchQNetwork(self.num_actions)(ss) + q_idn_inp = jnp.max(q_idn, axis=-1) + q_idn = q_idn.reshape(s.shape[0], -1) + q_mix = MixingNetwork()(q_idn_inp, g) + return jnp.concatenate([q_mix, q_idn], axis=-1) diff --git a/ltc/agents/qlbt.py b/ltc/agents/qlbt.py index c7d7f2a..35de728 100644 --- a/ltc/agents/qlbt.py +++ b/ltc/agents/qlbt.py @@ -10,8 +10,6 @@ from reinforced_lib.utils.experience_replay import ExperienceReplay from reinforced_lib.utils.jax_utils import forward -from ltc.sim.constants import Actions - class QLBT(DDQN): def __init__( @@ -41,13 +39,10 @@ def loss_fn( q_key, q_target_key = jax.random.split(key) rewards_ind = next_states[..., -1, -1] - b, n = rewards_ind.shape + n = rewards_ind.shape[-1] q_values, net_state = forward(q_network, params, state.net_state, q_key, states) q_tot, q_ind = q_values[..., :1], q_values[..., 1:n + 1] - q_values = q_values[..., n + 1:].reshape(b, n, -1) - new_actions = jnp.argmax(q_values, axis=-1) - next_states = next_states.at[..., -1, 0].set(new_actions) q_values_target, _ = forward(q_network, state.params_target, state.net_state_target, q_target_key, next_states) q_tot_target, q_ind_target = q_values_target[..., :1], q_values_target[..., 1:n + 1] @@ -149,6 +144,6 @@ def sample( q = q[0, n + 1:].reshape(n, act_space_size) max_q = (q == q.max(axis=-1, keepdims=True)).astype(float) - probs = (1 - state.epsilon) * max_q / jnp.sum(max_q, axis=-1, keepdims=True) + state.epsilon / q.shape[-1] + probs = (1 - state.epsilon) * max_q / jnp.sum(max_q, axis=-1, keepdims=True) + state.epsilon / act_space_size return jax.random.categorical(action_key, jnp.log(probs), axis=-1) diff --git a/ltc/run.py b/ltc/run.py index 6e1d10e..4d2b5b1 100644 --- a/ltc/run.py +++ b/ltc/run.py @@ -56,7 +56,7 @@ def rl_step_fn(c, _): drl_states, drl_actions = drl_step( c.drl_states, drl_keys, c.obs[:n_drl], c.actions[:n_drl], c.rewards[:n_drl + 1], - c.terminals[:n_drl], c.channel_state != 0 + c.terminals[:n_drl], False ) legacy_states, legacy_actions = legacy_step( c.legacy_states, legacy_keys, c.obs[n_drl:], c.actions[n_drl:], c.rewards[n_drl + 1:], @@ -89,7 +89,7 @@ def rl_step_fn(c, _): parser.add_argument('--n', type=int, default=5, help='Total number of agents in the simulation.') parser.add_argument('--n_drl', type=int, default=5, help='Number of DRL agents.') parser.add_argument('--n_epochs', type=int, default=1, help='Number of training epochs to run.') - parser.add_argument('--n_steps', type=int, default=100000, help='Number of steps per epoch.') + parser.add_argument('--n_steps', type=int, default=20000, help='Number of steps per epoch.') parser.add_argument('--window_size', type=int, default=10, 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_false', default=True, help='Whether to save the generated plots.') @@ -126,7 +126,7 @@ def rl_step_fn(c, _): experience_replay_steps=1, discount=0.5, epsilon=1.0, - epsilon_decay=0.9997, + epsilon_decay=0.998, epsilon_min=0.01, tau=0.02 ) diff --git a/ltc/utils/plots.py b/ltc/utils/plots.py index 7485f05..9092cf6 100644 --- a/ltc/utils/plots.py +++ b/ltc/utils/plots.py @@ -631,7 +631,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=1, aggregation=500): +def plot_first(filename, n_epochs=1, aggregation=100): with lz4.frame.open(filename, 'rb') as f: _, history = cloudpickle.load(f) history = jax.tree.map(lambda x: x[:n_epochs], asdict(history)) From 6cdf0dccf7533470909d3eb6b7a80a47c7c7620b Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Thu, 20 Nov 2025 22:50:25 +0100 Subject: [PATCH 13/15] Second attempt to fix QLBT --- ltc/agents/model.py | 17 +++++++++++------ ltc/agents/qlbt.py | 25 +++++++++---------------- ltc/run.py | 2 +- 3 files changed, 21 insertions(+), 23 deletions(-) diff --git a/ltc/agents/model.py b/ltc/agents/model.py index 402f033..f7a096c 100644 --- a/ltc/agents/model.py +++ b/ltc/agents/model.py @@ -61,7 +61,7 @@ class QLBTNetwork(nn.Module): num_actions: int @nn.compact - def __call__(self, s): + def __call__(self, s, epsilon=0.0): BatchQNetwork = nn.vmap( QNetwork, in_axes=1, out_axes=1, @@ -75,8 +75,13 @@ def __call__(self, s): d2lt = d2lt / jnp.maximum(d2lt.sum(axis=-1, keepdims=True), 1) g = jnp.concatenate([actions, d2lt], axis=-1) - q_idn = BatchQNetwork(self.num_actions)(ss) - q_idn_inp = jnp.max(q_idn, axis=-1) - q_idn = q_idn.reshape(s.shape[0], -1) - q_mix = MixingNetwork()(q_idn_inp, g) - return jnp.concatenate([q_mix, q_idn], axis=-1) + q_ind = BatchQNetwork(self.num_actions)(ss) + max_q = (q_ind == q_ind.max(axis=-1, keepdims=True)).astype(float) + probs = (1 - epsilon) * max_q / jnp.sum(max_q, axis=-1, keepdims=True) + epsilon / self.num_actions + acts = jax.random.categorical(self.make_rng('rlib'), jnp.log(probs), axis=-1) + q_ind_inp = jnp.take_along_axis(q_ind, acts[..., None], axis=-1).squeeze(-1) + + q_mix = MixingNetwork()(q_ind_inp, g) + q_vals = jnp.concatenate([q_mix, q_ind.reshape(s.shape[0], -1)], axis=-1) + + return q_vals, acts diff --git a/ltc/agents/qlbt.py b/ltc/agents/qlbt.py index 35de728..eecd65a 100644 --- a/ltc/agents/qlbt.py +++ b/ltc/agents/qlbt.py @@ -36,15 +36,17 @@ def loss_fn( beta: Scalar = 1.0 ) -> tuple[Scalar, dict]: states, _, rewards_tot, terminals, next_states = batch - q_key, q_target_key = jax.random.split(key) rewards_ind = next_states[..., -1, -1] - n = rewards_ind.shape[-1] + b, n = rewards_ind.shape - q_values, net_state = forward(q_network, params, state.net_state, q_key, states) + (q_values, _), net_state = forward(q_network, params, state.net_state, key, states, state.epsilon) q_tot, q_ind = q_values[..., :1], q_values[..., 1:n + 1] + q_values = q_values[..., n + 1:].reshape(b, n, -1) + new_actions = jnp.argmax(q_values, axis=-1) + next_states = next_states.at[..., -1, 0].set(new_actions) - q_values_target, _ = forward(q_network, state.params_target, state.net_state_target, q_target_key, next_states) + (q_values_target, _), _ = forward(q_network, state.params_target, state.net_state_target, key, next_states) q_tot_target, q_ind_target = q_values_target[..., :1], q_values_target[..., 1:n + 1] target_tot = rewards_tot + (1 - terminals) * discount * q_tot_target @@ -107,13 +109,12 @@ def update( ) -> DDQNState: filled_rewards = jnp.repeat(rewards[1:].reshape(-1, 1), env_state.shape[1], axis=1)[..., None] env_state = jnp.concatenate([env_state, filled_rewards], axis=-1) - n = env_state.shape[0] replay_buffer = er.append(state.replay_buffer, state.prev_env_state, 0, rewards[0], False, env_state) batch_key, network_key = jax.random.split(key) loss_params = (network_key, state, er.sample(replay_buffer, batch_key)) - params, net_state, opt_state = QLBT.gradient_step(state.params, loss_params, state.opt_state, n) + params, net_state, opt_state = QLBT.gradient_step(state.params, loss_params, state.opt_state, env_state.shape[0]) params_target, net_state_target = optax.incremental_update((params, net_state), (state.params_target, state.net_state_target), tau) return DDQNState( @@ -135,15 +136,7 @@ def sample( q_network: nn.Module, act_space_size: int ) -> int: - network_key, action_key = jax.random.split(key) dummy_state = jnp.zeros_like(env_state[..., 0])[..., None] env_state = jnp.concatenate([env_state, dummy_state], axis=-1) - - q, _ = forward(q_network, state.params, state.net_state, network_key, env_state) - n = (q.shape[-1] - 1) // (act_space_size + 1) - q = q[0, n + 1:].reshape(n, act_space_size) - - max_q = (q == q.max(axis=-1, keepdims=True)).astype(float) - probs = (1 - state.epsilon) * max_q / jnp.sum(max_q, axis=-1, keepdims=True) + state.epsilon / act_space_size - - return jax.random.categorical(action_key, jnp.log(probs), axis=-1) + (_, a), _ = forward(q_network, state.params, state.net_state, key, env_state, state.epsilon) + return a[0] diff --git a/ltc/run.py b/ltc/run.py index 4d2b5b1..73fb4a8 100644 --- a/ltc/run.py +++ b/ltc/run.py @@ -38,7 +38,7 @@ def agent_fn(state, key, obs, action, reward, terminal): def wait_fn(state, key, obs, action, reward, terminal): return state, jnp.full_like(action, Actions.CS.value) - return jax.lax.cond(jnp.any(wait), wait_fn, agent_fn, state, key, obs, action, reward, terminal) + return jax.lax.cond(wait, wait_fn, agent_fn, state, key, obs, action, reward, terminal) def init_traffic(traffic, key, n): From 7535489b875b213818603bc2bcedc554dba582ef Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Thu, 20 Nov 2025 23:03:40 +0100 Subject: [PATCH 14/15] Third attempt to fix QLBT --- ltc/run.py | 4 ++-- ltc/sim/process_output.py | 5 ++++- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/ltc/run.py b/ltc/run.py index 73fb4a8..de08fff 100644 --- a/ltc/run.py +++ b/ltc/run.py @@ -50,7 +50,7 @@ def init_traffic(traffic, key, n): def rl_step(drl_step, legacy_step, traffic_step, n, n_drl): def rl_step_fn(c, _): - key, drl_keys, legacy_keys, traffic_key = jax.random.split(c.key, 4) + key, drl_keys, legacy_keys, traffic_key, process_key = jax.random.split(c.key, 5) legacy_keys = jax.random.split(legacy_keys, n - n_drl) traffic_keys = jax.random.split(traffic_key, n) @@ -67,7 +67,7 @@ def rl_step_fn(c, _): traffic_states, new_frames = traffic_step(c.traffic_states, traffic_keys) buffer_states, channel_state, d2lt = simulate(c.buffer_states, new_frames, actions, c.d2lt) obs, rewards, throughputs, powers = process_output( - c.buffer_states, buffer_states, c.power_states, c.d2lt, c.throughputs, channel_state, c.obs, actions, c.terminals + process_key, c.buffer_states, buffer_states, c.power_states, c.d2lt, c.throughputs, channel_state, c.obs, actions, c.terminals ) terminals = jnp.logical_or(c.terminals, powers < 0) diff --git a/ltc/sim/process_output.py b/ltc/sim/process_output.py index 4902105..30532a5 100644 --- a/ltc/sim/process_output.py +++ b/ltc/sim/process_output.py @@ -32,7 +32,7 @@ def process_output_i(buffer_state, new_buffer_state, power_state, d2lt, idx, cha return obs, power -def process_output(buffer_states, new_buffer_states, power_states, d2lt, throughputs, channel_state, obs, actions, terminals): +def process_output(key, buffer_states, new_buffer_states, power_states, d2lt, throughputs, channel_state, obs, actions, terminals): tx_action = (actions == Actions.TX.value).astype(int) thr_t = (tx_action & (channel_state == 1)).astype(int) throughputs = jnp.roll(throughputs, -1, axis=1) @@ -40,6 +40,9 @@ def process_output(buffer_states, new_buffer_states, power_states, d2lt, through priority = new_buffer_states / (throughputs.mean(axis=1) + 1e-6) opt_action = (priority == priority.max()).astype(int) + probs = opt_action / opt_action.sum() + opt_action = jax.random.choice(key, actions.shape[0], p=probs) + opt_action = jnp.zeros_like(actions).at[opt_action].set(1) rewards_ind = 2 * (tx_action == opt_action).astype(float) - 1 rewards_ind = jnp.where(terminals, 0., rewards_ind) From 7a21377dd0ab47201036604c4b45a3e529789cbe Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Tue, 25 Nov 2025 09:35:38 +0100 Subject: [PATCH 15/15] Fourth attempt to fix QLBT --- ltc/agents/qlbt.py | 14 ++++++++------ ltc/run.py | 21 +++++++++++---------- 2 files changed, 19 insertions(+), 16 deletions(-) diff --git a/ltc/agents/qlbt.py b/ltc/agents/qlbt.py index eecd65a..7b0f4ed 100644 --- a/ltc/agents/qlbt.py +++ b/ltc/agents/qlbt.py @@ -38,13 +38,10 @@ def loss_fn( states, _, rewards_tot, terminals, next_states = batch rewards_ind = next_states[..., -1, -1] - b, n = rewards_ind.shape + _, n = rewards_ind.shape - (q_values, _), net_state = forward(q_network, params, state.net_state, key, states, state.epsilon) + (q_values, _), net_state = forward(q_network, params, state.net_state, key, states) q_tot, q_ind = q_values[..., :1], q_values[..., 1:n + 1] - q_values = q_values[..., n + 1:].reshape(b, n, -1) - new_actions = jnp.argmax(q_values, axis=-1) - next_states = next_states.at[..., -1, 0].set(new_actions) (q_values_target, _), _ = forward(q_network, state.params_target, state.net_state_target, key, next_states) q_tot_target, q_ind_target = q_values_target[..., :1], q_values_target[..., 1:n + 1] @@ -114,7 +111,12 @@ def update( batch_key, network_key = jax.random.split(key) loss_params = (network_key, state, er.sample(replay_buffer, batch_key)) - params, net_state, opt_state = QLBT.gradient_step(state.params, loss_params, state.opt_state, env_state.shape[0]) + params, net_state, opt_state = jax.lax.cond( + er.is_ready(replay_buffer), + lambda p, lp, os: QLBT.gradient_step(p, lp, os, env_state.shape[0]), + lambda p, lp, os: (p, lp[1].net_state, os), + state.params, loss_params, state.opt_state + ) params_target, net_state_target = optax.incremental_update((params, net_state), (state.params_target, state.net_state_target), tau) return DDQNState( diff --git a/ltc/run.py b/ltc/run.py index de08fff..9bd7d75 100644 --- a/ltc/run.py +++ b/ltc/run.py @@ -29,16 +29,17 @@ def init_agents(agent, key, n, apply_vmap): def agent_step(agent, state, key, obs, action, reward, terminal, wait): - def agent_fn(state, key, obs, action, reward, terminal): - update_key, sample_key = jax.random.split(key) - state = agent.update(state, update_key, obs, action, reward, terminal) - action = agent.sample(state, sample_key, obs) - return state, action - - def wait_fn(state, key, obs, action, reward, terminal): - return state, jnp.full_like(action, Actions.CS.value) + update_key, sample_key = jax.random.split(key) + + state = agent.update(state, update_key, obs, action, reward, terminal) + action = jax.lax.cond( + wait, + lambda *_: jnp.full_like(action, Actions.CS.value), + lambda state, key, obs: agent.sample(state, key, obs), + state, sample_key, obs + ) - return jax.lax.cond(wait, wait_fn, agent_fn, state, key, obs, action, reward, terminal) + return state, action def init_traffic(traffic, key, n): @@ -56,7 +57,7 @@ def rl_step_fn(c, _): drl_states, drl_actions = drl_step( c.drl_states, drl_keys, c.obs[:n_drl], c.actions[:n_drl], c.rewards[:n_drl + 1], - c.terminals[:n_drl], False + c.terminals[:n_drl], c.channel_state != 0 ) legacy_states, legacy_actions = legacy_step( c.legacy_states, legacy_keys, c.obs[n_drl:], c.actions[n_drl:], c.rewards[n_drl + 1:],