PyJama: Differentiable Jamming Simulator

• PyJama is a fully differentiable, open-source library that models jamming and anti-jamming strategies in realistic MIMO-OFDM wireless environments.
• It integrates NVIDIA Sionna and TensorFlow auto-diff to enable end-to-end gradient-based training of both jammers and receivers.
• Empirical results show that learned jamming power allocation can significantly increase error rates while challenging traditional anti-jamming techniques.

PyJama is a fully differentiable open-source library for the simulation and optimization of jamming and anti-jamming strategies in realistic MIMO-OFDM wireless communication systems. Integrated with NVIDIA Sionna’s TensorFlow-based physical-layer simulator, PyJama enables end-to-end gradient-based training of both jammers and receivers across complex link-level scenarios, including advanced waveform, channel, and mobility models. This infrastructure provides, for the first time, the ability to learn power allocation patterns and receiver countermeasures via stochastic gradient descent under rigorous physical-layer constraints (Ulbricht et al., 2024).

1. Motivation and Goals

Wireless communication systems are increasingly threatened by sophisticated jamming attacks targeting physical-layer protocols. Traditional research in this area is predominantly analytical and lacks the differentiable frameworks necessary to support end-to-end machine learning–based optimization. PyJama directly addresses this gap by providing a differentiable environment where both jamming and anti-jamming can be simulated, learned, and evaluated under realistic, high-fidelity system models derived from Sionna. This design enables:

• Stochastic gradient descent–based learning of optimal jamming power allocation and, in future extensions, jamming waveforms.
• Co-optimization of receiver architecture, including learned beamforming and equalization, by supporting backpropagation through the entire channel, jamming, and receiver chain.
• Realistic simulation leveraging state-of-the-art models for MIMO channels, OFDM waveforms, forward error correction (LDPC), and mobility.

(Ulbricht et al., 2024)

2. Architecture and System Components

PyJama is built modularly atop Sionna’s simulation graph, introducing new blocks specifically designed for jamming and anti-jamming:

• Jamming Modules: Frequency-domain jammers (modes: barrage, pilot-only, data-only, sparse) and time-domain jammers (currently barrage); symbol alphabets include uniform-disk, Gaussian, QAM, or custom distributions. The key parameter is the trainable power allocation vector $\boldsymbol{\rho}$, representing per-resource-element transmit power.
• Anti-Jamming Modules: POS (Projection Onto Signal subspace) effectively nulls the estimated jamming subspace, while IAN-LMMSE (Interference-As-Noise LMMSE) treats the jammer as colored noise. Channel state information (CSI) can be perfect or estimated using pilots, with infrastructure to support silence-based estimation for jamming environments.
• Training Infrastructure: Full support for TensorFlow auto-diff, optimizers (Adam), loss computations (L1, MSE), and power-constraint regularization.

Data flow in a typical simulation is as follows:

Step	Sionna Component	PyJama Extension
Transmitter	OFDM, MIMO, FEC	–
Channel	Realistic channel	–
Jamming	–	Frequency/time-domain jammer
Receiver	LS/LMMSE	POS, IAN-LMMSE anti-jamming
Decoder	LDPC, demapper	Iteration-loss for LDPC
Loss, Update	–	SGD-based optimizer

(Ulbricht et al., 2024)

3. Mathematical Formulation and Optimization

MIMO-Jamming Channel Model

The MIMO jamming channel in PyJama is modeled as: $$\mathbf{y}[k] = \sum_{\ell=0}^{L-1}\bigl(\mathbf{H}[k,\ell]\mathbf{s}[k-\ell] + \mathbf{J}[k,\ell]\mathbf{w}[k-\ell]\bigr) + \mathbf{n}[k],$$ where $\mathbf{s}[k]$ and $\mathbf{w}[k]$ are user and jammer signals, $\mathbf{H}[k,\ell]$ and $\mathbf{J}[k,\ell]$ are the respective channel impulse responses, and $\mathbf{n}[k]$ is AWGN.

Jamming Power Allocation and Optimization

The jammer’s per-OFDM-symbol power allocation is parametrized as $\boldsymbol{\rho} = [\rho_1,\dots,\rho_{N_s}]$, constrained by an average power limit: $$\frac{1}{N_s} \sum_{i=1}^{N_s} \rho_i \leq \rho_{\max}, \quad \rho_i \geq 0.$$ The loss function measuring receiver performance (bit or block error) is denoted $\ell(\mathbf{b},\hat{\mathbf{b}})$, where $\mathbf{b}$ are ground-truth bits, $\hat{\mathbf{b}}$ are soft receiver outputs. The end-to-end jamming optimization solves: $$\max_{\boldsymbol{\rho} \succeq 0}\; \mathbb{E}_{\rm data}[\ell(\mathbf{b},\hat{\mathbf{b}}(\boldsymbol{\rho}))]$$ subject to the power constraint, implemented via a hinge penalty: $$\mathcal{L}(\boldsymbol{\rho}) = -\mathbb{E}[\ell(\mathbf{b},\hat{\mathbf{b}})] + \lambda\max\Bigl(0, \frac{1}{N_s} \sum_i \rho_i - \rho_{\max}\Bigr).$$ PyJama does not use binary cross-entropy for maximization of error rates due to unbounded loss scaling, preferring L1 or MSE losses and their iteration-wise aggregation for LDPC decoders.

(Ulbricht et al., 2024)

4. Implementation and Simulation Workflow

PyJama provides an end-to-end reproducible workflow, from installation and system configuration to gradient-based training. Core steps include:

• Resource grid and pilot pattern configuration, including silent slots for estimation.
• Channel modeling via Sionna’s TDL/3GPP/geometry-based models.
• Jamming block instantiation with user-selectable modes and power allocation vectors.
• Receiver selection (conventional or anti-jamming).
• Layered decoding and loss computation.
• End-to-end optimization using TensorFlow’s automatic differentiation and Adam optimizer.

The following code sketch outlines the training loop for learning jamming power allocation:

optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3) for step in range(5000): b = tf.random.uniform((batch_size, info_bits), maxval=2, dtype=tf.int32) with tf.GradientTape() as tape: s = rg.mapper(b) X = rg.ofdm_mod(s) Y_uncjammed = channel(X) J = jammer(Y_uncjammed) Y = Y_uncjammed + J Y_fd = rg.ofdm_demod(Y) Y_eq = pos(Y_fd) if use_POS else data_eq(Y_fd) llr_list = ldpc_dec(Y_eq, rg.pilot_pattern) losses = [tf.reduce_mean(tf.abs(b - llr)) for llr in llr_list] data_loss = tf.add_n(losses) avg_power = tf.reduce_mean(jammer.power_allocation) pen = tf.nn.relu(avg_power - rho_max) * lambda_penalty total_loss = -data_loss + pen grads = tape.gradient(total_loss, [jammer.power_allocation]) optimizer.apply_gradients(zip(grads, [jammer.power_allocation]))

This enables simultaneous tuning of the jamming strategy under power constraints and realistic signal processing assumptions.

(Ulbricht et al., 2024)

5. Experimental Results and Key Findings

PyJama demonstrates several compelling scenarios:

• MIMO Jamming and Anti-jamming: In a 4-UE uplink to a 16-antenna BS with a +2 dB barrage jammer, conventional receivers without countermeasures exhibit error floors (BLER=1), while POS recovers performance at ~1–2 dB cost (due to CSI estimation and DoF loss).
• Mobility: Jamming effectiveness is more sensitive to jammer mobility than UE mobility; robust pilot design can mitigate UE-induced effects.
• Time-domain OFDM Jamming: CP-violating jammers (time-domain barrage) generate full-rank interference, eluding nulling by POS; frequency-domain models substantially underestimate the jammer’s impact.
• Learned Jamming Strategies: Learning the power allocation vector $\boldsymbol{\rho}$ produces non-trivial, interpretable strategies. Single-UE jammers focus on pilots, few-UE jammers exploit pilot/data trade-offs, and many-UE/jammers allocate to data slots. Learned strategies achieve up to a $7.2\times$ BER increase at 5 dB SNR (and $20$ dB effectiveness gain compared to uniform barrage).
• Bypassing Anti-jamming (POS): Learned jammers exploit POS’s channel estimation slot assignment by remaining silent during estimation, defeating projection-based approaches and matching the efficacy of much higher-power barrage attacks.

All results are fully reproducible via code and notebooks provided at pyjama.ethz.ch.

6. Limitations and Future Research Directions

PyJama currently supports only barrage-type jamming in the time domain and implements two anti-jamming receivers (POS, IAN-LMMSE). There is no joint optimization of jamming waveform structure or beamforming vectors, nor are reinforcement learning agents or adversarial minimax loops built-in. Simulation in time or frequency domains may miss certain jamming effects unless chosen with care. Future directions outlined include:

• Expanding jamming modules to support time-sparse, hybrid, or beamformed attacks.
• Integration of neural anti-jamming front-ends through end-to-end differentiable architectures.
• Adversarial training and reinforcement learning for real-time, adaptive jamming strategies.
• Supporting alternative physical-layer waveforms and higher-frequency (mmWave) mobility/propagation models leveraging Sionna-RT.
• Deeper co-design of pilot structure and jamming to address estimation vulnerabilities.

These extensions are enabled by the modular and differentiable design architecture of PyJama as an augmentation of Sionna’s simulation stack (Ulbricht et al., 2024).