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/__init__.py b/ltc/agents/__init__.py index 98d4e6f..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 -from ltc.agents.model import QNetwork +from ltc.agents.ddqn import BayesianDDQN, DDQN +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 7231f2d..66198d8 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..f7a096c 100644 --- a/ltc/agents/model.py +++ b/ltc/agents/model.py @@ -1,88 +1,87 @@ +import jax import jax.numpy as jnp 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 +def add_batch_dim(x): + return x[None, ...] if x.ndim == 3 else x -class Transformer(nn.Module): - num_layers: int - num_heads: int - ff_dim: int - dropout_rate: float - dtype: jnp.dtype +class QNetwork(nn.Module): + num_actions: int + rnn_dim = 32 + fc_dim = 32 @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) + def __call__(self, s): + scan_gru = nn.scan(nn.GRUCell, variable_broadcast='params', split_rngs={'params': False}, in_axes=1, out_axes=1) + h = scan_gru(self.rnn_dim).initialize_carry(jax.random.PRNGKey(0), s[:, 0].shape) - x = nn.LayerNorm(dtype=self.dtype)(x) - return x + _, 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) + return s -def add_batch_dim(x): - return x[None, ...] if x.ndim == 2 else x +class MixingNetwork(nn.Module): + fc_dim: int = 32 -class QNetwork(nn.Module): + @nn.compact + def __call__(self, qs, g): + b, n = qs.shape + _, g_feat = g.shape + + 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) + 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 + + +class QLBTNetwork(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' @nn.compact - def __call__(self, s, training=True): - x = add_batch_dim(s) - b, t, _ = x.shape - - 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)) - - x = nn.Dense(self.dim, dtype=self.dtype)(x) - x = jnp.concatenate([cls_tokens, x], axis=1) - x = x + pos_embed - - 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 + def __call__(self, s, epsilon=0.0): + 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 / jnp.maximum(d2lt.sum(axis=-1, keepdims=True), 1) + g = jnp.concatenate([actions, d2lt], 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 new file mode 100644 index 0000000..7b0f4ed --- /dev/null +++ b/ltc/agents/qlbt.py @@ -0,0 +1,144 @@ +from functools import partial +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 + + +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, + beta: Scalar = 1.0 + ) -> tuple[Scalar, dict]: + states, _, rewards_tot, terminals, next_states = batch + + rewards_ind = next_states[..., -1, -1] + _, n = rewards_ind.shape + + (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_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 + 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) + + 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(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) + return objective, aux, opt_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 = 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( + 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, + q_network: nn.Module, + act_space_size: int + ) -> int: + dummy_state = jnp.zeros_like(env_state[..., 0])[..., None] + env_state = jnp.concatenate([env_state, dummy_state], 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 e889747..9bd7d75 100644 --- a/ltc/run.py +++ b/ltc/run.py @@ -11,38 +11,36 @@ 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 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): - 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): +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_off(state, update_key, sample_key, obs, action, reward, terminal): - return state, Actions.IDLE.value - - return jax.lax.cond( - terminal, power_off, power_on, - state, update_key, sample_key, obs, action, reward, terminal + 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 state, action + def init_traffic(traffic, key, n): keys = jax.random.split(key, n) @@ -51,38 +49,36 @@ 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) - drl_keys = jax.random.split(drl_keys, n_drl) + 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) - 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 + 1], + 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:], + 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 = 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, throughputs, powers = process_output( + 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) - 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, + 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 @@ -91,13 +87,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('--window_size', type=int, default=20, help='Size of the observation window for each agent.') + parser.add_argument('--n_epochs', type=int, default=1, help='Number of training epochs to run.') + 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_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() @@ -110,36 +106,37 @@ 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) - 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) - drl = BayesianDDQN( - q_network=StochasticVariationalNetwork(QNetwork(num_actions, num_layers=4, dim=64, num_heads=4)), - 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.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) - 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) @@ -155,14 +152,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, throughputs, 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, @@ -177,5 +176,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/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..30532a5 100644 --- a/ltc/sim/process_output.py +++ b/ltc/sim/process_output.py @@ -4,97 +4,16 @@ 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, 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 = 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) - - -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 - 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 - - -def process_output_i(buffer_state, new_buffer_state, power_state, channel_state, obs, action, terminal): - _, _, ret_c, no_tx, _ = obs[-1] - args = (action, buffer_state, ret_c, channel_state, no_tx) - - 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) +def process_output_i(buffer_state, new_buffer_state, power_state, d2lt, idx, channel_state, obs, action, terminal): + *_, 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 + + 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 +25,42 @@ 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, d2lt[idx], ret_c, new_buffer_state]) obs = jnp.roll(obs, -1, axis=0) obs = obs.at[-1].set(obs_t) - return obs, reward, power + return obs, power + +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) + throughputs = throughputs.at[:, -1].set(thr_t) + + 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) + + reward_tot = jnp.where( + channel_state == -1, COLLISION_PENALTY, + jnp.where( + channel_state == 0, NO_TX_REWARD, + jnp.where( + (tx_action & (d2lt == d2lt.max())).sum() > 0, TX_REWARD, + d2lt[jnp.argmax(tx_action)] / jnp.maximum(d2lt.sum(), 1) + ) + ) + ) + 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 + ) -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) + return obs, rewards, throughputs, powers 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/plots.py b/ltc/utils/plots.py index 99975d8..9092cf6 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=100): 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('_') diff --git a/ltc/utils/scan_states.py b/ltc/utils/scan_states.py index f048a3e..b470be5 100644 --- a/ltc/utils/scan_states.py +++ b/ltc/utils/scan_states.py @@ -14,6 +14,8 @@ class Carry: traffic_states: ModelState buffer_states: jax.Array power_states: jax.Array + d2lt: jax.Array + throughputs: 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", ]