diff --git a/ltc/agents/__init__.py b/ltc/agents/__init__.py index 34a1957..c6f764a 100644 --- a/ltc/agents/__init__.py +++ b/ltc/agents/__init__.py @@ -1,3 +1,8 @@ from ltc.agents.dcf import DCF -from ltc.agents.deep import BayesianDDQN +from ltc.agents.deep import BayesianDDQN, DDQN +from ltc.agents.dlma import DLMANetwork from ltc.agents.drl import QNetwork, QNetworkDropout, StochasticVariationalNetwork +from ltc.agents.eb_aloha import EBALOHA +from ltc.agents.fw_aloha import FWALOHA +from ltc.agents.q_aloha import QALOHA +from ltc.agents.tdma import TDMA diff --git a/ltc/agents/dlma.py b/ltc/agents/dlma.py new file mode 100644 index 0000000..8ddc4ed --- /dev/null +++ b/ltc/agents/dlma.py @@ -0,0 +1,89 @@ +from functools import partial + +import jax +import jax.numpy as jnp +import flax.linen as nn + +from ltc.agents.drl import add_batch_dim + + +@jax.vmap +@jax.vmap +def prepare_observation(obs): + def prepare_action(obs): + return jax.nn.one_hot(obs[1], 2) + + def prepare_state(obs): + """ + **WARNING**: This function assumes that there are only two nodes in the network (1 DLMA and 1 legacy)! + + 1. ch = -1, act = 0 (collision) + 2. ch = 1, act = 0 (DLMA transmits) + 3. ch = 1, act = 1 (legacy transmits) + 4. ch = 0, act = 1 (idle) + + s = max(ch + 2 * act, 0) + + 1. s = max(-1 + 0, 0) = 0 (collision) + 2. s = max(1 + 0, 0) = 1 (DLMA transmits) + 3. s = max(1 + 2, 0) = 3 (legacy transmits) + 4. s = max(0 + 2, 0) = 2 (idle) + """ + + state = obs[:-1] # remove buffer state + state = state[0] + 2 * state[1] + state = jnp.maximum(state, 0) + state = jax.nn.one_hot(state, 4) + return state + + def prepare_rewards(obs): + """ + **WARNING**: This function assumes that there are only two nodes in the network (1 DLMA and 1 legacy)! + + 1. ch = -1, act = 0 (collision) -> DLMA reward = 0, legacy reward = 0 + 2. ch = 1, act = 0 (DLMA transmits) -> DLMA reward = 1, legacy reward = 0 + 3. ch = 1, act = 1 (legacy transmits) -> DLMA reward = 0, legacy reward = 1 + 4. ch = 0, act = 1 (idle) -> DLMA reward = 0, legacy reward = 0 + """ + + state = obs[:-1] # remove buffer state + no_tx = jnp.maximum(state[0], 0) + dlma_tx = no_tx * (1 - state[1]) + legacy_tx = no_tx * state[1] + reward = jnp.array([dlma_tx, legacy_tx]) + return reward + + action = prepare_action(obs) + state = prepare_state(obs) + reward = prepare_rewards(obs) + obs = jnp.concatenate([action, state, reward], axis=0) + return obs + + +class DLMANetwork(nn.Module): + num_actions: int = 2 + + @nn.compact + def __call__(self, s): + dense = partial(nn.Dense, kernel_init=nn.initializers.he_normal()) + + s = add_batch_dim(s) + s = prepare_observation(s) + + b, *_ = s.shape + x = s.reshape(b, -1) + + for _ in range(2): + x = dense(64)(x) + x = nn.relu(x) + + for _ in range(2): + x_res = x + x = dense(64)(x) + x = nn.relu(x) + x = dense(64)(x) + x = nn.relu(x) + x = x + x_res + + x = dense(self.num_actions)(x) + return x diff --git a/ltc/agents/eb_aloha.py b/ltc/agents/eb_aloha.py new file mode 100644 index 0000000..3114a0f --- /dev/null +++ b/ltc/agents/eb_aloha.py @@ -0,0 +1,57 @@ +from functools import partial + +import jax +import jax.numpy as jnp +from chex import dataclass +from reinforced_lib.agents import BaseAgent, AgentState + +from ltc.sim.constants import Actions + + +@dataclass +class EBALOHAState(AgentState): + backoff: int + ret_c: int + + +class EBALOHA(BaseAgent): + def __init__(self, window_size: int, max_backoff: int): + self.window_size = window_size + self.max_backoff = max_backoff + self.init = jax.jit(partial(self.init, window_size=window_size)) + self.update = jax.jit(partial(self.update, window_size=window_size, max_backoff=max_backoff)) + self.sample = jax.jit(self.sample) + + @staticmethod + def init(key, window_size: int): + ret_c = 0 + backoff = jax.random.randint(key, shape=(), minval=0, maxval=window_size) + return EBALOHAState(backoff=backoff, ret_c=ret_c) + + @staticmethod + def update(state, key, env_state, action, reward, terminal, window_size: int, max_backoff: int): + collision = (action == Actions.TX.value) & (reward == 0) + success = (action == Actions.TX.value) & (reward > 0) + + new_ret_c = jnp.where( + collision, + jnp.minimum(state.ret_c + 1, max_backoff), + jnp.where(success, 0, state.ret_c) + ) + window = window_size * (2 ** new_ret_c) + + reset_value = jax.random.randint(key, shape=(), minval=0, maxval=window) + new_backoff = state.backoff - 1 + new_backoff = jnp.where(new_backoff < 0, reset_value, new_backoff) + + return EBALOHAState(backoff=new_backoff, ret_c=new_ret_c) + + @staticmethod + def sample(state, key, env_state): + *_, buffer = env_state[-1] + + return jnp.where( + (buffer > 0) & (state.backoff == 0), + Actions.TX.value, + Actions.IDLE.value + ) diff --git a/ltc/agents/fw_aloha.py b/ltc/agents/fw_aloha.py new file mode 100644 index 0000000..fa338d8 --- /dev/null +++ b/ltc/agents/fw_aloha.py @@ -0,0 +1,43 @@ +from functools import partial + +import jax +import jax.numpy as jnp +from chex import dataclass +from reinforced_lib.agents import BaseAgent, AgentState + +from ltc.sim.constants import Actions + + +@dataclass +class FWALOHAState(AgentState): + backoff: int + + +class FWALOHA(BaseAgent): + def __init__(self, window_size: int): + self.window_size = window_size + self.init = jax.jit(partial(self.init, window_size=window_size)) + self.update = jax.jit(partial(self.update, window_size=window_size)) + self.sample = jax.jit(self.sample) + + @staticmethod + def init(key, window_size: int): + backoff = jax.random.randint(key, shape=(), minval=0, maxval=window_size) + return FWALOHAState(backoff=backoff) + + @staticmethod + def update(state, key, env_state, action, reward, terminal, window_size: int): + reset_value = jax.random.randint(key, shape=(), minval=0, maxval=window_size) + new_backoff = state.backoff - 1 + new_backoff = jnp.where(new_backoff < 0, reset_value, new_backoff) + return FWALOHAState(backoff=new_backoff) + + @staticmethod + def sample(state, key, env_state): + *_, buffer = env_state[-1] + + return jnp.where( + (buffer > 0) & (state.backoff == 0), + Actions.TX.value, + Actions.IDLE.value + ) diff --git a/ltc/agents/q_aloha.py b/ltc/agents/q_aloha.py new file mode 100644 index 0000000..1c422b5 --- /dev/null +++ b/ltc/agents/q_aloha.py @@ -0,0 +1,42 @@ +from functools import partial + +import jax +import jax.numpy as jnp +from chex import dataclass +from reinforced_lib.agents import BaseAgent, AgentState + +from ltc.sim.constants import Actions + + +@dataclass +class QALOHAState(AgentState): + pass + + +class QALOHA(BaseAgent): + def __init__(self, q): + self.init = jax.jit(self.init) + self.update = jax.jit(self.update) + self.sample = jax.jit(partial(self.sample, q=q)) + + @staticmethod + def init(key): + return QALOHAState() + + @staticmethod + def update(state, key, env_state, action, reward, terminal): + return state + + @staticmethod + def sample(state, key, env_state, q): + *_, buffer = env_state[-1] + + return jnp.where( + buffer == 0, + Actions.IDLE.value, + jnp.where( + jax.random.uniform(key) < q, + Actions.TX.value, + Actions.IDLE.value + ) + ) diff --git a/ltc/agents/tdma.py b/ltc/agents/tdma.py new file mode 100644 index 0000000..02c0f02 --- /dev/null +++ b/ltc/agents/tdma.py @@ -0,0 +1,44 @@ +from functools import partial + +import jax +import jax.numpy as jnp +from chex import dataclass +from reinforced_lib.agents import BaseAgent, AgentState + +from ltc.sim.constants import Actions + + +@dataclass +class TDMAState(AgentState): + counter: int + tdma_slots: jnp.ndarray + + +class TDMA(BaseAgent): + def __init__(self, state_size: int, assigned_slots: int): + self.state_size = state_size + self.assigned_slots = assigned_slots + self.init = jax.jit(partial(self.init, state_size=state_size, assigned_slots=assigned_slots)) + self.update = jax.jit(self.update) + self.sample = jax.jit(self.sample) + + @staticmethod + def init(key, state_size: int, assigned_slots: int): + tdma_slots = jnp.zeros(state_size, dtype=int).at[:assigned_slots].set(1) + tdma_slots = jax.random.permutation(key, tdma_slots) + return TDMAState(counter=0, tdma_slots=tdma_slots) + + @staticmethod + def update(state, key, env_state, action, reward, terminal): + new_count = (state.counter + 1) % state.tdma_slots.shape[0] + return TDMAState(counter=new_count, tdma_slots=state.tdma_slots) + + @staticmethod + def sample(state, key, env_state): + *_, buffer = env_state[-1] + + return jnp.where( + (buffer > 0) & (state.tdma_slots[state.counter] == 1), + Actions.TX.value, + Actions.IDLE.value + ) diff --git a/ltc/run.py b/ltc/run.py index 7abb8c2..8604858 100644 --- a/ltc/run.py +++ b/ltc/run.py @@ -1,17 +1,17 @@ import os os.environ['XLA_PYTHON_CLIENT_PREALLOCATE'] = 'false' -from dataclasses import replace from functools import partial +import argparse import cloudpickle import jax 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, QNetworkDropout, StochasticVariationalNetwork +from ltc.agents import BayesianDDQN, DDQN, DLMANetwork, QNetwork, QNetworkDropout, StochasticVariationalNetwork, DCF, QALOHA, EBALOHA, FWALOHA, TDMA from ltc.sim import InitialStateConf, cox_traffic, process_output, simulate from ltc.sim.constants import INITIAL_CAPACITY, Actions from ltc.utils.scan_states import Carry, Output @@ -25,7 +25,7 @@ def init_agents(agent, key, n): return states, step_fn -def agent_step(agent, state, key, obs, action, reward, terminal): +def agent_step(agent, state, key, obs, action, reward, terminal, buffer_state, active): update_key, sample_key = jax.random.split(key) def power_on(state, update_key, sample_key, obs, action, reward, terminal): @@ -37,7 +37,7 @@ def power_off(state, update_key, sample_key, obs, action, reward, terminal): return state, Actions.IDLE.value return jax.lax.cond( - terminal, power_off, power_on, + jnp.logical_or(jnp.logical_or(~active, terminal), buffer_state == 0), power_off, power_on, state, update_key, sample_key, obs, action, reward, terminal ) @@ -49,15 +49,95 @@ def init_traffic(traffic, key, n): return states, step_fn -def rl_step(drl_step, legacy_step, traffic_step, n, n_drl, n_bins=50): - def rl_step_fn(c, _): +def active_mask_from_count(count, n): + return jnp.arange(n) < count + + +def schedule_active_stations( + active, + step, + *, + one_shot_step=None, + one_shot_target=None, + station_change_interval=None, + station_change_delta=0, + station_change_start_step=0, + station_change_stop_step=None, + station_change_target=None, +): + n = active.shape[0] + next_active = active + + if one_shot_step is not None and one_shot_target is not None: + target = jnp.asarray(one_shot_target, dtype=jnp.int32) + one_shot_active = active_mask_from_count(target, n) + next_active = jnp.where(step == one_shot_step, one_shot_active, next_active) + + if station_change_interval is not None and station_change_delta != 0: + should_change = step >= station_change_start_step + if station_change_stop_step is not None: + should_change = jnp.logical_and(should_change, step <= station_change_stop_step) + should_change = jnp.logical_and( + should_change, (step - station_change_start_step) % station_change_interval == 0 + ) + + current_count = jnp.sum(next_active.astype(jnp.int32)) + updated_count = jnp.clip(current_count + station_change_delta, 0, n) + if station_change_target is not None: + target = jnp.asarray(station_change_target, dtype=jnp.int32) + updated_count = jnp.where( + station_change_delta > 0, + jnp.minimum(updated_count, target), + jnp.maximum(updated_count, target), + ) + + changed_active = active_mask_from_count(updated_count, n) + next_active = jnp.where(should_change, changed_active, next_active) + + return next_active + + +def rl_step( + drl_step, + legacy_step, + traffic_step, + n, + n_drl, + n_bins=50, + one_shot_step=None, + one_shot_target=None, + station_change_interval=None, + station_change_delta=0, + station_change_start_step=0, + station_change_stop_step=None, + station_change_target=None, +): + def rl_step_fn(c, step): + active = schedule_active_stations( + c.active, + step, + one_shot_step=one_shot_step, + one_shot_target=one_shot_target, + station_change_interval=station_change_interval, + station_change_delta=station_change_delta, + station_change_start_step=station_change_start_step, + station_change_stop_step=station_change_stop_step, + station_change_target=station_change_target, + ) + key, drl_keys, legacy_keys, traffic_key = jax.random.split(c.key, 4) drl_keys = jax.random.split(drl_keys, n_drl) legacy_keys = jax.random.split(legacy_keys, n - n_drl) 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.buffer_states[:n_drl], active[:n_drl] + ) + legacy_states, legacy_actions = legacy_step( + c.legacy_states, legacy_keys, c.obs[n_drl:], c.actions[n_drl:], c.rewards[n_drl:], + c.terminals[n_drl:], c.buffer_states[n_drl:], active[n_drl:] + ) actions = jnp.concatenate([drl_actions, legacy_actions]) traffic_states, new_frames = traffic_step(c.traffic_states, traffic_keys) @@ -65,62 +145,132 @@ def rl_step_fn(c, _): obs, rewards, powers = process_output(c.buffer_states, buffer_states, c.power_states, channel_state, c.obs, actions, c.terminals) terminals = jnp.logical_or(c.terminals, powers < 0) - params = c.drl_states.params['model'] if 'model' in c.drl_states.params else c.drl_states.params - flat_params, _ = jax.tree.flatten(params) - flat_params = jax.tree.map(lambda x: x.reshape(n_drl, -1), flat_params) - flat_params = jnp.hstack(flat_params) - hist, bin_edges = jax.vmap(jnp.histogram, in_axes=(0, None))(flat_params, n_bins) + if n_drl > 0: + params = c.drl_states.params['model'] if 'model' in c.drl_states.params else c.drl_states.params + flat_params, _ = jax.tree.flatten(params) + flat_params = jax.tree.map(lambda x: x.reshape(n_drl, -1), flat_params) + flat_params = jnp.hstack(flat_params) + hist, bin_edges = jax.vmap(jnp.histogram, in_axes=(0, None))(flat_params, n_bins) + else: + hist, bin_edges = None, None c = Carry( drl_states, legacy_states, traffic_states, buffer_states, powers, - channel_state, key, obs, actions, rewards, terminals + channel_state, key, obs, actions, rewards, terminals, active ) o = Output( legacy_states, obs, actions, rewards, terminals, buffer_states, powers, - (new_frames > 0).astype(int), channel_state, hist, bin_edges + (new_frames > 0).astype(int), channel_state, active, hist, bin_edges ) return c, o return rl_step_fn +def setup_args(): + parser = argparse.ArgumentParser(description="Run the RL network simulation with configurable parameters.") + parser.add_argument('--n', type=int, default=5, help='Initial number of agents in the simulation.') + parser.add_argument('--n_drl', type=int, default=5, help='Number of DRL agents in the simulation.') + parser.add_argument('--n_final', type=int, help='Final number of agents in the simulation.') + parser.add_argument('--n_steps', type=int, default=2000, help='Number of steps.') + parser.add_argument('--window_size', type=int, default=1, help='Size of the observation window for each agent.') + parser.add_argument('--seed', type=int, default=42, help='Random seed for reproducibility.') + parser.add_argument('--save_plots', action='store_true', default=False, help='Whether to save the generated plots.') + parser.add_argument('--loc', type=float, default=5.0, help='loc traffic generator parameter.') + parser.add_argument('--scale', type=float, default=0.0, help='scale traffic generator parameter') + parser.add_argument('--f3dB', type=float, default=1.0, help='f3dB traffic generator parameter') + parser.add_argument('--station_change_interval', type=int, help='Apply periodic station-count changes every N global steps.') + parser.add_argument('--station_change_delta', type=int, default=0, help='Stations to add/remove per interval tick. Negative values remove stations.') + parser.add_argument('--station_change_start_step', type=int, help='Global step when periodic station changes start. Defaults to interval.') + parser.add_argument('--station_change_stop_step', type=int, help='Global step when periodic station changes stop.') + parser.add_argument('--traffic_type', type=str, default='saturated', choices=['constant', 'saturated', 'bursty', 'custom'],help="Traffic model to use: 'constant', 'saturated', or 'bursty'.") + parser.add_argument('--legacy_type', type=str, default='tdma', choices=['q-aloha', 'eb-aloha', 'fw-aloha', 'tdma'], help="Legacy agent type to use: 'q-aloha', 'eb-aloha', 'fw-aloha', or 'tdma'.") + args = parser.parse_args() + return args + + if __name__ == '__main__': - n, n_drl = 10, 5 - n_epochs, n_steps = 50, 10000 - window_size = 5 - seed = 42 - traffic_type = 'saturated' # 'constant', 'saturated', 'bursty' + args = setup_args() + + n_drl = args.n_drl + n_init = args.n + has_n_final = args.n_final is not None + n_final = args.n_final if has_n_final else n_init + n = max(n_init, n_final) + n_steps = args.n_steps + window_size = args.window_size + seed = args.seed + traffic_type = args.traffic_type + legacy_type = args.legacy_type + station_change_interval = args.station_change_interval + station_change_delta = args.station_change_delta + station_change_start_step = args.station_change_start_step + station_change_stop_step = args.station_change_stop_step + + if station_change_interval is not None and station_change_interval <= 0: + raise ValueError('--station_change_interval must be positive.') + if station_change_interval is None: + if station_change_start_step is not None or station_change_stop_step is not None: + raise ValueError('--station_change_start_step/--station_change_stop_step require --station_change_interval.') + else: + if station_change_delta == 0: + raise ValueError('--station_change_delta must be non-zero when --station_change_interval is used.') + if station_change_start_step is None: + station_change_start_step = station_change_interval + + one_shot_step = None + one_shot_target = None + station_change_target = n_final if has_n_final else None + if station_change_interval is None and n_final != n_init: + one_shot_step = n_steps // 2 + one_shot_target = n_final + + loc = args.loc + scale = args.scale + f3dB = args.f3dB key = jax.random.key(seed) actions = jnp.zeros(n, dtype=int) buffer_states = jnp.zeros(n, dtype=int) power_states = jnp.full(n, INITIAL_CAPACITY, dtype=int) channel_state = 0 - obs = jnp.zeros((n, window_size, 5), dtype=int).at[:, -1].set(INITIAL_CAPACITY) + obs = jnp.zeros((n, window_size, 3), dtype=int) rewards = jnp.zeros(n) terminals = jnp.full(n, False, dtype=bool) + active = jnp.ones(n, dtype=bool).at[n_init:].set(False) - drl = BayesianDDQN( - q_network=StochasticVariationalNetwork(QNetworkDropout()), + drl = DDQN( + q_network=DLMANetwork(), obs_space_shape=obs.shape[1:], - act_space_size=3, - optimizer=optax.adam(1e-4), - experience_replay_buffer_size=10000, - experience_replay_batch_size=128, - experience_replay_steps=5, - discount=1.0, - epsilon=1.0, - epsilon_decay=0.999, - epsilon_min=0.001, - tau=0.01 + act_space_size=2, + optimizer=optax.rmsprop(1e-2), + experience_replay_buffer_size=500, + experience_replay_batch_size=32, + experience_replay_steps=1, + discount=0.9, + epsilon=0.1, + epsilon_decay=0.9999, + epsilon_min=0.005, + tau=0.02 ) key, init_key = jax.random.split(key) drl_states, drl_step = init_agents(drl, init_key, n_drl) drl_states = drl_states.replace(prev_env_state=drl_states.prev_env_state.astype(int)) - dcf = DCF() key, init_key = jax.random.split(key) - legacy_states, legacy_step = init_agents(dcf, init_key, n - n_drl) + + if legacy_type == 'q-aloha': + legacy = QALOHA(q=0.5) + elif legacy_type == 'eb-aloha': + legacy = EBALOHA(window_size=4, max_backoff=2) + elif legacy_type == 'fw-aloha': + legacy = FWALOHA(window_size=4) + elif legacy_type == 'tdma': + legacy = TDMA(state_size=10, assigned_slots=5) + else: + raise ValueError(f'Unknown legacy type: {legacy_type}') + + legacy_states, legacy_step = init_agents(legacy, init_key, n - n_drl) key, init_key = jax.random.split(key) @@ -130,32 +280,45 @@ def rl_step_fn(c, _): traffic = cox_traffic(f3dB=1.0, loc=5.0, scale=0.0, initial_state=InitialStateConf.ZERO) elif traffic_type == 'bursty': traffic = cox_traffic(f3dB=0.1, loc=-5.0, scale=5.0, initial_state=InitialStateConf.ZERO) + elif traffic_type == 'custom': + traffic = cox_traffic(f3dB=f3dB, loc=loc, scale=scale, initial_state=InitialStateConf.ZERO) else: raise ValueError(f'Unknown traffic type: {traffic_type}') traffic_states, traffic_step = init_traffic(traffic, init_key, n) - rl_step_fn = jax.jit(rl_step(drl_step, legacy_step, traffic_step, n, n_drl)) + rl_step_fn = jax.jit( + rl_step( + drl_step, + legacy_step, + traffic_step, + n, + n_drl, + one_shot_step=one_shot_step, + one_shot_target=one_shot_target, + station_change_interval=station_change_interval, + station_change_delta=station_change_delta, + station_change_start_step=station_change_start_step, + station_change_stop_step=station_change_stop_step, + station_change_target=station_change_target, + ) + ) + rl_step_fn = scan_tqdm(n_steps)(rl_step_fn) init_carry = Carry( drl_states, legacy_states, traffic_states, buffer_states, power_states, - channel_state, key, obs, actions, rewards, terminals + channel_state, key, obs, actions, rewards, terminals, active ) all_outputs = [] - for _ in trange(n_epochs): - carry, output = jax.lax.scan(rl_step_fn, init_carry, length=n_steps) - init_carry = replace( - init_carry, - drl_states=carry.drl_states, - key=carry.key - ) - all_outputs.append(output) + carry, output = jax.lax.scan(rl_step_fn, init_carry, jnp.arange(n_steps, dtype=jnp.int32)) + all_outputs.append(output) all_outputs = jax.tree.map(lambda *x: jnp.stack(x), *all_outputs) - filename = f'history_{n}_{n_drl}_{seed}.pkl.lz4' + filename = f'history_{n_init}_{n_final}_{seed}.pkl.lz4' with lz4.frame.open(filename, 'wb') as f: - cloudpickle.dump((init_carry.drl_states, all_outputs), f) + cloudpickle.dump((carry.drl_states, all_outputs), f) - plot_all(filename) - plot_first(filename) + if args.save_plots: + plot_all(filename) + plot_first(filename) diff --git a/ltc/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..b232e6e 100644 --- a/ltc/sim/process_output.py +++ b/ltc/sim/process_output.py @@ -4,97 +4,10 @@ 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 = channel_state == 1 reward = jnp.where(terminal, 0., reward) - channel_state = jnp.where(action == Actions.CS.value, channel_state, -1) power = jnp.where( action == Actions.TX.value, power_state - TX_CONSUMPTION, jnp.where( @@ -106,7 +19,7 @@ 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([channel_state, action, new_buffer_state]) obs = jnp.roll(obs, -1, axis=0) obs = obs.at[-1].set(obs_t) diff --git a/ltc/utils/scan_states.py b/ltc/utils/scan_states.py index f048a3e..51bdf93 100644 --- a/ltc/utils/scan_states.py +++ b/ltc/utils/scan_states.py @@ -20,6 +20,7 @@ class Carry: actions: jax.Array rewards: jax.Array terminals: jax.Array + active: jax.Array @jax.tree_util.register_dataclass @@ -34,5 +35,6 @@ class Output: power_states: jax.Array new_frames: jax.Array channel_state: int + active: jax.Array weights_histogram: jax.Array weights_bin_edges: jax.Array diff --git a/pyproject.toml b/pyproject.toml index 93dfc63..5e1b048 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,6 +21,7 @@ authors = [ ] dependencies = [ + "jax-tqdm==0.4.0", "reinforced-lib==2.0.0", "tqdm~=4.67.1", ]