Train a Simple GAN on 1D Gaussian Data

import numpy as np

def train_gan_1d(real_mean: float = 5.0, real_std: float = 1.0,
                 n_epochs: int = 5000, batch_size: int = 64,
                 lr: float = 0.01):
    g_w = np.random.randn() * 0.1
    g_b = np.random.randn() * 0.1
    d_w = np.random.randn() * 0.1
    d_b = np.random.randn() * 0.1

    def sigmoid(x):
        return 1.0 / (1.0 + np.exp(-np.clip(x, -500, 500)))

    def generator(z):
        return g_w * z + g_b

    def discriminator(x):
        return sigmoid(d_w * x + d_b)

    for epoch in range(n_epochs):
        real_data = np.random.normal(real_mean, real_std, batch_size)
        z = np.random.randn(batch_size)
        fake_data = generator(z)

        d_real = discriminator(real_data)
        d_fake = discriminator(fake_data)

        d_loss_real = -np.mean(np.log(d_real + 1e-8))
        d_loss_fake = -np.mean(np.log(1 - d_fake + 1e-8))

        dd_real_dw = np.mean(real_data * d_real * (1 - d_real))
        dd_real_db = np.mean(d_real * (1 - d_real))
        dd_fake_dw = -np.mean(fake_data * d_fake * (1 - d_fake))
        dd_fake_db = -np.mean(d_fake * (1 - d_fake))

        d_w += lr * (dd_real_dw + dd_fake_dw)
        d_b += lr * (dd_real_db + dd_fake_db)

        z = np.random.randn(batch_size)
        fake_data = generator(z)
        d_fake = discriminator(fake_data)

        dg_dw = np.mean(d_w * z * d_fake * (1 - d_fake)) / (np.mean(d_fake) + 1e-8)
        dg_db = np.mean(d_w * d_fake * (1 - d_fake)) / (np.mean(d_fake) + 1e-8)
        g_w += lr * dg_dw
        g_b += lr * dg_db

    return g_w, g_b, d_w, d_b