pseudocode.txt

encoder = Encoder()
decoder = Decoder()
codebook = Codebook(k, d)
hmm = HMM(num_states=k)
optimizer_vqvae = Adam([...])
optimizer_hmm = Adam([...])

def train_vqvae_hmm(epochs, data_loader):
    for epoch in range(epochs):
        for x_seq in data_loader:
            z_e_seq = encoder(x_seq)
            z_q_seq, code_indices = quantize(z_e_seq, codebook)
            z_q_seq_st = z_e_seq + (z_q_seq - z_e_seq).detach()
            x_hat_seq = decoder(z_q_seq_st)

            recon_loss = MSE(x_hat_seq, x_seq)
            commit_loss = beta * MSE(z_e_seq, z_q_seq.detach())
            codebook_loss = MSE(z_q_seq, z_e_seq.detach())
            loss_vqvae = recon_loss + commit_loss + codebook_loss

            optimizer_vqvae.zero_grad()
            loss_vqvae.backward()
            optimizer_vqvae.step()

        # Train HMM on the sequence of discrete code indices
        all_code_indices = collect_indices_from_dataset(encoder, codebook, data_loader)
        hmm.train_em(all_code_indices)

def sample_from_vqvae_hmm(seq_len):
    z_idx_seq = hmm.sample(seq_len)
    z_q_seq = codebook.lookup(z_idx_seq)
    return decoder(z_q_seq)