How does modulation frequency appear in the modulation spectra?
This problem is precisely what synchrosqueezing wavelet transform was invented for, and indeed maps it with great precision. I'm still developing it, and first pre-release is expected today or tomorrow, but currently you can use it with all default settings to estimate the amplitude modulation information by plotting the appropriate row:

# printed output
1998  # carrier frequency
200   # modulator frequency

I'll be writing a detailed explanation of synchrosqueezing in context of CWT in a separate Q&amp;A 'soon', but for now you can refer to sources in relevant functions' docstrings. The change to 200 (or less, relative to carrier_freq) was necessitated from limitations of working only with default values, but you can still use defaults with additional re-interpreting of max_row_fft.
As a tip, forget STFT for this task, CWT and CWT-based synchrosqueezing are superior. For more relevant intuition, see this answer (downvotes were due to a related 'fiasco').

How's it work?
I'll use STFT to explain; firstly you'll need more segments in total to capture the modulator; default segment length is 256, and overlap half of it, so to make more windows let's use noverlap=224 and don't touch nfft (no real need) - plot abs(Zxx), and the row at which it is maximal:

The modulator is seen more clearly now, and the plot on the right captures it (the AM &quot;envelope&quot;). If it's accurate, then the FFT of that row should produce a peak at the modulator's frequency - and it does:

To understand how this works one must understand first and foremost how the DFT works, and then the STFT; I'll leave good learning sources below, but also in a nutshell: STFT takes a windowed DFT over the signal, each window centered at a different point of the signal; if the window decays away from center, then signal points farther from center contribute less to spectrum's computation, thus localizing frequencies.
Each row of STFT uses the same window on a different frequency complex sinusoid (i.e. DFT basis), overall STFT forming a time-frequency representation of the signal. The row in STFT has an associated frequency, returned by scipy.signal.stft, which in this example is the carrier's frequency. As STFT values at that row capture the signal contents due to that row's associated frequency components, the row essentially is the original signal (in this case), thus taking its absolute value gives us the (carrier's amplitude) modulator signal, whose frequency we can find via DFT (FFT). One caveat, it's not just one row, as nearby frequencies also correlate to produce nonzero values; this is what synchrosqueezing addresses.

FT
DFT
STFT
DFT further intuition
DFT vs DTFT, zero-padding


&quot;Row&quot; clarification -- the &quot;row&quot; is the original signal by approximation (in this case - only one 'effective' row), except also having a 90-shifted imaginary component, so taking absolute value yields the enveloping (AM) signal. This is clear from a &quot;continuous STFT&quot; (shift window by 1 sample); it won't be exact (but note that the frequency is exact) as again, the rest of information is in other rows, and reconstruction involves extra steps (done by inverse STFT). -- Code


Code: (-- STFT version)
import numpy as np
import matplotlib.pyplot as plt
from ssqueezepy import synsq_cwt_fwd

#%%## AM signal generation ###################################################
duration = 1 # in seconds
fs = 44100 # Hz
carrier_freq = 2000 # Hz
mod_freq = 200
t = np.linspace(0,duration,endpoint=True,num=fs)
audio = (1 + 1 * np.cos(2*np.pi*mod_freq*t)) * np.cos(2*np.pi*carrier_freq*t)

#%%###########################################################################
Tx, fs, *_ = synsq_cwt_fwd(audio, t=t)
aTx = np.abs(Tx)

#%%## Plot synchrosqueezed CWT ###############################################
plt.imshow(np.flipud(aTx), aspect='auto', cmap='bone')
yt = [&quot;%.2f&quot; % f for f in fs[np.linspace(0, len(fs) - 1, 6).astype('int32')]][::-1]
plt.gca().set_yticklabels(['', *yt])
plt.show()

#%%## Find carrier freq ######################################################
max_row_idx = np.where(aTx == aTx.max())[0]
max_row = aTx[max_row_idx].squeeze()

# print peak's frequency
print(fs[max_row_idx])
#%%###########################################################################
# plot amplitude modulator
plt.plot(max_row[:800]); plt.show()

# find modulator frequency
max_row_fft = np.abs(np.fft.rfft(max_row))
plt.plot(max_row_fft); plt.show()

# peak at 200; exclude dc term
peak_bin = np.argmax(max_row_fft[1:]) + 1
print(peak_bin)

Note: amplitude values in above plot are off due to mistaken normalization used in MATLAB repository; this'll be fixed.