diff --git a/.gitignore b/.gitignore index d86e4ff..290b7ae 100644 --- a/.gitignore +++ b/.gitignore @@ -169,3 +169,6 @@ cython_debug/ # PyPI configuration file .pypirc + +scripts/ +plots/ \ No newline at end of file diff --git a/ltc/agents/dcf.py b/ltc/agents/dcf.py index 7231f2d..96ee116 100644 --- a/ltc/agents/dcf.py +++ b/ltc/agents/dcf.py @@ -3,7 +3,7 @@ from chex import dataclass from reinforced_lib.agents import BaseAgent, AgentState -from ltc.sim.constants import Actions +from ltc.sim.constants import Actions, ObsIdx @dataclass @@ -44,22 +44,28 @@ def double_cw(): backoff = jax.random.randint(key, (), 0, state.cw) return DCFState(cw=cw, backoff=backoff) - buffer, channel, ret_c, _, _ = env_state[-1] + buffer_count = env_state[-1][ObsIdx.BUFFER_PACKET_COUNT] + channel_busy = env_state[-1][ObsIdx.CHANNEL_LAST_CS_BUSY] + ret_c = env_state[-1][ObsIdx.STATUS_RETRY_COUNTER] + + is_tx = action == Actions.TX.value + is_success = jax.lax.bitwise_and(is_tx, jax.lax.bitwise_or(reward > 0, ret_c == 0)) + is_failure = jax.lax.bitwise_and(is_tx, reward < 0) return jax.lax.cond( - buffer == 0, + buffer_count == 0, reset, lambda: jax.lax.cond( - jax.lax.bitwise_and(state.backoff > 0, channel == 0), + jax.lax.bitwise_and(state.backoff > 0, channel_busy == 0), countdown, lambda: jax.lax.cond( - jax.lax.bitwise_and(state.backoff > 0, channel != 0), + jax.lax.bitwise_and(state.backoff > 0, channel_busy != 0), freeze, lambda: jax.lax.cond( - jax.lax.bitwise_or(reward > 0, ret_c == 0), + is_success, reset, lambda: jax.lax.cond( - reward < 0, + is_failure, double_cw, freeze ) @@ -70,13 +76,14 @@ def double_cw(): @staticmethod def sample(state, key, env_state): - buffer, channel, _, _, _ = env_state[-1] + buffer_count = env_state[-1][ObsIdx.BUFFER_PACKET_COUNT] + channel_busy = env_state[-1][ObsIdx.CHANNEL_LAST_CS_BUSY] return jnp.where( - buffer == 0, + buffer_count == 0, Actions.IDLE.value, jnp.where( - jax.lax.bitwise_and(state.backoff == 0, channel == 0), + jax.lax.bitwise_and(state.backoff == 0, channel_busy == 0), Actions.TX.value, Actions.CS.value ) diff --git a/ltc/run.py b/ltc/run.py index e16da42..da15dd2 100644 --- a/ltc/run.py +++ b/ltc/run.py @@ -1,8 +1,10 @@ import os os.environ['XLA_PYTHON_CLIENT_PREALLOCATE'] = 'false' os.environ['XLA_FLAGS'] = '--xla_gpu_enable_triton_gemm=false' +os.environ['JAX_ENABLE_X64'] = 'true' import argparse +from dataclasses import dataclass from functools import partial import cloudpickle @@ -13,13 +15,34 @@ from tqdm import trange from ltc.agents import BayesianDDQN, DCF, QNetwork, StochasticVariationalNetwork -from ltc.sim import InitialStateConf, cox_traffic, process_output, simulate -from ltc.sim.constants import INITIAL_CAPACITY, Actions +from ltc.sim import InitialStateConf, cox_traffic, normalize_obs, process_output, simulate, channel_state_selector +from ltc.sim.constants import * from ltc.utils.history import build_history_filename, ensure_clean_git_worktree, get_short_commit_hash -from ltc.utils.scan_states import Carry, Output +from ltc.utils.structs import * from ltc.utils.plots import plot_all, plot_first +@dataclass(frozen=True) +class StationScheduleConfig: + one_shot_step: int | None = None + one_shot_target: int | None = None + station_change_interval: int | None = None + station_change_delta: int = 0 + station_change_start_step: int = 0 + station_change_stop_step: int | None = None + station_change_target: int | None = None + + +@dataclass(frozen=True) +class RLStepConfig: + n_bins: int = 50 + schedule: StationScheduleConfig = StationScheduleConfig() + obs_enable_cs_tx_same_type: bool = False + obs_enable_tx_collision_other: bool = False + action_space_variant: str = 'txcs_duration_set' + max_duration: int = 5 + legacy_tx_duration: int = 5 + def init_agents(agent, key, n): keys = jax.random.split(key, n) @@ -52,40 +75,30 @@ def init_traffic(traffic, key, n): return states, step_fn -def schedule_active_stations( - active, - step, - *, - one_shot_step=None, - one_shot_target=None, - station_change_interval=None, - station_change_delta=0, - station_change_start_step=0, - station_change_stop_step=None, - station_change_target=None, -): +def schedule_active_stations(active, step, schedule): n = active.shape[0] next_active = active - if one_shot_step is not None and one_shot_target is not None: - target = jnp.asarray(one_shot_target, dtype=jnp.int32) + if schedule.one_shot_step is not None and schedule.one_shot_target is not None: + target = jnp.asarray(schedule.one_shot_target, dtype=jnp.int32) one_shot_active = jnp.arange(n) < target - next_active = jnp.where(step == one_shot_step, one_shot_active, next_active) + next_active = jnp.where(step == schedule.one_shot_step, one_shot_active, next_active) - if station_change_interval is not None and station_change_delta != 0: - should_change = step >= station_change_start_step - if station_change_stop_step is not None: - should_change = jnp.logical_and(should_change, step <= station_change_stop_step) + if schedule.station_change_interval is not None and schedule.station_change_delta != 0: + should_change = step >= schedule.station_change_start_step + if schedule.station_change_stop_step is not None: + should_change = jnp.logical_and(should_change, step <= schedule.station_change_stop_step) should_change = jnp.logical_and( - should_change, (step - station_change_start_step) % station_change_interval == 0 + should_change, + (step - schedule.station_change_start_step) % schedule.station_change_interval == 0, ) current_count = jnp.sum(next_active.astype(jnp.int32)) - updated_count = jnp.clip(current_count + station_change_delta, 0, n) - if station_change_target is not None: - target = jnp.asarray(station_change_target, dtype=jnp.int32) + updated_count = jnp.clip(current_count + schedule.station_change_delta, 0, n) + if schedule.station_change_target is not None: + target = jnp.asarray(schedule.station_change_target, dtype=jnp.int32) updated_count = jnp.where( - station_change_delta > 0, + schedule.station_change_delta > 0, jnp.minimum(updated_count, target), jnp.maximum(updated_count, target), ) @@ -96,42 +109,323 @@ def schedule_active_stations( 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 action_space_size(action_space_variant, max_duration): + if action_space_variant == 'txcs_duration_set': + return 2 * max_duration + if action_space_variant == 'tx_stage_commit': + return 3 + raise ValueError(f'Unknown action-space variant: {action_space_variant}') + + +def decode_drl_actions(raw_actions, staged_tx, action_space_variant, max_duration): + raw_actions = raw_actions.astype(jnp.int32) + + if action_space_variant == 'txcs_duration_set': + is_tx = raw_actions < max_duration + duration_idx = jnp.where(is_tx, raw_actions, raw_actions - max_duration) + action_type = jnp.where(is_tx, Actions.TX.value, Actions.CS.value) + duration = duration_idx + 1 + return action_type.astype(jnp.int32), duration.astype(jnp.int32), staged_tx + + if action_space_variant == 'tx_stage_commit': + is_stage = raw_actions == 1 + is_commit = raw_actions == 2 + staged_next = jnp.where(is_stage, staged_tx + 1, staged_tx) + commit_duration = jnp.clip(staged_tx, 1, max_duration) + has_stage = staged_tx > 0 + action_type = jnp.where(is_commit & has_stage, Actions.TX.value, Actions.CS.value).astype(jnp.int32) + duration = jnp.where(is_commit & has_stage, commit_duration, 1).astype(jnp.int32) + staged_next = jnp.where(is_commit, 0, staged_next).astype(jnp.int32) + return action_type, duration, staged_next + + raise ValueError(f'Unknown action-space variant: {action_space_variant}') + + +def decode_legacy_actions(raw_actions, legacy_tx_duration): + raw_actions = raw_actions.astype(jnp.int32) + action_type = jnp.where( + raw_actions == Actions.TX.value, + Actions.TX.value, + jnp.where(raw_actions == Actions.CS.value, Actions.CS.value, Actions.IDLE.value), + ).astype(jnp.int32) + duration = jnp.where( + raw_actions == Actions.TX.value, + legacy_tx_duration, + 1, + ).astype(jnp.int32) + return action_type, duration + + +def window_mean(hist, axis=None): + valid = hist >= 0.0 + count = jnp.sum(valid, axis=axis) + return jnp.where(count > 0, jnp.sum(jnp.where(valid, hist, 0.0), axis=axis) / count, -1.0) + + +def roll_append(arr, val): + return jnp.roll(arr, -1, axis=-1).at[..., -1].set(val) + + +def update_packet_tracking( + macro, slot_executing, slot_actions, macro_remaining, + tx_collision_other_now, agent_ready_now, ret_c, buffer_states, ): - def rl_step_coroutine(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, - ) + tx_executing_pre = slot_actions == Actions.TX.value + channel_state_prelim = channel_state_selector(slot_actions) + + sub_window_start = (macro_remaining % TX_SLOTS == 0) & slot_executing & tx_executing_pre + sub_collision_flag = jnp.where(sub_window_start, False, macro.sub_collision_flag) + sub_collision_flag = jnp.where( + tx_executing_pre & (channel_state_prelim == -1), True, sub_collision_flag, + ) + sub_collision_other_flag = jnp.where(sub_window_start, False, macro.sub_collision_other_flag) + sub_collision_other_flag = jnp.where( + tx_executing_pre & (channel_state_prelim == -1) & (tx_collision_other_now > 0.5), + True, sub_collision_other_flag, + ) + tx_packet_mask = ((macro_remaining - 1) % TX_SLOTS == 0) & slot_executing & tx_executing_pre + tx_success_mask = tx_packet_mask & ~sub_collision_flag + + collision_with_other = tx_collision_other_now > 0.5 + is_first_tx_slot = agent_ready_now & tx_executing_pre + header_collision = jnp.where(is_first_tx_slot, channel_state_prelim == -1, macro.header_collision) + header_collision_other = jnp.where( + is_first_tx_slot, (channel_state_prelim == -1) & collision_with_other, macro.header_collision_other, + ) + done_now = jnp.logical_and(slot_executing, macro_remaining == 1) + ack_collision_now = done_now & tx_executing_pre & (channel_state_prelim == -1) + ack_collision_other_now = ack_collision_now & collision_with_other + + initial_ret_c = jnp.where(is_first_tx_slot, ret_c, macro.initial_ret_c) + tx_started_empty = jnp.where( + is_first_tx_slot, jnp.all(buffer_states == EMPTY_PACKET_ID, axis=1), macro.tx_started_empty, + ) + + return ( + sub_collision_flag, sub_collision_other_flag, + tx_packet_mask, tx_success_mask, + header_collision, header_collision_other, + done_now, ack_collision_now, ack_collision_other_now, + initial_ret_c, tx_started_empty, + ) + + +def scale_tx_durations(action_types, durations): + return jnp.where( + action_types == Actions.TX.value, + durations * TX_SLOTS, + durations, + ) + + +def compute_obs_features( + obs_tracker, tx_hist, n_drl, + new_frames, planned_packets, successful_packets, channel_state, tx_mask, + cs_tx_same_type_now, tx_collision_other_now, +): + n = obs_tracker.arrival_hist.shape[0] + + arrival_hist = roll_append(obs_tracker.arrival_hist, new_frames.astype(jnp.float32)) + planned_tx_hist = roll_append(obs_tracker.planned_tx_hist, planned_packets.astype(jnp.float32)) + success_tx_hist = roll_append(obs_tracker.success_tx_hist, successful_packets.astype(jnp.float32)) + channel_busy_hist = roll_append(obs_tracker.channel_busy_hist, (channel_state != 0).astype(jnp.float32)) + collision_hist = roll_append(obs_tracker.collision_hist, (channel_state == -1).astype(jnp.float32)) + tx_hist = jnp.roll(tx_hist, -1, axis=0).at[-1].set(tx_mask.astype(jnp.int32)) + + planned_valid = planned_tx_hist >= 0.0 + planned_sum = jnp.sum(jnp.where(planned_valid, planned_tx_hist, 0.0), axis=1) + success_sum = jnp.sum(jnp.where(planned_valid, success_tx_hist, 0.0), axis=1) + back_pct = jnp.where(planned_sum > 0, success_sum / planned_sum, -1.0) + + tx_valid_rows = tx_hist[:, 0] >= 0 + unique_ltc_tx_window = jnp.where( + jnp.any(tx_valid_rows), + jnp.sum(jnp.any(jnp.where(tx_valid_rows[:, None], tx_hist[:, :n_drl], 0) > 0, axis=0).astype(jnp.float32)), + -1.0, + ) + + features = ObsFeatureInputs( + traffic_mean_arrival_rate=window_mean(arrival_hist, axis=1), + channel_occupancy_pct_window=jnp.broadcast_to(window_mean(channel_busy_hist), (n,)), + channel_collisions_pct_window=jnp.broadcast_to(window_mean(collision_hist), (n,)), + back_pct=back_pct, + unique_ltc_tx_window=jnp.broadcast_to(unique_ltc_tx_window, (n,)), + cs_tx_same_type_now=cs_tx_same_type_now, + tx_collision_other_now=tx_collision_other_now, + ) + + return (arrival_hist, planned_tx_hist, success_tx_hist, channel_busy_hist, collision_hist, tx_hist), features + + +def update_macro_state_from_ready_agents( + macro_state, obs_tracker, ready_drl, ready_legacy, + decoded_actions, n_drl +): + macro_action_types = macro_state.action_types + macro_remaining = macro_state.remaining + staged_tx = obs_tracker.staged_tx + + macro_action_types = macro_action_types.at[:n_drl].set( + jnp.where(ready_drl, decoded_actions.drl_action_types, macro_state.action_types[:n_drl]) + ) + drl_scaled_durations = scale_tx_durations(decoded_actions.drl_action_types, decoded_actions.drl_durations) + macro_remaining = macro_remaining.at[:n_drl].set( + jnp.where(ready_drl, drl_scaled_durations, macro_state.remaining[:n_drl]) + ) + staged_tx = staged_tx.at[:n_drl].set( + jnp.where(ready_drl, decoded_actions.drl_staged_next, obs_tracker.staged_tx[:n_drl]) + ) + + macro_action_types = macro_action_types.at[n_drl:].set( + jnp.where(ready_legacy, decoded_actions.legacy_action_types, macro_state.action_types[n_drl:]) + ) + legacy_scaled_durations = scale_tx_durations( + decoded_actions.legacy_action_types, decoded_actions.legacy_durations + ) + macro_remaining = macro_remaining.at[n_drl:].set( + jnp.where(ready_legacy, legacy_scaled_durations, macro_state.remaining[n_drl:]) + ) + + return macro_action_types, macro_remaining, staged_tx + + +def build_traffic(traffic_type, f3dB, loc, scale): + if traffic_type == 'custom': + params = {'f3dB': f3dB, 'loc': loc, 'scale': scale} + else: + presets = { + 'constant': {'f3dB': 1.0, 'loc': -1.0, 'scale': 0.0}, + 'saturated': {'f3dB': 1.0, 'loc': 5.0, 'scale': 0.0}, + 'bursty': {'f3dB': 0.1, 'loc': -5.0, 'scale': 5.0}, + } + params = presets.get(traffic_type) + if params is None: + raise ValueError(f'Unknown traffic type: {traffic_type}') + + return cox_traffic(initial_state=InitialStateConf.ZERO, **params) + + +def rl_step(drl_step, legacy_step, traffic_step, n, n_drl, config=RLStepConfig()): + def rl_step_fn(c, step): + active = schedule_active_stations(c.active, step, config.schedule) key, drl_keys, legacy_keys, traffic_key, reward_key = jax.random.split(c.key, 5) drl_keys = jax.random.split(drl_keys, n_drl) legacy_keys = jax.random.split(legacy_keys, n - n_drl) traffic_keys = jax.random.split(traffic_key, n) - drl_states, drl_actions = yield drl_step( - c.drl_states, drl_keys, c.obs[:n_drl], c.actions[:n_drl], c.rewards[:n_drl], c.terminals[:n_drl], active[:n_drl] + ready = c.macro.remaining == 0 + drl_ready = jnp.logical_and(active[:n_drl], ready[:n_drl]) + legacy_ready = jnp.logical_and(active[n_drl:], ready[n_drl:]) + + obs_norm = normalize_obs(c.obs, queue_size=c.buffer_states.shape[1]) + drl_states, drl_actions_raw = drl_step( + c.drl_states, drl_keys, obs_norm[:n_drl], c.actions[:n_drl], c.rewards[:n_drl], c.terminals[:n_drl], drl_ready ) - legacy_states, legacy_actions = legacy_step( - c.legacy_states, legacy_keys, c.obs[n_drl:], c.actions[n_drl:], c.rewards[n_drl:], c.terminals[n_drl:], active[n_drl:] + legacy_states, legacy_actions_raw = legacy_step( + c.legacy_states, legacy_keys, obs_norm[n_drl:], c.actions[n_drl:], c.rewards[n_drl:], c.terminals[n_drl:], legacy_ready + ) + + raw_actions = c.actions + raw_actions = raw_actions.at[:n_drl].set(jnp.where(drl_ready, drl_actions_raw, c.actions[:n_drl])) + raw_actions = raw_actions.at[n_drl:].set(jnp.where(legacy_ready, legacy_actions_raw, c.actions[n_drl:])) + + drl_action_types, drl_durations, drl_staged_next = decode_drl_actions( + raw_actions[:n_drl], c.obs_tracker.staged_tx[:n_drl], config.action_space_variant, config.max_duration + ) + legacy_action_types, legacy_durations = decode_legacy_actions( + raw_actions[n_drl:], config.legacy_tx_duration + ) + + decoded_actions = ActionDecodingResults( + drl_action_types=drl_action_types, + drl_durations=drl_durations, + drl_staged_next=drl_staged_next, + legacy_action_types=legacy_action_types, + legacy_durations=legacy_durations, + ) + + macro_action_types, macro_remaining, staged_tx = update_macro_state_from_ready_agents( + c.macro, c.obs_tracker, drl_ready, legacy_ready, + decoded_actions, n_drl + ) + + slot_executing = jnp.logical_and(active, macro_remaining > 0) + slot_actions = jnp.where(slot_executing, macro_action_types, Actions.IDLE.value) + tx_mask = slot_actions == Actions.TX.value + + is_ltc = jnp.arange(n) < n_drl + tx_idx = jnp.argmax(tx_mask.astype(jnp.int32)) + cs_tx_same_type_now = (is_ltc[tx_idx] == is_ltc).astype(jnp.float32) + diff_type = is_ltc[:, None] != is_ltc[None, :] + tx_collision_other_now = jnp.any(diff_type & tx_mask[None, :], axis=1).astype(jnp.float32) + + agent_ready_now = jnp.concatenate([drl_ready, legacy_ready]) + ret_c = c.obs[:, -1, ObsIdx.STATUS_RETRY_COUNTER].astype(jnp.int32) + ( + sub_collision_flag, sub_collision_other_flag, + tx_packet_mask, tx_success_mask, + header_collision, header_collision_other, + done_now, ack_collision_now, ack_collision_other_now, + initial_ret_c, tx_started_empty, + ) = update_packet_tracking( + c.macro, slot_executing, slot_actions, macro_remaining, + tx_collision_other_now, agent_ready_now, ret_c, c.buffer_states, ) - actions = jnp.concatenate([drl_actions, legacy_actions]) traffic_states, new_frames = traffic_step(c.traffic_states, traffic_keys) - buffer_states, channel_state = simulate(c.buffer_states, new_frames, actions) + ( + buffer_states, + buffer_birth_steps, + channel_state, + packet_seqs, + planned_packets, + successful_packets, + ) = simulate(c.buffer_states, c.obs_tracker.buffer_birth_steps, new_frames, slot_actions, c.packet_seqs, step, + tx_packet_mask=tx_packet_mask, tx_success_mask=tx_success_mask, active=active) + + (arrival_hist, planned_tx_hist, success_tx_hist, channel_busy_hist, collision_hist, tx_hist), obs_features = \ + compute_obs_features( + c.obs_tracker, c.tx_hist, n_drl, + new_frames, planned_packets, successful_packets, channel_state, tx_mask, + cs_tx_same_type_now, tx_collision_other_now, + ) + + packet_collided = tx_packet_mask & sub_collision_flag + k_tx_accum = c.macro.tx_success_accum + successful_packets.astype(jnp.float32) + k_tx_planned = c.macro.tx_planned_accum + tx_packet_mask.astype(jnp.float32) + k_coll_ltc = c.macro.k_coll_ltc + jnp.where(packet_collided & ~sub_collision_other_flag, 1.0, 0.0) + k_coll_coex = c.macro.k_coll_coex + jnp.where(packet_collided & sub_collision_other_flag, 1.0, 0.0) + + obs_config = ObsFeatureConfig( + enable_cs_tx_same_type=config.obs_enable_cs_tx_same_type, + enable_tx_collision_other=config.obs_enable_tx_collision_other, + ) obs, rewards, powers = process_output( - c.buffer_states, buffer_states, c.power_states, channel_state, c.obs, actions, c.terminals, reward_key + c.buffer_states, buffer_states, buffer_birth_steps, c.power_states, channel_state, c.obs, slot_actions, + c.terminals, reward_key, step, obs_features, obs_config, + roll_obs=done_now, + done_now=done_now, + k_tx=k_tx_accum, + k_tx_planned=k_tx_planned, + k_coll_ltc=k_coll_ltc, + k_coll_coex=k_coll_coex, + header_collision=header_collision, + header_collision_other=header_collision_other, + ack_collision=ack_collision_now, + ack_collision_other=ack_collision_other_now, + initial_ret_c=initial_ret_c, + tx_started_empty=tx_started_empty, + tx_packet_mask=tx_packet_mask, + tx_success_mask=tx_success_mask, ) + + # Reset per-macro accumulators when action completes + k_tx_accum = jnp.where(done_now, 0.0, k_tx_accum) + k_tx_planned = jnp.where(done_now, 0.0, k_tx_planned) + k_coll_ltc = jnp.where(done_now, 0.0, k_coll_ltc) + k_coll_coex = jnp.where(done_now, 0.0, k_coll_coex) + macro_remaining = jnp.where(slot_executing, macro_remaining - 1, macro_remaining) terminals = jnp.logical_or(c.terminals, powers < 0) if n_drl > 0: @@ -139,46 +433,59 @@ def rl_step_coroutine(c, step): flat_params, _ = jax.tree.flatten(params) flat_params = jax.tree.map(lambda x: x.reshape(n_drl, -1), flat_params) flat_params = jnp.hstack(flat_params) - hist, bin_edges = jax.vmap(jnp.histogram, in_axes=(0, None))(flat_params, n_bins) + hist, bin_edges = jax.vmap(jnp.histogram, in_axes=(0, None))(flat_params, config.n_bins) else: hist, bin_edges = None, None + macro_state = MacroActionState( + remaining=macro_remaining, + action_types=macro_action_types, + tx_success_accum=k_tx_accum, + tx_planned_accum=k_tx_planned, + k_coll_ltc=k_coll_ltc, + k_coll_coex=k_coll_coex, + sub_collision_flag=sub_collision_flag, + sub_collision_other_flag=sub_collision_other_flag, + header_collision=header_collision, + header_collision_other=header_collision_other, + initial_ret_c=initial_ret_c, + tx_started_empty=tx_started_empty, + ) + obs_tracker_state = ObsTrackerState( + buffer_birth_steps=buffer_birth_steps, + arrival_hist=arrival_hist, + planned_tx_hist=planned_tx_hist, + success_tx_hist=success_tx_hist, + channel_busy_hist=channel_busy_hist, + collision_hist=collision_hist, + staged_tx=staged_tx, + ) c = Carry( - drl_states, legacy_states, traffic_states, buffer_states, powers, - channel_state, key, obs, actions, rewards, terminals, active + drl_states, legacy_states, traffic_states, packet_seqs, buffer_states, + tx_hist, powers, channel_state, key, obs, raw_actions, rewards, terminals, active, + macro_state, obs_tracker_state ) o = Output( - legacy_states, obs, actions, rewards, terminals, buffer_states, powers, - (new_frames > 0).astype(int), channel_state, active, hist, bin_edges + legacy_states, obs, raw_actions, rewards, terminals, buffer_states, powers, + new_frames, channel_state, active, hist, bin_edges, successful_packets ) - yield c, o - - def rl_step_fn(*args, **kwargs): - gen = rl_step_coroutine(*args, **kwargs) - intercepted = next(gen) - return gen.send(intercepted) + return c, o - def pre_rl_fn(*args, **kwargs): - gen = rl_step_coroutine(*args, **kwargs) - intercepted = next(gen) - return intercepted + return rl_step_fn - def post_rl_fn(intermediate, *args, **kwargs): - gen = rl_step_coroutine(*args, **kwargs) - # Unused computations will be DCE-ed - _ = next(gen) - return gen.send(intermediate) - return rl_step_fn, pre_rl_fn, post_rl_fn - - -def setup_args(): +if __name__ == '__main__': parser = argparse.ArgumentParser(description="Run the RL network simulation with configurable parameters.") parser.add_argument('--n', type=int, default=5, help='Initial number of agents in the simulation.') parser.add_argument('--n_final', type=int, help='Final number of agents in the simulation.') - parser.add_argument('--n_epochs', type=int, default=50, help='Number of training epochs to run.') - parser.add_argument('--n_steps', type=int, default=2000, help='Number of steps per epoch.') + parser.add_argument('--n_epochs', type=int, default=100, help='Number of training epochs to run.') + parser.add_argument('--n_steps', type=int, default=3000, help='Number of steps per epoch.') parser.add_argument('--window_size', type=int, default=1, help='Size of the observation window for each agent.') + parser.add_argument('--queue_size', type=int, default=16, help='Fixed packet-queue length for each station.') + parser.add_argument('--obs_stats_window', type=int, default=1000, help='Window length for rolling observation statistics.') + parser.add_argument('--obs_enable_cs_tx_same_type', action='store_true', default=False, help='Enable CS me-vs-other transmitter feature.') + parser.add_argument('--obs_enable_tx_collision_other', action='store_true', default=False, help='Enable TX collision-involves-other feature.') + parser.add_argument('--action_space_variant', type=str, default='txcs_duration_set', choices=['txcs_duration_set', 'tx_stage_commit'], help='Macro action-space variant.') parser.add_argument('--seed', type=int, default=42, help='Random seed for reproducibility.') parser.add_argument('--save_plots', action='store_true', default=False, help='Whether to save the generated plots.') parser.add_argument('--loc', type=float, default=5.0, help='loc traffic generator parameter.') @@ -190,64 +497,46 @@ def setup_args(): parser.add_argument('--station_change_stop_step', type=int, help='Global step when periodic station changes stop.') parser.add_argument('--traffic_type', type=str, default='saturated', choices=['constant', 'saturated', 'bursty', 'custom'],help="Traffic model to use: 'constant', 'saturated', 'bursty', or 'custom'.") parser.add_argument('--legacy_type', type=str, default='tdma', choices=['q-aloha', 'eb-aloha', 'fw-aloha', 'tdma'], help="Legacy agent type to use: 'q-aloha', 'eb-aloha', 'fw-aloha', or 'tdma'.") + parser.add_argument('--skip_git_check', action='store_true', default=False, help='Skip clean git worktree check.') args = parser.parse_args() - return args - -if __name__ == '__main__': - args = setup_args() - ensure_clean_git_worktree() + if not args.skip_git_check: + ensure_clean_git_worktree() commit_hash = get_short_commit_hash() n_init = args.n - has_n_final = args.n_final is not None - n_final = args.n_final if has_n_final else n_init + n_final = args.n_final if args.n_final is not None else n_init n = max(n_init, n_final) - n_epochs = args.n_epochs - n_steps = args.n_steps - total_steps = n_epochs * n_steps - window_size = args.window_size - seed = args.seed - traffic_type = args.traffic_type - - loc = args.loc - scale = args.scale - f3dB = args.f3dB - station_change_interval = args.station_change_interval - station_change_delta = args.station_change_delta - station_change_start_step = args.station_change_start_step - station_change_stop_step = args.station_change_stop_step + total_steps = args.n_epochs * args.n_steps - if station_change_interval is not None and station_change_interval <= 0: + if args.queue_size <= 0: + raise ValueError('--queue_size must be positive.') + if args.obs_stats_window <= 0: + raise ValueError('--obs_stats_window must be positive.') + + station_change_start_step = args.station_change_start_step + if args.station_change_interval is not None and args.station_change_interval <= 0: raise ValueError('--station_change_interval must be positive.') - if station_change_interval is None: - if station_change_start_step is not None or station_change_stop_step is not None: + if args.station_change_interval is None: + if args.station_change_start_step is not None or args.station_change_stop_step is not None: raise ValueError('--station_change_start_step/--station_change_stop_step require --station_change_interval.') else: - if station_change_delta == 0: + if args.station_change_delta == 0: raise ValueError('--station_change_delta must be non-zero when --station_change_interval is used.') if station_change_start_step is None: - station_change_start_step = station_change_interval + station_change_start_step = args.station_change_interval one_shot_step = None one_shot_target = None - station_change_target = n_final if has_n_final else None - if station_change_interval is None and n_final != n_init: + station_change_target = n_final if args.n_final is not None else None + if args.station_change_interval is None and n_final != n_init: one_shot_step = total_steps // 2 one_shot_target = n_final - key = jax.random.key(seed) - num_actions = 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) - rewards = jnp.zeros(n) - terminals = jnp.full(n, False, dtype=bool) - active = jnp.ones(n, dtype=bool).at[n_init:].set(False) - - lr_schedule = optax.cosine_decay_schedule(init_value=1e-4, decay_steps=60000, alpha=0.01) + key = jax.random.key(args.seed) + num_actions = action_space_size(args.action_space_variant, MAX_MACRO_DURATION) + + lr_schedule = optax.cosine_decay_schedule(init_value=1e-4, decay_steps=total_steps, alpha=0.01) optimizer = optax.chain( optax.clip_by_global_norm(1.0), optax.adam(lr_schedule, b1=0.95, b2=0.95) @@ -255,7 +544,7 @@ def setup_args(): drl = BayesianDDQN( q_network=StochasticVariationalNetwork(QNetwork(num_actions, num_layers=1, dim=64, num_heads=4)), - obs_space_shape=obs.shape[1:], + obs_space_shape=(args.window_size, OBS_SIZE), act_space_size=num_actions, optimizer=optimizer, experience_replay_buffer_size=30000, @@ -264,12 +553,12 @@ def setup_args(): discount=0.95, epsilon=1.0, epsilon_decay=0.999, - epsilon_min=0.0, + epsilon_min=0.05, tau=0.05 ) key, init_key = jax.random.split(key) drl_states, drl_step = init_agents(drl, init_key, n) - drl_states = drl_states.replace(prev_env_state=drl_states.prev_env_state.astype(int)) + drl_states = drl_states.replace(prev_env_state=drl_states.prev_env_state.astype(jnp.float32)) dcf = DCF() key, init_key = jax.random.split(key) @@ -277,47 +566,76 @@ def setup_args(): key, init_key = jax.random.split(key) - if traffic_type == 'constant': - traffic = cox_traffic(f3dB=1.0, loc=-1.0, scale=0.0, initial_state=InitialStateConf.ZERO) - elif traffic_type == 'saturated': - traffic = cox_traffic(f3dB=1.0, loc=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 = build_traffic(args.traffic_type, args.f3dB, args.loc, args.scale) traffic_states, traffic_step = init_traffic(traffic, init_key, n) - rl_step_fn, _, _ = rl_step( - drl_step, - legacy_step, - traffic_step, - n, - n, - one_shot_step=one_shot_step, - one_shot_target=one_shot_target, - station_change_interval=station_change_interval, - station_change_delta=station_change_delta, - station_change_start_step=station_change_start_step, - station_change_stop_step=station_change_stop_step, - station_change_target=station_change_target, + step_config = RLStepConfig( + schedule=StationScheduleConfig( + one_shot_step=one_shot_step, + one_shot_target=one_shot_target, + station_change_interval=args.station_change_interval, + station_change_delta=args.station_change_delta, + station_change_start_step=station_change_start_step, + station_change_stop_step=args.station_change_stop_step, + station_change_target=station_change_target, + ), + obs_enable_cs_tx_same_type=args.obs_enable_cs_tx_same_type, + obs_enable_tx_collision_other=args.obs_enable_tx_collision_other, + action_space_variant=args.action_space_variant, + max_duration=MAX_MACRO_DURATION, + legacy_tx_duration=LEGACY_TX_DURATION, ) - rl_step_fn = jax.jit(rl_step_fn) + rl_step_fn = jax.jit(rl_step(drl_step, legacy_step, traffic_step, n, n, config=step_config)) + carry = Carry( - drl_states, legacy_states, traffic_states, buffer_states, power_states, - channel_state, key, obs, actions, rewards, terminals, active + drl_states=drl_states, + legacy_states=legacy_states, + traffic_states=traffic_states, + packet_seqs=jnp.zeros(n, dtype=jnp.int64), + buffer_states=jnp.full((n, args.queue_size), EMPTY_PACKET_ID, dtype=jnp.int64), + tx_hist=jnp.full((args.obs_stats_window, n), -1, dtype=jnp.int32), + power_states=jnp.full(n, INITIAL_CAPACITY, dtype=jnp.int32), + channel_state=0, + key=key, + obs=jnp.zeros((n, args.window_size, OBS_SIZE), dtype=jnp.float32), + actions=jnp.zeros(n, dtype=jnp.int32), + rewards=jnp.zeros(n), + terminals=jnp.full(n, False, dtype=bool), + active=jnp.ones(n, dtype=bool).at[n_init:].set(False), + macro=MacroActionState( + remaining=jnp.zeros(n, dtype=jnp.int32), + action_types=jnp.full(n, Actions.IDLE.value, dtype=jnp.int32), + tx_success_accum=jnp.zeros(n, dtype=jnp.float32), + tx_planned_accum=jnp.zeros(n, dtype=jnp.float32), + k_coll_ltc=jnp.zeros(n, dtype=jnp.float32), + k_coll_coex=jnp.zeros(n, dtype=jnp.float32), + sub_collision_flag=jnp.zeros(n, dtype=bool), + sub_collision_other_flag=jnp.zeros(n, dtype=bool), + header_collision=jnp.zeros(n, dtype=bool), + header_collision_other=jnp.zeros(n, dtype=bool), + initial_ret_c=jnp.zeros(n, dtype=jnp.int32), + tx_started_empty=jnp.zeros(n, dtype=bool), + ), + obs_tracker=ObsTrackerState( + buffer_birth_steps=jnp.full((n, args.queue_size), -1, dtype=jnp.int32), + arrival_hist=jnp.full((n, args.obs_stats_window), -1.0, dtype=jnp.float32), + planned_tx_hist=jnp.full((n, args.obs_stats_window), -1.0, dtype=jnp.float32), + success_tx_hist=jnp.full((n, args.obs_stats_window), -1.0, dtype=jnp.float32), + channel_busy_hist=jnp.full(args.obs_stats_window, -1.0, dtype=jnp.float32), + collision_hist=jnp.full(args.obs_stats_window, -1.0, dtype=jnp.float32), + staged_tx=jnp.zeros(n, dtype=jnp.int32), + ), ) + all_outputs = [] - for epoch in trange(n_epochs): - global_steps = jnp.arange(epoch * n_steps, (epoch + 1) * n_steps, dtype=jnp.int32) + for epoch in trange(args.n_epochs): + global_steps = jnp.arange(epoch * args.n_steps, (epoch + 1) * args.n_steps, dtype=jnp.int32) carry, output = jax.lax.scan(rl_step_fn, carry, xs=global_steps) all_outputs.append(output) all_outputs = jax.tree.map(lambda *x: jnp.stack(x), *all_outputs) - filename = build_history_filename(n_init, n_final, seed, commit_hash) + filename = build_history_filename(n_init, n_final, args.seed, commit_hash) with lz4.frame.open(filename, 'wb') as f: cloudpickle.dump((carry.drl_states, all_outputs, vars(args)), f) diff --git a/ltc/sim/__init__.py b/ltc/sim/__init__.py index 325b814..f428fa0 100644 --- a/ltc/sim/__init__.py +++ b/ltc/sim/__init__.py @@ -1,3 +1,3 @@ from ltc.sim.traffic import InitialStateConf, ModelState, TrafficModel, cox_traffic -from ltc.sim.process_output import process_output -from ltc.sim.sim import simulate +from ltc.sim.process_output import process_output, normalize_obs +from ltc.sim.sim import simulate, channel_state_selector diff --git a/ltc/sim/constants.py b/ltc/sim/constants.py index 203828e..a4affa3 100644 --- a/ltc/sim/constants.py +++ b/ltc/sim/constants.py @@ -10,13 +10,19 @@ class Actions(Enum): IDLE = 2 +""" +Time scaling: one TX slot occupies TX_SLOTS mini-slots. +During one TX, CS agents can observe TX_SLOTS times. +""" +TX_SLOTS = 20 + """ Simulation parameters """ MAX_RETRANSMISSION = 8 -SAFE_IDLE_PERIOD = 25 +SAFE_IDLE_PERIOD = 25 * TX_SLOTS SAFE_IDLE_PERIOD_STD = 3 -PENALIZED_IDLE_PERIOD = 25 +PENALIZED_IDLE_PERIOD = 25 * TX_SLOTS """ Rewards and penalties @@ -25,9 +31,15 @@ class Actions(Enum): EMPTY_BUFFER_REWARD = 0.5 NO_TX_REWARD = 0.0 NO_TX_PENALTY = -1.0 -EMPTY_TX_PENALTY = -0.5 COLLISION_PENALTY = -1.0 -MAX_RETRANSMISSION_PENALTY = -1.0 + +TX_SIZE_THRESHOLD = 3 +TX_SIZE_PENALTY_WINDOW = 2 +TX_SIZE_PENALTY = -1.0 +TX_LTC_COLLISION_PENALTY = -1.0 +TX_COEX_COLLISION_PENALTY = -5.0 +TX_EMPTY_BUFFER_PENALTY = -0.5 +TX_MAX_RETRANSMISSION_PENALTY = -1.0 """ Power consumption @@ -42,3 +54,42 @@ class Actions(Enum): https://ieeexplore.ieee.org/document/8930559 """ TAU = 5.484 * 1e-3 + +""" +Queue representation +""" +EMPTY_PACKET_ID = -1 + +""" +Macro-action defaults +""" +MAX_MACRO_DURATION = 5 +LEGACY_TX_DURATION = 5 + +""" +Observation normalization constants +""" +OBS_UNIQUE_LTC_TX_NORM = 10 +OBS_TRAFFIC_ARRIVAL_NORM = 10 + + +class ObsIdx: + BUFFER_OCCUPANCY_PCT = 0 + BUFFER_PACKET_COUNT = 1 + BUFFER_OLDEST_AGE_NORM = 2 + TRAFFIC_MEAN_ARRIVAL_RATE = 3 + CHANNEL_LAST_CS_BUSY = 4 + CHANNEL_LAST_CS_TX_SAME_TYPE = 5 + CHANNEL_LAST_TX_COLLISION_OTHER = 6 + CHANNEL_OCCUPANCY_PCT_WINDOW = 7 + CHANNEL_COLLISIONS_PCT_WINDOW = 8 + ACTION_TX = 9 + ACTION_CS = 10 + ACTION_IDLE = 11 + ACTION_BACK_PCT = 12 + STATUS_RETRY_COUNTER = 13 + STATUS_NO_TX_COUNTER = 14 + STATUS_UNIQUE_LTC_TX_WINDOW = 15 + + +OBS_SIZE = 16 diff --git a/ltc/sim/process_output.py b/ltc/sim/process_output.py index ee55f3e..360e532 100644 --- a/ltc/sim/process_output.py +++ b/ltc/sim/process_output.py @@ -2,11 +2,12 @@ import jax.numpy as jnp from ltc.sim.constants import * +from ltc.sim.sim import is_buffer_empty def no_transmission(args): _, buffer_state, _, _, _, _ = args - return jax.lax.cond(buffer_state == 0, idle_empty_buffer, idle_full_buffer, args) + return jax.lax.cond(is_buffer_empty(buffer_state), idle_empty_buffer, idle_full_buffer, args) def idle_empty_buffer(_): @@ -19,7 +20,7 @@ def idle_empty_buffer(_): def idle_full_buffer(args): _, _, _, _, no_tx, key = args rnd_factor = jax.random.normal(key) * SAFE_IDLE_PERIOD_STD - rnd_factor = jnp.round(rnd_factor).astype(int) + rnd_factor = jnp.round(rnd_factor).astype(jnp.int32) safe_period = SAFE_IDLE_PERIOD + rnd_factor args += (safe_period,) return jax.lax.cond(no_tx < safe_period, no_transmission_short, no_transmission_long, args) @@ -40,56 +41,190 @@ def no_transmission_long(args): 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 reset_counters(_): + ret_c = 0 + no_tx = 0 + return 0.0, ret_c, no_tx -def transmission_with_collision(args): - _, _, ret_c, _, _, _ = args - return jax.lax.cond(ret_c < MAX_RETRANSMISSION, retransmission, max_retransmission_collision, args) +def tx_macro_reward( + k_tx, k_tx_planned, k_coll_ltc, k_coll_coex, tx_started_empty, initial_ret_c, + header_collision, ack_collision, header_collision_other, ack_collision_other, +): + hdr_ack = header_collision | ack_collision + k_tx_eff = jnp.where(hdr_ack, 0.0, k_tx) + # Header/ACK collision makes all frames unreadable: all planned sub-windows become collisions + collision_with_other = (header_collision & header_collision_other) | (ack_collision & ack_collision_other) + k_coll_ltc_eff = jnp.where(hdr_ack, jnp.where(~collision_with_other, k_tx_planned, 0.0), k_coll_ltc) + k_coll_coex_eff = jnp.where(hdr_ack, jnp.where(collision_with_other, k_tx_planned, 0.0), k_coll_coex) -def max_retransmission_collision(_): - reward = MAX_RETRANSMISSION_PENALTY - ret_c = 0 - no_tx = 0 - return reward, ret_c, no_tx + k_coll = k_coll_ltc_eff + k_coll_coex_eff + any_collision = k_coll > 0 + net_data_slots = jnp.maximum(k_tx_eff * TX_SLOTS - 2.0, k_tx_eff) + success_reward = jnp.where( + k_tx_eff > 0, + TX_REWARD * net_data_slots / TX_SLOTS, + 0.0, + ) -def retransmission(args): - _, _, ret_c, _, _, _ = args - reward = COLLISION_PENALTY - ret_c = ret_c + 1 - no_tx = 0 - return reward, ret_c, no_tx + size_penalty = jnp.where( + k_tx_planned >= TX_SIZE_THRESHOLD, + TX_SIZE_PENALTY * jnp.minimum(1.0, (k_tx_planned - TX_SIZE_THRESHOLD + 1.0) / TX_SIZE_PENALTY_WINDOW), + 0.0, + ) + # Collision penalty scales with absolute count; max-retry penalty is additive + collision_reward = jnp.where( + any_collision, + TX_LTC_COLLISION_PENALTY * k_coll_ltc_eff + + TX_COEX_COLLISION_PENALTY * k_coll_coex_eff + + jnp.where(initial_ret_c >= MAX_RETRANSMISSION, TX_MAX_RETRANSMISSION_PENALTY, 0.0), + 0.0, + ) -def transmission_without_collision(args): - _, buffer_state, _, _, _, _ = args - return jax.lax.cond(buffer_state > 0, successful_transmission, empty_buffer_transmission, args) + # Empty-buffer penalty scales with planned duration + empty_penalty = jnp.where( + jnp.logical_and(tx_started_empty, ~any_collision), + TX_EMPTY_BUFFER_PENALTY * k_tx_planned, + 0.0, + ) + return success_reward + size_penalty + collision_reward + empty_penalty -def empty_buffer_transmission(_): - reward = EMPTY_TX_PENALTY - ret_c = 0 - no_tx = 0 - return reward, ret_c, no_tx +def oldest_packet_age_norm(buffer_state, buffer_birth_state, current_step): + valid = buffer_state != EMPTY_PACKET_ID + oldest_birth = jnp.min(jnp.where(valid, buffer_birth_state, current_step)) + age = jnp.where(jnp.any(valid), current_step - oldest_birth, -1) + return jnp.where(age >= 0, age.astype(jnp.float32), -1.0) -def successful_transmission(args): - _, _, ret_c, _, _, _ = args - reward = TX_REWARD - ret_c = 0 - no_tx = 0 - return reward, ret_c, no_tx +def build_observation_entry( + buffer_state, new_buffer_state, new_buffer_birth_state, action, channel_state, + current_step, obs_features, obs_config, ret_c, no_tx +): + buffer_count = jnp.sum((new_buffer_state != EMPTY_PACKET_ID).astype(jnp.int32)) + buffer_occupancy_pct = buffer_count.astype(jnp.float32) / new_buffer_state.shape[0] + packet_age_norm = oldest_packet_age_norm(new_buffer_state, new_buffer_birth_state, current_step) -def process_output_i(buffer_state, new_buffer_state, power_state, channel_state, obs, action, terminal, key): - _, _, ret_c, no_tx, _ = obs[-1] - args = (action, buffer_state, ret_c, channel_state, no_tx, key) + cs_busy = jnp.where(action == Actions.CS.value, channel_state != 0, -1).astype(jnp.float32) + cs_tx_same_type = jnp.where( + obs_config.enable_cs_tx_same_type, + jnp.where(jnp.logical_and(action == Actions.CS.value, channel_state == 1), obs_features.cs_tx_same_type_now, -1.0), + -1.0, + ) + tx_collision_other = jnp.where( + obs_config.enable_tx_collision_other, + jnp.where(jnp.logical_and(action == Actions.TX.value, channel_state == -1), obs_features.tx_collision_other_now, -1.0), + -1.0, + ) - reward, ret_c, no_tx = jax.lax.cond(action == Actions.TX.value, transmission, no_transmission, args) + action_tx = jnp.where(action == Actions.TX.value, 1.0, 0.0) + action_cs = jnp.where(action == Actions.CS.value, 1.0, 0.0) + action_idle = jnp.where(action == Actions.IDLE.value, 1.0, 0.0) + + return jnp.array([ + buffer_occupancy_pct, + buffer_count.astype(jnp.float32), + packet_age_norm, + obs_features.traffic_mean_arrival_rate, + cs_busy, + cs_tx_same_type.astype(jnp.float32), + tx_collision_other.astype(jnp.float32), + obs_features.channel_occupancy_pct_window, + obs_features.channel_collisions_pct_window, + action_tx, + action_cs, + action_idle, + obs_features.back_pct, + ret_c.astype(jnp.float32), + no_tx.astype(jnp.float32), + obs_features.unique_ltc_tx_window, + ], dtype=jnp.float32) + + +def normalize_obs(obs, queue_size): + obs = obs.at[..., ObsIdx.BUFFER_PACKET_COUNT].divide(queue_size) + raw_age = obs[..., ObsIdx.BUFFER_OLDEST_AGE_NORM] + obs = obs.at[..., ObsIdx.BUFFER_OLDEST_AGE_NORM].set( + jnp.where(raw_age >= 0, jnp.minimum(raw_age / SAFE_IDLE_PERIOD, 1.0), raw_age) + ) + obs = obs.at[..., ObsIdx.TRAFFIC_MEAN_ARRIVAL_RATE].divide(OBS_TRAFFIC_ARRIVAL_NORM) + obs = obs.at[..., ObsIdx.STATUS_RETRY_COUNTER].divide(MAX_RETRANSMISSION) + obs = obs.at[..., ObsIdx.STATUS_NO_TX_COUNTER].divide(SAFE_IDLE_PERIOD + PENALIZED_IDLE_PERIOD) + obs = obs.at[..., ObsIdx.STATUS_UNIQUE_LTC_TX_WINDOW].divide(OBS_UNIQUE_LTC_TX_NORM) + return obs + + +def process_output_i( + buffer_state, + new_buffer_state, + new_buffer_birth_state, + power_state, + channel_state, + obs, + action, + terminal, + key, + current_step, + obs_features, + obs_config, + roll_obs, + done_now, + k_tx, + k_tx_planned, + k_coll_ltc, + k_coll_coex, + header_collision, + header_collision_other, + ack_collision, + ack_collision_other, + initial_ret_c, + tx_started_empty, + tx_packet_mask, + tx_success_mask, +): + old_ret_c = obs[-1, ObsIdx.STATUS_RETRY_COUNTER].astype(jnp.int64) + no_tx = obs[-1, ObsIdx.STATUS_NO_TX_COUNTER].astype(jnp.int64) + + is_tx = action == Actions.TX.value + + # CS/IDLE: get slot reward, no_tx update, and ret_c (resets if buffer empty) + cs_args = (action, buffer_state, old_ret_c, channel_state, no_tx, key) + cs_reward, cs_ret_c, cs_no_tx = no_transmission(cs_args) + + # TX: no_tx always resets (agent is transmitting); slot reward is 0 (macro reward below) + no_tx = jnp.where(is_tx, jnp.zeros_like(no_tx), cs_no_tx) + slot_reward = jnp.where(is_tx, 0.0, cs_reward) + + # ret_c for TX: update once per sub-window completion (tx_packet_mask), not every slot + tx_sub_success = is_tx & tx_packet_mask & tx_success_mask + tx_sub_coll_retry = is_tx & tx_packet_mask & ~tx_success_mask & (old_ret_c < MAX_RETRANSMISSION) + tx_sub_coll_reset = is_tx & tx_packet_mask & ~tx_success_mask & (old_ret_c >= MAX_RETRANSMISSION) + + ret_c = jnp.where(is_tx, old_ret_c, cs_ret_c) # CS uses cs logic; TX holds by default + ret_c = jnp.where(tx_sub_success, 0, ret_c) + ret_c = jnp.where(tx_sub_coll_retry, old_ret_c + 1, ret_c) + ret_c = jnp.where(tx_sub_coll_reset, 0, ret_c) + + if obs_config.enable_tx_collision_other: + k_coll_ltc_r, k_coll_coex_r = k_coll_ltc, k_coll_coex + hdr_other_r, ack_other_r = header_collision_other, ack_collision_other + else: + k_coll_ltc_r, k_coll_coex_r = k_coll_ltc + k_coll_coex, jnp.zeros_like(k_coll_coex) + hdr_other_r, ack_other_r = jnp.zeros_like(header_collision_other), jnp.zeros_like(ack_collision_other) + + macro_reward = tx_macro_reward( + k_tx, k_tx_planned, k_coll_ltc_r, k_coll_coex_r, tx_started_empty, initial_ret_c, + header_collision, ack_collision, hdr_other_r, ack_other_r, + ) + reward = jnp.where( + action == Actions.TX.value, + jnp.where(done_now, macro_reward, 0.0), + slot_reward, + ) reward = jnp.where(terminal, 0., reward) channel_state = jnp.where(action == Actions.CS.value, channel_state, -1) @@ -104,14 +239,90 @@ def process_output_i(buffer_state, new_buffer_state, power_state, channel_state, ) ) - obs_t = jnp.array([new_buffer_state, channel_state, ret_c, no_tx, action]) - obs = jnp.roll(obs, -1, axis=0) - obs = obs.at[-1].set(obs_t) + obs_t = build_observation_entry( + buffer_state, new_buffer_state, new_buffer_birth_state, action, channel_state, + current_step, obs_features, obs_config, ret_c, no_tx + ) + obs = jax.lax.cond( + roll_obs, + lambda x: jnp.roll(x, -1, axis=0).at[-1].set(obs_t), + lambda x: x.at[-1].set(obs_t), + obs, + ) return obs, reward, power -def process_output(buffer_states, new_buffer_states, power_states, channel_state, obs, actions, terminals, key): - return jax.vmap(process_output_i, in_axes=(0, 0, 0, None, 0, 0, 0, None))( - buffer_states, new_buffer_states, power_states, channel_state, obs, actions, terminals, key +def process_output( + buffer_states, + new_buffer_states, + new_buffer_birth_states, + power_states, + channel_state, + obs, + actions, + terminals, + key, + current_step, + obs_features, + obs_config, + roll_obs=True, + done_now=False, + k_tx=None, + k_tx_planned=None, + k_coll_ltc=None, + k_coll_coex=None, + header_collision=None, + header_collision_other=None, + ack_collision=None, + ack_collision_other=None, + initial_ret_c=None, + tx_started_empty=None, + tx_packet_mask=None, + tx_success_mask=None, +): + n = actions.shape[0] + keys = jax.random.split(key, n) + roll_obs = jnp.broadcast_to(jnp.asarray(roll_obs), actions.shape) + done_now = jnp.broadcast_to(jnp.asarray(done_now), actions.shape) + k_tx = jnp.broadcast_to(jnp.asarray(k_tx if k_tx is not None else 0.0, dtype=jnp.float32), (n,)) + k_tx_planned = jnp.broadcast_to(jnp.asarray(k_tx_planned if k_tx_planned is not None else 0.0, dtype=jnp.float32), (n,)) + k_coll_ltc = jnp.broadcast_to(jnp.asarray(k_coll_ltc if k_coll_ltc is not None else 0.0, dtype=jnp.float32), (n,)) + k_coll_coex = jnp.broadcast_to(jnp.asarray(k_coll_coex if k_coll_coex is not None else 0.0, dtype=jnp.float32), (n,)) + header_collision = jnp.broadcast_to(jnp.asarray(header_collision if header_collision is not None else False), (n,)) + header_collision_other = jnp.broadcast_to(jnp.asarray(header_collision_other if header_collision_other is not None else False), (n,)) + ack_collision = jnp.broadcast_to(jnp.asarray(ack_collision if ack_collision is not None else False), (n,)) + ack_collision_other = jnp.broadcast_to(jnp.asarray(ack_collision_other if ack_collision_other is not None else False), (n,)) + initial_ret_c = jnp.broadcast_to(jnp.asarray(initial_ret_c if initial_ret_c is not None else 0, dtype=jnp.int32), (n,)) + tx_started_empty = jnp.broadcast_to(jnp.asarray(tx_started_empty if tx_started_empty is not None else False), (n,)) + tx_packet_mask = jnp.broadcast_to(jnp.asarray(tx_packet_mask if tx_packet_mask is not None else False), (n,)) + tx_success_mask = jnp.broadcast_to(jnp.asarray(tx_success_mask if tx_success_mask is not None else False), (n,)) + + return jax.vmap(process_output_i, in_axes=(0, 0, 0, 0, None, 0, 0, 0, 0, None, 0, None, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0))( + buffer_states, + new_buffer_states, + new_buffer_birth_states, + power_states, + channel_state, + obs, + actions, + terminals, + keys, + current_step, + obs_features, + obs_config, + roll_obs, + done_now, + k_tx, + k_tx_planned, + k_coll_ltc, + k_coll_coex, + header_collision, + header_collision_other, + ack_collision, + ack_collision_other, + initial_ret_c, + tx_started_empty, + tx_packet_mask, + tx_success_mask, ) diff --git a/ltc/sim/sim.py b/ltc/sim/sim.py index 369c854..47d044f 100644 --- a/ltc/sim/sim.py +++ b/ltc/sim/sim.py @@ -1,16 +1,10 @@ +import jax import jax.numpy as jnp -from ltc.sim.constants import Actions +from ltc.sim.constants import Actions, EMPTY_PACKET_ID def channel_state_selector(actions): - """ - Returns a state of the channel due to how many STA transmit in the same time. - :param actions: vector of stations that transmit at a given time - :return: - int: -1 if more than one STA transmit, 1 if exactly one STA transmit, 0 if noone transmit at the moment. - """ - ones_count = jnp.sum(actions == Actions.TX.value) return jnp.where( ones_count > 1, -1, @@ -18,36 +12,98 @@ def channel_state_selector(actions): ) -def buffer_clearing(buffer_states, actions): - """ - Updates the buffer_states vector based on the action vector. - Sets the value to 0 in buffer_states where there is a 1 in both vectors. - :param buffer_states: vector with binary value of STAs' buffer occupation. - 0 - STA's buffer is empty - 1 - STA's buffer is full - :param actions: binary vector describes STAs' actions - 0 - channel sensing - 1 - transmission - :return: - jnp.ndarray: updated buffer_states. - """ - return jnp.where((buffer_states == 1) & (actions == Actions.TX.value), 0, buffer_states) - - -def add_new_frames(buffer_states, new_frames): - """ - Updates the buffer_states by adding new frames from generator. - :param buffer_states: vector with binary value of STAs' buffer occupation. - :param new_frames: vector with new frames generated for each STA. - :return: - jnp.ndarray: updated buffer_states. - """ - new_frames = (new_frames > 0).astype(int) - return jnp.bitwise_or(buffer_states, new_frames) - - -def simulate(buffer_states, new_frames, actions): +def cantor_pairing(a, b): + s = a + b + return (s * (s + 1)) // 2 + b + + +def is_buffer_empty(buffer_state): + return jnp.all(buffer_state == EMPTY_PACKET_ID) + + +def remove_packets_from_queue(queue_state, remove_mask): + n = queue_state.shape[0] + valid = queue_state != EMPTY_PACKET_ID + remove_mask = jnp.asarray(remove_mask, dtype=bool) + keep = valid & ~remove_mask + + idx = jnp.arange(n) + sort_key = jnp.where(keep, idx, n + idx) + order = jnp.argsort(sort_key) + compacted = queue_state[order] + kept_count = jnp.sum(keep.astype(jnp.int32)) + tail_mask = jnp.arange(n) >= kept_count + return jnp.where(tail_mask, EMPTY_PACKET_ID, compacted) + + +def remove_transmitted_packets(buffer_states, tx_masks): + return jax.vmap(remove_packets_from_queue)(buffer_states, tx_masks) + + +def enqueue_generated_packets(buffer_state, buffer_birth_state, new_frames, station_id, packet_seq, current_step): + q = buffer_state.shape[0] + new_frames = jnp.maximum(new_frames.astype(jnp.int32), 0) + occupied = jnp.sum((buffer_state != EMPTY_PACKET_ID).astype(jnp.int32)) + + kept_new = jnp.minimum(new_frames, q) + kept_old = jnp.minimum(occupied, q - kept_new) + old_start = occupied - kept_old + new_start = packet_seq + (new_frames - kept_new) + + idx = jnp.arange(q, dtype=jnp.int32) + + take_old = idx < kept_old + old_idx = jnp.clip(old_start + idx, 0, jnp.maximum(occupied - 1, 0)) + old_packets = jnp.where(take_old, buffer_state[old_idx], EMPTY_PACKET_ID) + old_births = jnp.where(take_old, buffer_birth_state[old_idx], -1) + + new_rel = idx - kept_old + take_new = jnp.logical_and(idx >= kept_old, new_rel < kept_new) + new_local_ids = new_start + jnp.maximum(new_rel, 0).astype(jnp.int64) + new_packets = cantor_pairing(station_id.astype(jnp.int64), new_local_ids) + new_births = jnp.full((q,), current_step, dtype=jnp.int32) + + next_state = jnp.where(take_old, old_packets, jnp.where(take_new, new_packets, EMPTY_PACKET_ID)) + next_birth = jnp.where(take_old, old_births, jnp.where(take_new, new_births, -1)) + next_packet_seq = packet_seq + new_frames + return next_state, next_birth, next_packet_seq + + +def add_new_frames(buffer_states, buffer_birth_steps, new_frames, packet_seqs, current_step): + station_ids = jnp.arange(buffer_states.shape[0], dtype=jnp.int64) + return jax.vmap(enqueue_generated_packets, in_axes=(0, 0, 0, 0, 0, None))( + buffer_states, buffer_birth_steps, new_frames, station_ids, packet_seqs, current_step + ) + + +def simulate(buffer_states, buffer_birth_steps, new_frames, actions, packet_seqs, current_step, + tx_packet_mask, tx_success_mask, active=None): + if active is None: + active = jnp.ones_like(actions, dtype=bool) + + active = jnp.asarray(active, dtype=bool) + actions = jnp.where(active, actions, Actions.IDLE.value) + new_frames = jnp.where(active, new_frames, 0) + channel_state = channel_state_selector(actions) - 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 + pre_nonempty = jnp.any(buffer_states != EMPTY_PACKET_ID, axis=1) + packet_tx = tx_success_mask & pre_nonempty & active + planned_packets = (tx_packet_mask & active).astype(jnp.int32) + successful_packets = packet_tx.astype(jnp.int32) + queue_size = buffer_states.shape[1] + head_mask = jnp.arange(queue_size) == 0 + tx_masks = packet_tx[:, None] & head_mask[None, :] + + buffer_states = remove_transmitted_packets(buffer_states, tx_masks) + buffer_birth_steps = remove_transmitted_packets(buffer_birth_steps, tx_masks) + buffer_states, buffer_birth_steps, packet_seqs = add_new_frames( + buffer_states, buffer_birth_steps, new_frames, packet_seqs, current_step + ) + return ( + buffer_states, + buffer_birth_steps, + channel_state, + packet_seqs, + planned_packets, + successful_packets, + ) diff --git a/ltc/utils/history.py b/ltc/utils/history.py index f1dc106..4b731ca 100644 --- a/ltc/utils/history.py +++ b/ltc/utils/history.py @@ -5,7 +5,7 @@ HISTORY_SUFFIX = ".pkl.lz4" -def _run_git_command(*args: str) -> str: +def run_git_command(*args: str) -> str: result = subprocess.run( ["git", *args], check=True, @@ -16,7 +16,7 @@ def _run_git_command(*args: str) -> str: def ensure_clean_git_worktree() -> None: - status = _run_git_command("status", "--porcelain", "--untracked-files=no") + status = run_git_command("status", "--porcelain", "--untracked-files=no") if status: raise RuntimeError( "Refusing to run with tracked uncommitted changes. Commit or stash changes first." @@ -24,7 +24,7 @@ def ensure_clean_git_worktree() -> None: def get_short_commit_hash() -> str: - return _run_git_command("rev-parse", "--short", "HEAD") + return run_git_command("rev-parse", "--short", "HEAD") def build_history_filename(n: int, n_drl: int, seed: int, commit_hash: str | None = None) -> str: diff --git a/ltc/utils/plot_actions.py b/ltc/utils/plot_actions.py new file mode 100644 index 0000000..91634c0 --- /dev/null +++ b/ltc/utils/plot_actions.py @@ -0,0 +1,415 @@ +import argparse + +import cloudpickle +import lz4.frame +import matplotlib.patches as mpatches +import matplotlib.pyplot as plt +import numpy as np + +from ltc.sim.constants import Actions, MAX_MACRO_DURATION +from ltc.utils.history import parse_history_filename, unpack_history + +TX_COLOR = np.array([0.2, 0.75, 0.2 ]) # green +CS_COLOR = np.array([1.0, 0.85, 0.0 ]) # yellow +IDLE_COLOR = np.array([0.92, 0.92, 0.92]) # light gray +CH_SUCCESS = np.array([0.2, 0.75, 0.2 ]) # green +CH_IDLE = np.array([1.0, 0.85, 0.0 ]) # yellow +CH_COLLISION= np.array([0.9, 0.15, 0.15]) # red + + +def _x_ticks(ax, n_steps): + ticks = np.arange(0, n_steps, max(1, n_steps // 20)) + ax.set_xticks(ticks) + ax.set_xticklabels(ticks, fontsize=7, rotation=45) + + +def _vlines(ax, n_steps, every=10): + for x in range(every, n_steps, every): + ax.axvline(x - 0.5, color='black', linewidth=0.4, alpha=0.4) + + +def decode_action_type(raw_action, action_space_variant, max_duration): + if action_space_variant in { + 'txcs_duration_set', + 'txcs_to_discrete_duration', + 'txcs_to_continuous_duration', + 'txcs_to_geometric_duration', + }: + return np.where(raw_action < max_duration, Actions.TX.value, Actions.CS.value) + if action_space_variant == 'tx_stage_commit': + return np.where(raw_action == 0, Actions.CS.value, Actions.TX.value) + return raw_action + + +def resolve_max_duration(metadata, raw_actions): + """ + Resolve max_duration for plotting. + + Older history files may not store max_duration in metadata. In that case, + infer it from observed raw actions (best-effort) instead of falling back to + current code constant, which can misdecode old runs. + """ + md = metadata or {} + md_val = md.get('max_duration') + if md_val is not None: + try: + md_int = int(md_val) + if md_int > 0: + return md_int + except (TypeError, ValueError): + pass + + # txcs_duration_set has action range [0, 2*max_duration-1]. + # Infer max_duration from observed maximum action index. + raw_max = int(np.max(raw_actions)) + return max(1, (raw_max + 2) // 2) + + +# ── 1. Actions ──────────────────────────────────────────────────────────────── + +def plot_action_slots(history, metadata, n, n_drl, seed, epoch=0, max_steps=300): + action_space_variant = (metadata or {}).get('action_space_variant', 'txcs_duration_set') + raw = np.array(history.actions[epoch, :max_steps]) # (steps, n_agents) + max_duration = resolve_max_duration(metadata, raw) + n_steps, n_agents = raw.shape + has_slot = hasattr(history, 'slot_actions') and history.slot_actions is not None + if has_slot: + types = np.array(history.slot_actions[epoch, :max_steps]).T + else: + types = decode_action_type(raw.T, action_space_variant, max_duration) + + grid = np.zeros((n_agents, n_steps, 3)) + grid[types == Actions.TX.value] = TX_COLOR + grid[types == Actions.CS.value] = CS_COLOR + grid[types == Actions.IDLE.value] = IDLE_COLOR + + fig, ax = plt.subplots(figsize=(max(14, n_steps / 30), max(2, n_agents * 0.7 + 1.2))) + ax.imshow(grid, aspect='auto', interpolation='nearest', origin='upper') + + ax.set_yticks(range(n_agents)) + ax.set_yticklabels( + [f'DRL {i}' if i < n_drl else f'Legacy {i - n_drl}' for i in range(n_agents)], + fontsize=9, + ) + _x_ticks(ax, n_steps) + ax.set_xlabel('Slot', fontsize=10) + ax.set_ylabel('Agent', fontsize=10) + ax.set_title(f'Actions — epoch {epoch}, {n_steps} slots [{action_space_variant}]', fontsize=11) + ax.legend(handles=[ + mpatches.Patch(color=TX_COLOR, label='TX'), + mpatches.Patch(color=CS_COLOR, label='CS'), + mpatches.Patch(color=IDLE_COLOR, label='IDLE'), + ], loc='upper right', fontsize=9) + + plt.tight_layout() + out = f'action_slots_{n}_{n_drl}_{seed}_e{epoch}.png' + plt.savefig(out, dpi=150, bbox_inches='tight') + print(f'Saved: {out}') + plt.show() + + +# ── 2. Channel state ────────────────────────────────────────────────────────── + +def plot_channel_slots(history, n, n_drl, seed, epoch=0, max_steps=300): + ch = np.array(history.channel_state[epoch, :max_steps]) # (steps,) values: -1, 0, 1 + n_steps = len(ch) + + grid = np.zeros((1, n_steps, 3)) + grid[0, ch == 1] = CH_SUCCESS + grid[0, ch == 0] = CH_IDLE + grid[0, ch == -1] = CH_COLLISION + + fig, ax = plt.subplots(figsize=(max(14, n_steps / 30), 1.4)) + ax.imshow(grid, aspect='auto', interpolation='nearest', origin='upper') + + ax.set_yticks([0]) + ax.set_yticklabels(['Channel'], fontsize=9) + _x_ticks(ax, n_steps) + ax.set_xlabel('Slot', fontsize=10) + ax.set_title(f'Channel state — epoch {epoch}, {n_steps} slots', fontsize=11) + ax.legend(handles=[ + mpatches.Patch(color=CH_SUCCESS, label='TX success'), + mpatches.Patch(color=CH_IDLE, label='Idle'), + mpatches.Patch(color=CH_COLLISION, label='Collision'), + ], loc='upper right', fontsize=9) + + plt.tight_layout() + out = f'channel_slots_{n}_{n_drl}_{seed}_e{epoch}.png' + plt.savefig(out, dpi=150, bbox_inches='tight') + print(f'Saved: {out}') + plt.show() + + +# ── 3. Rewards ──────────────────────────────────────────────────────────────── + +def plot_reward_slots(history, n, n_drl, seed, epoch=0, max_steps=300): + rewards = np.array(history.rewards[epoch, :max_steps]) # (steps, n_agents) + n_steps, n_agents = rewards.shape + + # diverging colormap: red=negative, white=0, green=positive + vmax = max(np.abs(rewards).max(), 1e-6) + + fig, axes = plt.subplots( + n_agents, 1, + figsize=(max(14, n_steps / 30), max(3, n_agents * 1.1 + 0.8)), + sharex=True, + ) + if n_agents == 1: + axes = [axes] + + for i, ax in enumerate(axes): + r = rewards[:, i].reshape(1, -1) + im = ax.imshow(r, aspect='auto', interpolation='nearest', + cmap='RdYlGn', vmin=-vmax, vmax=vmax, origin='upper') + label = f'DRL {i}' if i < n_drl else f'Legacy {i - n_drl}' + ax.set_yticks([0]) + ax.set_yticklabels([label], fontsize=9) + if i == n_agents - 1: + _x_ticks(ax, n_steps) + ax.set_xlabel('Slot', fontsize=10) + + fig.colorbar(im, ax=axes, orientation='vertical', fraction=0.015, pad=0.01, label='Reward') + fig.suptitle(f'Rewards per agent — epoch {epoch}, {n_steps} slots', fontsize=11, y=1.01) + + plt.tight_layout() + out = f'reward_slots_{n}_{n_drl}_{seed}_e{epoch}.png' + plt.savefig(out, dpi=150, bbox_inches='tight') + print(f'Saved: {out}') + plt.show() + + +# ── 4. Combined ─────────────────────────────────────────────────────────────── + +def plot_combined(history, metadata, n, n_drl, seed, epoch=0, max_steps=300, grid_every=10): + action_space_variant = (metadata or {}).get('action_space_variant', 'txcs_duration_set') + raw = np.array(history.actions[epoch, :max_steps]) # (steps, n_agents) + max_duration = resolve_max_duration(metadata, raw) + has_slot = hasattr(history, 'slot_actions') and history.slot_actions is not None + slot_act = np.array(history.slot_actions[epoch, :max_steps]) if has_slot else None + ch = np.array(history.channel_state[epoch, :max_steps]) # (steps,) + rewards = np.array(history.rewards[epoch, :max_steps]) # (steps, n_agents) + has_packets = hasattr(history, 'successful_packets') and history.successful_packets is not None + pkts = np.array(history.successful_packets[epoch, :max_steps]) if has_packets else None + + n_steps, n_agents = raw.shape + + # row layout: actions×n_agents | channel×1 | rewards×n_agents | packets×n_agents + n_rows = n_agents + 1 + n_agents + (n_agents if has_packets else 0) + heights = [1] * n_agents + [1] + [1] * n_agents + ([1] * n_agents if has_packets else []) + + fig, axes = plt.subplots( + n_rows, 1, + figsize=(max(16, n_steps / 25), n_rows * 0.55 + 1.0), + gridspec_kw={'height_ratios': heights, 'hspace': 0.08}, + sharex=True, + ) + + # ── actions ── (slot_actions if available, else decoded raw macro-actions) + if has_slot: + types = slot_act.T # (n_agents, n_steps) — already TX/CS/IDLE values + else: + types = decode_action_type(raw.T, action_space_variant, max_duration) + action_grid = np.zeros((n_agents, n_steps, 3)) + action_grid[types == Actions.TX.value] = TX_COLOR + action_grid[types == Actions.CS.value] = CS_COLOR + action_grid[types == Actions.IDLE.value] = IDLE_COLOR + + for i in range(n_agents): + ax = axes[i] + ax.imshow(action_grid[i:i+1], aspect='auto', interpolation='nearest', origin='upper') + _vlines(ax, n_steps, grid_every) + label = f'DRL {i}' if i < n_drl else f'Legacy {i - n_drl}' + ax.set_yticks([0]) + ax.set_yticklabels([label], fontsize=8) + ax.tick_params(left=False) + if i == 0: + ax.set_title( + f'Simulation overview — epoch {epoch}, {n_steps} slots [{action_space_variant}]', + fontsize=11, pad=6, + ) + + # divider label + axes[n_agents - 1].spines['bottom'].set_linewidth(1.5) + + # ── channel ── + ch_grid = np.zeros((1, n_steps, 3)) + ch_grid[0, ch == 1] = CH_SUCCESS + ch_grid[0, ch == 0] = CH_IDLE + ch_grid[0, ch == -1] = CH_COLLISION + + ax_ch = axes[n_agents] + ax_ch.imshow(ch_grid, aspect='auto', interpolation='nearest', origin='upper') + _vlines(ax_ch, n_steps, grid_every) + ax_ch.set_yticks([0]) + ax_ch.set_yticklabels(['Channel'], fontsize=8) + ax_ch.tick_params(left=False) + ax_ch.spines['bottom'].set_linewidth(1.5) + + # ── rewards ── + vmax = max(np.abs(rewards).max(), 1e-6) + reward_ims = [] + for i in range(n_agents): + ax = axes[n_agents + 1 + i] + r = rewards[:, i].reshape(1, -1) + im = ax.imshow(r, aspect='auto', interpolation='nearest', + cmap='RdYlGn', vmin=-vmax, vmax=vmax, origin='upper') + _vlines(ax, n_steps, grid_every) + reward_ims.append(im) + label = f'R DRL {i}' if i < n_drl else f'R Leg {i - n_drl}' + ax.set_yticks([0]) + ax.set_yticklabels([label], fontsize=8) + ax.tick_params(left=False) + + fig.colorbar(reward_ims[-1], ax=axes[n_agents + 1:n_agents + 1 + n_agents], orientation='vertical', + fraction=0.012, pad=0.01, label='Reward') + + # ── packets ── + if has_packets: + PKT_COLOR = np.array([0.2, 0.75, 0.2]) + NO_PKT_COLOR = np.array([0.93, 0.93, 0.93]) + for i in range(n_agents): + ax = axes[n_agents + 1 + n_agents + i] + pkt_grid = np.where(pkts[:, i:i+1].T == 1, 1.0, 0.0) + rgb = np.where(pkts[:, i:i+1].T[..., None] == 1, PKT_COLOR, NO_PKT_COLOR) + ax.imshow(rgb.reshape(1, n_steps, 3), aspect='auto', interpolation='nearest', origin='upper') + _vlines(ax, n_steps, grid_every) + label = f'Pkt DRL {i}' if i < n_drl else f'Pkt Leg {i - n_drl}' + ax.set_yticks([0]) + ax.set_yticklabels([label], fontsize=8) + ax.tick_params(left=False) + + # ── x ticks only on last row ── + _x_ticks(axes[-1], n_steps) + axes[-1].set_xlabel('Slot', fontsize=10) + + # ── shared legend (actions + channel) ── + legend_handles = [ + mpatches.Patch(color=TX_COLOR, label='TX'), + mpatches.Patch(color=CS_COLOR, label='CS'), + mpatches.Patch(color=IDLE_COLOR, label='IDLE'), + mpatches.Patch(color=CH_SUCCESS, label='Ch: success'), + mpatches.Patch(color=CH_IDLE, label='Ch: idle'), + mpatches.Patch(color=CH_COLLISION, label='Ch: collision'), + ] + if has_packets: + legend_handles.append(mpatches.Patch(color=np.array([0.2, 0.75, 0.2]), label='Packet sent')) + legend_handles.append(mpatches.Patch(color=np.array([0.93, 0.93, 0.93]), label='No packet')) + fig.legend(handles=legend_handles, loc='upper right', fontsize=8, ncol=2, bbox_to_anchor=(0.88, 1.0)) + + out = f'overview_{n}_{n_drl}_{seed}_e{epoch}.png' + plt.savefig(out, dpi=150, bbox_inches='tight') + print(f'Saved: {out}') + plt.show() + + +# ── 5. Raw data inspector ───────────────────────────────────────────────────── + +def inspect(history, metadata, n, n_drl, seed, epoch=0, max_steps=20): + from ltc.sim.constants import ObsIdx + + n_epochs = history.actions.shape[0] + n_steps_total = history.actions.shape[1] + n_agents = history.actions.shape[2] + + print("=" * 70) + print(f"FILE: n={n} n_drl={n_drl} seed={seed}") + print(f"SHAPE: {n_epochs} epochs × {n_steps_total} steps × {n_agents} agents") + print(f"ARGS: {metadata}") + print("=" * 70) + + raw = np.array(history.actions[epoch, :max_steps]) # (steps, agents) + ch = np.array(history.channel_state[epoch, :max_steps]) # (steps,) + rewards = np.array(history.rewards[epoch, :max_steps]) # (steps, agents) + obs = np.array(history.observations[epoch, :max_steps]) # (steps, agents, win, obs_size) + buf = np.array(history.buffer_states[epoch, :max_steps]) # (steps, agents, queue) + has_pkts = hasattr(history, 'successful_packets') and history.successful_packets is not None + pkts_raw = np.array(history.successful_packets[epoch, :max_steps]) if has_pkts else None + has_slot = hasattr(history, 'slot_actions') and history.slot_actions is not None + slot_raw = np.array(history.slot_actions[epoch, :max_steps]) if has_slot else None + + action_space_variant = (metadata or {}).get('action_space_variant', 'txcs_duration_set') + max_duration = resolve_max_duration(metadata, raw) + macro_types = decode_action_type(raw.T, action_space_variant, max_duration).T # (steps, agents) + + action_names = {Actions.TX.value: 'TX', Actions.CS.value: 'CS', Actions.IDLE.value: 'IDLE'} + ch_names = {1: 'SUCCESS', 0: 'IDLE', -1: 'COLLISION'} + + print(f"\n── Epoch {epoch}, first {max_steps} slots ──\n") + col_hdr = "raw/macro/slot/rew/pkt" if has_pkts else "raw/macro/slot/rew" + header = f"{'Slot':>5} {'Channel':>10} " + " ".join(f"A{i}({col_hdr})" for i in range(n_agents)) + print(header) + print("-" * len(header)) + + for t in range(min(max_steps, n_steps_total)): + ch_str = ch_names.get(int(ch[t]), str(ch[t])) + agent_cols = [] + for i in range(n_agents): + macro_str = action_names.get(int(macro_types[t, i]), str(macro_types[t, i])) + slot_str = action_names.get(int(slot_raw[t, i]), '?') if has_slot else '?' + pkt_part = f"/{int(pkts_raw[t,i])}" if has_pkts else "" + agent_cols.append(f"{int(raw[t,i]):>3}/{macro_str:<4}/{slot_str:<4}/{rewards[t,i]:>+6.2f}{pkt_part}") + print(f"{t:>5} {ch_str:>10} " + " ".join(agent_cols)) + + print(f"\n── Observations (last obs window, epoch {epoch}, first {max_steps} steps) ──\n") + obs_names = [name for name, _ in sorted(vars(ObsIdx).items(), key=lambda x: x[1]) + if not name.startswith('_')] + for t in range(min(max_steps, n_steps_total)): + print(f" Step {t}:") + for i in range(n_agents): + label = f'DRL {i}' if i < n_drl else f'Legacy {i-n_drl}' + vals = obs[t, i, -1] # last window entry + obs_str = " ".join(f"{obs_names[j]}={vals[j]:.2f}" for j in range(len(obs_names))) + print(f" [{label}] {obs_str}") + + print(f"\n── Buffer states (epoch {epoch}, first {max_steps} steps) ──\n") + for t in range(min(max_steps, n_steps_total)): + bufs = " ".join(f"A{i}:{list(buf[t,i])}" for i in range(n_agents)) + print(f" Step {t}: {bufs}") + + print("\n── Summary statistics ──\n") + has_pkts_hist = hasattr(history, 'successful_packets') and history.successful_packets is not None + for i in range(n_agents): + label = f'DRL {i}' if i < n_drl else f'Legacy {i-n_drl}' + all_types = decode_action_type( + np.array(history.actions[:, :, i]).flatten(), + action_space_variant, max_duration, + ) + tx_pct = 100 * np.mean(all_types == Actions.TX.value) + cs_pct = 100 * np.mean(all_types == Actions.CS.value) + total_r = float(np.array(history.rewards[:, :, i]).sum()) + mean_r = float(np.array(history.rewards[:, :, i]).mean()) + pkt_str = '' + if has_pkts_hist: + total_pkts = int(np.array(history.successful_packets[:, :, i]).sum()) + pkt_str = f' packets={total_pkts}' + print(f" {label:12s} TX={tx_pct:.1f}% CS={cs_pct:.1f}% " + f"total_reward={total_r:.1f} mean_reward={mean_r:.4f}{pkt_str}") + + ch_all = np.array(history.channel_state).flatten() + print(f"\n Channel: success={100*np.mean(ch_all==1):.1f}% " + f"idle={100*np.mean(ch_all==0):.1f}% " + f"collision={100*np.mean(ch_all==-1):.1f}%") + + +# ── CLI ─────────────────────────────────────────────────────────────────────── + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description='Analyze simulation history file') + parser.add_argument('filename', help='Path to .pkl.lz4 history file') + parser.add_argument('--epoch', type=int, default=0, help='Epoch index to visualize') + parser.add_argument('--steps', type=int, default=300, help='Max slots to show') + parser.add_argument('--grid', type=int, default=10, help='Vertical grid every N slots (0=off)') + parser.add_argument('--inspect', action='store_true', help='Print raw data instead of plotting') + parser.add_argument('--inspect_steps', type=int, default=20, help='Slots to print in inspect mode') + args = parser.parse_args() + + n, n_drl, seed, _ = parse_history_filename(args.filename) + + with lz4.frame.open(args.filename, 'rb') as f: + _, history, metadata = unpack_history(cloudpickle.load(f)) + + if args.inspect: + inspect(history, metadata, n, n_drl, seed, epoch=args.epoch, max_steps=args.inspect_steps) + else: + plot_combined(history, metadata, n, n_drl, seed, epoch=args.epoch, max_steps=args.steps, + grid_every=args.grid if args.grid > 0 else 10**9) diff --git a/ltc/utils/plots.py b/ltc/utils/plots.py index 913c689..188ad92 100644 --- a/ltc/utils/plots.py +++ b/ltc/utils/plots.py @@ -8,9 +8,9 @@ import matplotlib.pyplot as plt import numpy as np -from ltc.utils.scan_states import Output +from ltc.utils.structs import Output from ltc.utils.history import parse_history_filename, unpack_history -from ltc.sim.constants import NO_TX_REWARD, Actions, INITIAL_CAPACITY +from ltc.sim.constants import NO_TX_REWARD, Actions, INITIAL_CAPACITY, EMPTY_PACKET_ID SAVE_FORMAT = 'normal' @@ -304,14 +304,15 @@ def plot_buffer_states(buffer_states, terminals, n, n_drl, seed, name): else: xs = np.linspace(0, n_epochs * n_steps, n_epochs) + 1 - buffer_states = np.where(1 - terminals, buffer_states, 0).sum(axis=1) - buffer_states = buffer_states / (1 - terminals).sum(axis=1) + queue_fill = np.count_nonzero(buffer_states != EMPTY_PACKET_ID, axis=-1) / buffer_states.shape[-1] + queue_fill = np.where(1 - terminals, queue_fill, 0).sum(axis=1) + queue_fill = queue_fill / (1 - terminals).sum(axis=1) for i in range(n_drl): - plt.plot(xs, buffer_states[:, i], color='red') + plt.plot(xs, queue_fill[:, i], color='red') for i in range(n_drl, n): - plt.plot(xs, buffer_states[:, i], color='blue', linestyle='--') + plt.plot(xs, queue_fill[:, i], color='blue', linestyle='--') plt.plot([], color='blue', linestyle='--', label='Legacy') plt.plot([], color='red', label='DRL') @@ -507,8 +508,9 @@ def plot_channel_access_delay(buffer_states, new_frames, terminals, n, n_drl, se else: xs = np.linspace(0, n_epochs * n_steps, n_epochs) + 1 - buffer_states, new_frames = (1 - terminals) * buffer_states, (1 - terminals) * new_frames - cad = buffer_states.sum(axis=1) / (new_frames.sum(axis=1) + 1e-6) + queue_fill = np.count_nonzero(buffer_states != EMPTY_PACKET_ID, axis=-1) + queue_fill, new_frames = (1 - terminals) * queue_fill, (1 - terminals) * new_frames + cad = queue_fill.sum(axis=1) / (new_frames.sum(axis=1) + 1e-6) for i in range(n_drl): plt.plot(xs, cad[:, i], color='red') @@ -539,13 +541,15 @@ def plot_xnor(actions, channel_states, buffer_states, terminals, n, n_drl, seed, else: xs = np.linspace(0, n_epochs * n_steps, n_epochs) + 1 + empty = np.all(buffer_states == EMPTY_PACKET_ID, axis=-1) + non_empty = ~empty idle_and_empty = jax.lax.bitwise_and( jax.lax.bitwise_not(terminals), - jax.lax.bitwise_and(actions == Actions.IDLE.value, buffer_states == 0) + jax.lax.bitwise_and(actions == Actions.IDLE.value, empty) ).sum(axis=1) tx_and_full = jax.lax.bitwise_and( jax.lax.bitwise_and(jax.lax.bitwise_not(terminals), np.expand_dims(channel_states, axis=-1) == 1), - jax.lax.bitwise_and(actions == Actions.TX.value, buffer_states != 0) + jax.lax.bitwise_and(actions == Actions.TX.value, non_empty) ).sum(axis=1) xnor = idle_and_empty + tx_and_full diff --git a/ltc/utils/scan_states.py b/ltc/utils/scan_states.py deleted file mode 100644 index 51bdf93..0000000 --- a/ltc/utils/scan_states.py +++ /dev/null @@ -1,40 +0,0 @@ -from dataclasses import dataclass - -import jax -from reinforced_lib.agents import AgentState - -from ltc.sim import ModelState - - -@jax.tree_util.register_dataclass -@dataclass -class Carry: - drl_states: AgentState - legacy_states: AgentState - traffic_states: ModelState - buffer_states: jax.Array - power_states: jax.Array - channel_state: int - key: jax.random.PRNGKey - obs: jax.Array - actions: jax.Array - rewards: jax.Array - terminals: jax.Array - active: jax.Array - - -@jax.tree_util.register_dataclass -@dataclass -class Output: - legacy_states: AgentState - observations: jax.Array - actions: jax.Array - rewards: jax.Array - terminals: jax.Array - buffer_states: jax.Array - power_states: jax.Array - new_frames: jax.Array - channel_state: int - active: jax.Array - weights_histogram: jax.Array - weights_bin_edges: jax.Array diff --git a/ltc/utils/structs.py b/ltc/utils/structs.py new file mode 100644 index 0000000..71d0059 --- /dev/null +++ b/ltc/utils/structs.py @@ -0,0 +1,101 @@ +from dataclasses import dataclass +from typing import Any +import jax +from reinforced_lib.agents import AgentState + + +@jax.tree_util.register_dataclass +@dataclass +class MacroActionState: + remaining: jax.Array + action_types: jax.Array + tx_success_accum: jax.Array # k_tx: successful subframes + tx_planned_accum: jax.Array # k_tx_planned: planned subframes (= TX macro duration) + k_coll_ltc: jax.Array # collided subframes vs LTC + k_coll_coex: jax.Array # collided subframes vs coex + sub_collision_flag: jax.Array # True if any collision in current TX_SLOTS sub-window + sub_collision_other_flag: jax.Array # True if any coex collision in current TX_SLOTS sub-window + header_collision: jax.Array # True if first mini-slot of macro TX collided + header_collision_other: jax.Array # True if the header collision was with a non-LTC station + initial_ret_c: jax.Array # ret_c captured at macro TX start + tx_started_empty: jax.Array # True if buffer was empty when macro TX began + + +@jax.tree_util.register_dataclass +@dataclass +class ObsTrackerState: + buffer_birth_steps: jax.Array + arrival_hist: jax.Array + planned_tx_hist: jax.Array + success_tx_hist: jax.Array + channel_busy_hist: jax.Array + collision_hist: jax.Array + staged_tx: jax.Array + + +@jax.tree_util.register_dataclass +@dataclass +class ObsFeatureInputs: + traffic_mean_arrival_rate: jax.Array + channel_occupancy_pct_window: jax.Array + channel_collisions_pct_window: jax.Array + back_pct: jax.Array + unique_ltc_tx_window: jax.Array + cs_tx_same_type_now: jax.Array + tx_collision_other_now: jax.Array + + +@jax.tree_util.register_dataclass +@dataclass +class ObsFeatureConfig: + enable_cs_tx_same_type: bool + enable_tx_collision_other: bool + + +@jax.tree_util.register_dataclass +@dataclass +class ActionDecodingResults: + drl_action_types: jax.Array + drl_durations: jax.Array + drl_staged_next: jax.Array + legacy_action_types: jax.Array + legacy_durations: jax.Array + + +@jax.tree_util.register_dataclass +@dataclass +class Carry: + drl_states: AgentState + legacy_states: AgentState + traffic_states: Any + packet_seqs: jax.Array + buffer_states: jax.Array + tx_hist: jax.Array + power_states: jax.Array + channel_state: int + key: jax.random.PRNGKey + obs: jax.Array + actions: jax.Array + rewards: jax.Array + terminals: jax.Array + active: jax.Array + macro: MacroActionState + obs_tracker: ObsTrackerState + + +@jax.tree_util.register_dataclass +@dataclass +class Output: + legacy_states: AgentState + observations: jax.Array + actions: jax.Array + rewards: jax.Array + terminals: jax.Array + buffer_states: jax.Array + power_states: jax.Array + new_frames: jax.Array + channel_state: int + active: jax.Array + weights_histogram: jax.Array + weights_bin_edges: jax.Array + successful_packets: jax.Array diff --git a/pyproject.toml b/pyproject.toml index 93dfc63..470fd8c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "ltc" -version = "0.1.0" +version = "0.2.0" requires-python = ">=3.12, <4" description = "Learning to communicate" diff --git a/test/sim/test_action_space.py b/test/sim/test_action_space.py new file mode 100644 index 0000000..141b2fe --- /dev/null +++ b/test/sim/test_action_space.py @@ -0,0 +1,94 @@ +import unittest + +import jax.numpy as jnp + +from ltc.run import decode_drl_actions, decode_legacy_actions +from ltc.sim.constants import Actions, TX_REWARD, TX_SLOTS, TX_LTC_COLLISION_PENALTY, TX_COEX_COLLISION_PENALTY, TX_MAX_RETRANSMISSION_PENALTY, TX_EMPTY_BUFFER_PENALTY, MAX_RETRANSMISSION, TX_SIZE_THRESHOLD, TX_SIZE_PENALTY, TX_SIZE_PENALTY_WINDOW +from ltc.sim.process_output import tx_macro_reward + + +class ActionSpaceDecodeTestCase(unittest.TestCase): + def test_duration_set_decode(self): + raw = jnp.array([0, 4, 5, 9], dtype=jnp.int32) + action_type, duration, staged = decode_drl_actions( + raw, jnp.zeros((4,), dtype=jnp.int32), "txcs_duration_set", 5 + ) + self.assertTrue(jnp.array_equal(action_type, jnp.array([Actions.TX.value, Actions.TX.value, Actions.CS.value, Actions.CS.value]))) + self.assertTrue(jnp.array_equal(duration, jnp.array([1, 5, 1, 5]))) + self.assertTrue(jnp.array_equal(staged, jnp.zeros((4,), dtype=jnp.int32))) + + def test_stage_commit_decode(self): + raw = jnp.array([1, 2, 2], dtype=jnp.int32) # stage, commit, commit + staged = jnp.array([0, 2, 1], dtype=jnp.int32) + action_type, duration, staged_next = decode_drl_actions( + raw, staged, "tx_stage_commit", 5 + ) + self.assertTrue(jnp.array_equal(action_type, jnp.array([Actions.CS.value, Actions.TX.value, Actions.TX.value]))) + self.assertTrue(jnp.array_equal(duration, jnp.array([1, 2, 1]))) + self.assertTrue(jnp.array_equal(staged_next, jnp.array([1, 0, 0]))) + + def test_legacy_decode(self): + raw = jnp.array([Actions.TX.value, Actions.CS.value, Actions.IDLE.value], dtype=jnp.int32) + action_type, duration = decode_legacy_actions(raw, legacy_tx_duration=5) + self.assertTrue(jnp.array_equal(action_type, raw)) + self.assertTrue(jnp.array_equal(duration, jnp.array([5, 1, 1], dtype=jnp.int32))) + + def _macro_reward(self, **kwargs): + defaults = dict( + k_tx=jnp.array(0.0), k_tx_planned=jnp.array(0.0), + k_coll_ltc=jnp.array(0.0), k_coll_coex=jnp.array(0.0), + tx_started_empty=jnp.array(False), initial_ret_c=jnp.array(0), + header_collision=jnp.array(False), ack_collision=jnp.array(False), + header_collision_other=jnp.array(False), ack_collision_other=jnp.array(False), + ) + defaults.update(kwargs) + return tx_macro_reward(**defaults) + + def test_tx_macro_reward_branches(self): + # success: 1 sub-window, no collision + r = self._macro_reward(k_tx=jnp.array(1.0), k_tx_planned=jnp.array(1.0)) + expected = TX_REWARD * (1.0 * TX_SLOTS - 2.0) / TX_SLOTS + self.assertAlmostEqual(float(r), expected, places=6) + + # coex collision, ret_c < MAX: 1 collided sub-window → full coex penalty + r = self._macro_reward(k_coll_coex=jnp.array(1.0), k_tx_planned=jnp.array(1.0)) + self.assertAlmostEqual(float(r), TX_COEX_COLLISION_PENALTY * 1.0, places=6) + + # LTC collision + max retransmission: collision penalty + max-retry penalty (additive) + r = self._macro_reward( + k_coll_ltc=jnp.array(1.0), k_tx_planned=jnp.array(1.0), + initial_ret_c=jnp.array(MAX_RETRANSMISSION), + ) + expected = TX_LTC_COLLISION_PENALTY * 1.0 + TX_MAX_RETRANSMISSION_PENALTY + self.assertAlmostEqual(float(r), expected, places=6) + + # empty buffer, no collision: penalty scales with planned duration + r = self._macro_reward(tx_started_empty=jnp.array(True), k_tx_planned=jnp.array(2.0)) + self.assertAlmostEqual(float(r), TX_EMPTY_BUFFER_PENALTY * 2.0, places=6) + + # header collision (LTC): all 3 planned sub-windows → 3 LTC collisions + # size penalty applies: k_tx_planned=3, ramp=(3-2)/2=0.5 + r = self._macro_reward( + k_tx=jnp.array(3.0), k_tx_planned=jnp.array(3.0), + header_collision=jnp.array(True), header_collision_other=jnp.array(False), + ) + expected = TX_LTC_COLLISION_PENALTY * 3.0 + TX_SIZE_PENALTY * 0.5 + self.assertAlmostEqual(float(r), expected, places=6) + + # header collision (coex): all 2 planned sub-windows → 2 coex collisions + r = self._macro_reward( + k_tx=jnp.array(2.0), k_tx_planned=jnp.array(2.0), + header_collision=jnp.array(True), header_collision_other=jnp.array(True), + ) + self.assertAlmostEqual(float(r), TX_COEX_COLLISION_PENALTY * 2.0, places=6) + + # ACK collision (LTC), 2 planned sub-windows: all flagged → 2 LTC collisions + r = self._macro_reward( + k_tx=jnp.array(1.0), k_coll_ltc=jnp.array(1.0), k_tx_planned=jnp.array(2.0), + ack_collision=jnp.array(True), ack_collision_other=jnp.array(False), + ) + self.assertAlmostEqual(float(r), TX_LTC_COLLISION_PENALTY * 2.0, places=6) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/sim/test_process_output.py b/test/sim/test_process_output.py index 963cc4f..f77015a 100644 --- a/test/sim/test_process_output.py +++ b/test/sim/test_process_output.py @@ -5,85 +5,149 @@ from ltc.sim.constants import * from ltc.sim.process_output import * +from ltc.utils.structs import ObsFeatureInputs, ObsFeatureConfig KEY = jax.random.PRNGKey(0) +EMPTY = EMPTY_PACKET_ID class ProcessOutputTestCase(unittest.TestCase): def test_no_transmission(self): - args = (Actions.IDLE.value, 1, 5, 0, 1, KEY) + args = (Actions.IDLE.value, jnp.array([1, EMPTY, EMPTY]), jnp.array(5, dtype=jnp.int64), 0, jnp.array(1, dtype=jnp.int64), KEY) reward, r, no_tx = no_transmission(args) self.assertEqual(reward, 0.0) self.assertEqual(r, 5) - def test_transmission_without_collision_empty_buffer(self): - args = (Actions.TX.value, 0, 3, 1, 1, KEY) - reward, r, no_tx = transmission_without_collision(args) - self.assertEqual(reward, EMPTY_TX_PENALTY) - self.assertEqual(r, 0) - - def test_transmission_without_collision_successful(self): - args = (Actions.TX.value, 1, 2, 1, 1, KEY) - reward, r, no_tx = transmission_without_collision(args) - self.assertAlmostEqual(reward, TX_REWARD) - self.assertEqual(r, 0) + def _make_tx_process_output(self, ret_c_init, channel_state, tx_packet_mask, tx_success_mask): + """Helper: one TX agent, one step, returns obs ret_c.""" + buf = jnp.array([[5, EMPTY, EMPTY]]) + obs = jnp.zeros((1, 3, OBS_SIZE), dtype=jnp.float32) + obs = obs.at[0, -1, ObsIdx.STATUS_RETRY_COUNTER].set(ret_c_init) + obs_features = ObsFeatureInputs( + traffic_mean_arrival_rate=jnp.array([1.0]), + channel_occupancy_pct_window=jnp.array([0.0]), + channel_collisions_pct_window=jnp.array([0.0]), + back_pct=jnp.array([0.0]), + unique_ltc_tx_window=jnp.array([0.0]), + cs_tx_same_type_now=jnp.array([0.0]), + tx_collision_other_now=jnp.array([0.0]), + ) + obs_config = ObsFeatureConfig(enable_cs_tx_same_type=False, enable_tx_collision_other=False) + result_obs, _, _ = process_output( + buffer_states=buf, new_buffer_states=buf, new_buffer_birth_states=jnp.array([[-1,-1,-1]]), + power_states=jnp.array([100]), channel_state=channel_state, + obs=obs, actions=jnp.array([Actions.TX.value]), terminals=jnp.array([False]), + key=KEY, current_step=0, obs_features=obs_features, obs_config=obs_config, + tx_packet_mask=jnp.array([tx_packet_mask]), + tx_success_mask=jnp.array([tx_success_mask]), + ) + return int(result_obs[0, -1, ObsIdx.STATUS_RETRY_COUNTER]) - def test_transmission_with_collision_retransmission(self): - args = (Actions.TX.value, 1, 2, -1, 1, KEY) - reward, r, no_tx = transmission_with_collision(args) - self.assertEqual(reward, COLLISION_PENALTY) + def test_tx_sub_window_collision_increments_ret_c(self): + r = self._make_tx_process_output(ret_c_init=2, channel_state=-1, tx_packet_mask=True, tx_success_mask=False) self.assertEqual(r, 3) - def test_transmission_with_collision_max_retransmission(self): - args = (Actions.TX.value, 1, MAX_RETRANSMISSION, -1, 1, KEY) - reward, r, no_tx = transmission_with_collision(args) - self.assertEqual(reward, MAX_RETRANSMISSION_PENALTY) + def test_tx_sub_window_collision_at_max_resets_ret_c(self): + r = self._make_tx_process_output(ret_c_init=MAX_RETRANSMISSION, channel_state=-1, tx_packet_mask=True, tx_success_mask=False) self.assertEqual(r, 0) - def test_transmission_collision_path(self): - args = (Actions.TX.value, 1, 7, -1, 1, KEY) - reward, r, no_tx = transmission(args) - self.assertEqual(reward, COLLISION_PENALTY) + def test_tx_sub_window_collision_below_max(self): + r = self._make_tx_process_output(ret_c_init=7, channel_state=-1, tx_packet_mask=True, tx_success_mask=False) self.assertEqual(r, 8) - def test_transmission_successful_path(self): - args = (Actions.TX.value, 1, 1, 1, 1, KEY) - reward, r, no_tx = transmission(args) - self.assertAlmostEqual(reward, TX_REWARD) + def test_tx_sub_window_success_resets_ret_c(self): + r = self._make_tx_process_output(ret_c_init=3, channel_state=1, tx_packet_mask=True, tx_success_mask=True) self.assertEqual(r, 0) - def test_transmission_empty_buffer(self): - args = (Actions.TX.value, 0, 1, 1, 1, KEY) - reward, r, no_tx = transmission(args) - self.assertEqual(reward, EMPTY_TX_PENALTY) - self.assertEqual(r, 0) + def test_tx_mid_macro_ret_c_unchanged(self): + r = self._make_tx_process_output(ret_c_init=5, channel_state=-1, tx_packet_mask=False, tx_success_mask=False) + self.assertEqual(r, 5) def test_process_rl_output(self): - buffer_states = jnp.array([0, 0, 1, 0]) - new_buffer_states = jnp.array([0, 0, 1, 0]) + buffer_states = jnp.array([ + [EMPTY, EMPTY, EMPTY], + [EMPTY, EMPTY, EMPTY], + [9, EMPTY, EMPTY], + [EMPTY, EMPTY, EMPTY], + ]) + new_buffer_birth_states = jnp.array([ + [-1, -1, -1], + [-1, -1, -1], + [3, -1, -1], + [-1, -1, -1], + ], dtype=jnp.int32) + new_buffer_states = jnp.array([ + [EMPTY, EMPTY, EMPTY], + [EMPTY, EMPTY, EMPTY], + [9, EMPTY, EMPTY], + [EMPTY, EMPTY, EMPTY], + ]) power_states = jnp.array([100, 100, 100, 100]) actions = jnp.array([Actions.TX.value, Actions.CS.value, Actions.TX.value, Actions.TX.value]) terminals = jnp.array([False, False, False, False]) channel_state = -1 - obs = jnp.array([[ - [1, 1, 6, 0, 100], - [1, 1, 7, 0, 100], - [1, 1, 8, 0, 100] - ]] * 4) - - expected_obs_0 = jnp.array([ - [1, 1, 7, 0, 100], - [1, 1, 8, 0, 100], - [0, -1, 0, 0, Actions.TX.value], - ]) - expected_R_0 = -1.0 + obs = jnp.zeros((4, 3, OBS_SIZE), dtype=jnp.float32) + obs = obs.at[:, -1, ObsIdx.STATUS_RETRY_COUNTER].set(8) + expected_R_0 = TX_COEX_COLLISION_PENALTY * 1.0 + TX_MAX_RETRANSMISSION_PENALTY # -5.0 + -1.0 expected_power = power_states + jnp.array([-TX_CONSUMPTION, -CS_CONSUMPTION, -TX_CONSUMPTION, -TX_CONSUMPTION]) + obs_features = ObsFeatureInputs( + traffic_mean_arrival_rate=jnp.array([1.0, 2.0, 3.0, 4.0], dtype=jnp.float32), + channel_occupancy_pct_window=jnp.full((4,), 0.25, dtype=jnp.float32), + channel_collisions_pct_window=jnp.full((4,), 0.5, dtype=jnp.float32), + back_pct=jnp.array([0.0, 0.6, 1.0, 0.0], dtype=jnp.float32), + unique_ltc_tx_window=jnp.full((4,), 2.0, dtype=jnp.float32), + cs_tx_same_type_now=jnp.array([0.0, 1.0, 0.0, 0.0], dtype=jnp.float32), + tx_collision_other_now=jnp.array([1.0, 0.0, 1.0, 1.0], dtype=jnp.float32), + ) + obs_config = ObsFeatureConfig( + enable_cs_tx_same_type=True, + enable_tx_collision_other=True, + ) + + # Agent 0: TX, collision, ret_c=8 >= MAX_RETRANSMISSION, coex collision → coex penalty + max-retry penalty + done_now = jnp.array([True, True, True, True]) + k_tx = jnp.zeros(4, dtype=jnp.float32) + k_coll_ltc = jnp.zeros(4, dtype=jnp.float32) + k_coll_coex = jnp.array([1.0, 0.0, 1.0, 1.0], dtype=jnp.float32) # coex collision for TX agents + header_collision = jnp.zeros(4, dtype=bool) + header_collision_other = jnp.zeros(4, dtype=bool) + ack_collision = jnp.zeros(4, dtype=bool) + ack_collision_other = jnp.zeros(4, dtype=bool) + initial_ret_c = jnp.array([8, 0, 8, 8], dtype=jnp.int32) + tx_started_empty = jnp.array([True, False, False, True], dtype=bool) + result_obs, result_R, power = process_output( - buffer_states, new_buffer_states, power_states, channel_state, obs, actions, terminals, KEY + buffer_states=buffer_states, + new_buffer_states=new_buffer_states, + new_buffer_birth_states=new_buffer_birth_states, + power_states=power_states, + channel_state=channel_state, + obs=obs, + actions=actions, + terminals=terminals, + key=KEY, + current_step=10, + obs_features=obs_features, + obs_config=obs_config, + done_now=done_now, + k_tx=k_tx, + k_coll_ltc=k_coll_ltc, + k_coll_coex=k_coll_coex, + header_collision=header_collision, + header_collision_other=header_collision_other, + ack_collision=ack_collision, + ack_collision_other=ack_collision_other, + initial_ret_c=initial_ret_c, + tx_started_empty=tx_started_empty, ) - self.assertTrue(jnp.array_equal(result_obs[0], expected_obs_0)) + self.assertEqual(result_obs.shape[-1], OBS_SIZE) + self.assertEqual(result_obs[1, -1, ObsIdx.CHANNEL_LAST_CS_BUSY], 1.0) + self.assertEqual(result_obs[1, -1, ObsIdx.CHANNEL_LAST_CS_TX_SAME_TYPE], -1.0) + self.assertEqual(result_obs[0, -1, ObsIdx.CHANNEL_LAST_TX_COLLISION_OTHER], 1.0) + self.assertEqual(result_obs[2, -1, ObsIdx.BUFFER_PACKET_COUNT], 1.0) + self.assertEqual(result_obs[2, -1, ObsIdx.ACTION_BACK_PCT], 1.0) self.assertEqual(result_R[0], expected_R_0) self.assertTrue(jnp.array_equal(power, expected_power)) diff --git a/test/sim/test_sim.py b/test/sim/test_sim.py index 065caef..96bbbee 100644 --- a/test/sim/test_sim.py +++ b/test/sim/test_sim.py @@ -2,11 +2,13 @@ import jax.numpy as jnp -from ltc.sim.constants import Actions +from ltc.sim.constants import Actions, EMPTY_PACKET_ID from ltc.sim.sim import * class SimulateTestCase(unittest.TestCase): + QUEUE = 4 + def test_channel_state_selector(self): # Test no actions actions = jnp.array([Actions.CS.value, Actions.CS.value, Actions.CS.value, Actions.CS.value]) @@ -20,72 +22,87 @@ def test_channel_state_selector(self): actions = jnp.array([Actions.TX.value, Actions.TX.value, Actions.CS.value, Actions.CS.value]) self.assertEqual(channel_state_selector(actions), -1) - def test_buffer_clearing(self): - # Test clearing buffers when actions match - buffer_states = jnp.array([1, 0, 1, 0]) - actions = jnp.array([Actions.TX.value, Actions.CS.value, Actions.CS.value, Actions.CS.value]) - expected = jnp.array([0, 0, 1, 0]) - result = buffer_clearing(buffer_states, actions) - self.assertTrue(jnp.array_equal(result, expected)) - - # Test no clearing when no matching actions - buffer_states = jnp.array([1, 0, 1, 0]) - actions = jnp.array([Actions.CS.value, Actions.CS.value, Actions.CS.value, Actions.CS.value]) - expected = jnp.array([1, 0, 1, 0]) - result = buffer_clearing(buffer_states, actions) + def test_remove_transmitted_packets(self): + buffer_states = jnp.array([ + [10, 11, 12, EMPTY_PACKET_ID], + [20, 21, EMPTY_PACKET_ID, EMPTY_PACKET_ID], + ]) + tx_masks = jnp.array([ + [1, 0, 1, 0], # remove 1st and 3rd + [0, 1, 0, 0], # remove 2nd + ]) + expected = jnp.array([ + [11, EMPTY_PACKET_ID, EMPTY_PACKET_ID, EMPTY_PACKET_ID], + [20, EMPTY_PACKET_ID, EMPTY_PACKET_ID, EMPTY_PACKET_ID], + ]) + result = remove_transmitted_packets(buffer_states, tx_masks) self.assertTrue(jnp.array_equal(result, expected)) def test_add_new_frames(self): - # Test adding new frames - buffer_states = jnp.array([1, 0, 1, 0]) - new_frames = jnp.array([0, 1, 0, 0]) - expected = jnp.array([1, 1, 1, 0]) - result = add_new_frames(buffer_states, new_frames) - self.assertTrue(jnp.array_equal(result, expected)) - - # Test no new frames added - buffer_states = jnp.array([1, 0, 1, 0]) - new_frames = jnp.array([0, 0, 0, 0]) - expected = jnp.array([1, 0, 1, 0]) - result = add_new_frames(buffer_states, new_frames) - self.assertTrue(jnp.array_equal(result, expected)) + buffer_states = jnp.array([ + [EMPTY_PACKET_ID, EMPTY_PACKET_ID, EMPTY_PACKET_ID, EMPTY_PACKET_ID], + [100, 101, EMPTY_PACKET_ID, EMPTY_PACKET_ID], + ], dtype=jnp.int32) + buffer_birth_steps = jnp.array([ + [-1, -1, -1, -1], + [1, 2, -1, -1], + ], dtype=jnp.int32) + packet_seqs = jnp.array([0, 2], dtype=jnp.int32) + new_frames = jnp.array([3, 5], dtype=jnp.int32) + + result, result_births, new_packet_seqs = add_new_frames( + buffer_states, buffer_birth_steps, new_frames, packet_seqs, current_step=10 + ) + self.assertEqual(new_packet_seqs[0], 3) + self.assertEqual(new_packet_seqs[1], 7) + self.assertEqual(jnp.sum(result[0] != EMPTY_PACKET_ID), 3) + # Overflow should cap to queue capacity and drop oldest first. + self.assertEqual(jnp.sum(result[1] != EMPTY_PACKET_ID), self.QUEUE) + self.assertEqual(int(result[1][0]), 13) + self.assertEqual(int(result_births[1][0]), 10) def test_simulate(self): - # Test full simulation - buffer_states = jnp.array([1, 1, 1, 0]) - new_frames = jnp.array([0, 1, 1, 1]) + buffer_states = jnp.array([ + [7, 8, EMPTY_PACKET_ID, EMPTY_PACKET_ID], + [EMPTY_PACKET_ID, EMPTY_PACKET_ID, EMPTY_PACKET_ID, EMPTY_PACKET_ID], + [30, EMPTY_PACKET_ID, EMPTY_PACKET_ID, EMPTY_PACKET_ID], + [40, 41, 42, EMPTY_PACKET_ID], + ], dtype=jnp.int32) + buffer_birth_steps = jnp.array([ + [1, 2, -1, -1], + [-1, -1, -1, -1], + [5, -1, -1, -1], + [2, 3, 4, -1], + ], dtype=jnp.int32) + packet_seqs = jnp.array([2, 0, 1, 3], dtype=jnp.int32) + new_frames = jnp.array([0, 2, 1, 0], dtype=jnp.int32) actions = jnp.array([Actions.TX.value, Actions.CS.value, Actions.CS.value, Actions.CS.value]) - expected_buffer_states = jnp.array([0, 1, 1, 1]) - expected_channel_state = 1 - - result_buffer_states, result_channel_state = simulate(buffer_states, new_frames, actions) - self.assertTrue(jnp.array_equal(result_buffer_states, expected_buffer_states)) - self.assertEqual(result_channel_state, expected_channel_state) - - # Test simulation with no actions - buffer_states = jnp.array([1, 0, 1, 0]) - new_frames = jnp.array([0, 1, 0, 1]) - actions = jnp.array([Actions.CS.value, Actions.CS.value, Actions.CS.value, Actions.CS.value]) - - expected_buffer_states = jnp.array([1, 1, 1, 1]) - expected_channel_state = 0 - - result_buffer_states, result_channel_state = simulate(buffer_states, new_frames, actions) - self.assertTrue(jnp.array_equal(result_buffer_states, expected_buffer_states)) - self.assertEqual(result_channel_state, expected_channel_state) - - # Test simulation with multiple actions - buffer_states = jnp.array([0, 0, 0, 0]) - new_frames = jnp.array([0, 1, 0, 1]) - actions = jnp.array([Actions.TX.value, Actions.CS.value, Actions.CS.value, Actions.TX.value]) - - expected_buffer_states = jnp.array([0, 1, 0, 1]) - expected_channel_state = -1 - - result_buffer_states, result_channel_state = simulate(buffer_states, new_frames, actions) - self.assertTrue(jnp.array_equal(result_buffer_states, expected_buffer_states)) - self.assertEqual(result_channel_state, expected_channel_state) + tx_packet_mask = jnp.array([True, False, False, False]) + tx_success_mask = jnp.array([True, False, False, False]) + ( + result_buffer_states, + result_buffer_birth_steps, + result_channel_state, + result_packet_seqs, + planned_packets, + successful_packets, + ) = simulate( + buffer_states, buffer_birth_steps, new_frames, actions, packet_seqs, current_step=10, + tx_packet_mask=tx_packet_mask, tx_success_mask=tx_success_mask, + ) + + self.assertEqual(result_channel_state, 1) + # Successful TX on station 0 removes queue head (7). + self.assertEqual(int(result_buffer_states[0][0]), 8) + self.assertTrue(jnp.array_equal(result_buffer_states[0][1:], jnp.array([EMPTY_PACKET_ID] * 3))) + self.assertTrue(jnp.array_equal(result_buffer_birth_steps[0][1:], jnp.array([-1, -1, -1], dtype=jnp.int32))) + # New arrivals were appended. + self.assertEqual(jnp.sum(result_buffer_states[1] != EMPTY_PACKET_ID), 2) + self.assertEqual(jnp.sum(result_buffer_states[2] != EMPTY_PACKET_ID), 2) + self.assertTrue(jnp.array_equal(result_packet_seqs, packet_seqs + new_frames)) + self.assertTrue(jnp.array_equal(planned_packets, jnp.array([1, 0, 0, 0], dtype=jnp.int32))) + self.assertTrue(jnp.array_equal(successful_packets, jnp.array([1, 0, 0, 0], dtype=jnp.int32))) if __name__ == "__main__": diff --git a/test/sim/test_traffic.py b/test/sim/test_traffic.py index e0e68b5..13f71f9 100644 --- a/test/sim/test_traffic.py +++ b/test/sim/test_traffic.py @@ -61,7 +61,7 @@ def test_stats(self): estimated_acf=np.corrcoef(yjax[:-1]-yjax.mean(), yjax[1:]-yjax.mean())[1,0] - self.assertAlmostEqual(yjax.mean(), stats.mean, places=3) + self.assertAlmostEqual(yjax.mean(), stats.mean, places=2) self.assertAlmostEqual(yjax.var(), stats.variance, places=1) self.assertAlmostEqual(estimated_acf, stats.acf_lag1, places=2) ...