Bottleneck Replay Buffer pada Training World Model

Pelatihan arsitektur World Model dan Joint-Embedding Predictive Architecture (JEPA)—seperti implementasi LeMario—memerlukan sampling kontinu dari ribuan trajectory game interaktif. Setiap episode menghasilkan ratusan frame yang berisi observasi multidimensi, aksi agent, reward, dan flag terminal.

Ketika dataset replay buffer mencapai puluhan juta baris, dua masalah performa utama muncul pada database relasional:

  1. Overhead Index Storage B-Tree: B-Tree menyimpan referensi untuk setiap tuple. Pada tabel time-series atau step trajectory berukuran puluhan gigabyte, ukuran index B-Tree sering kali melebihi kapasitas RAM buffer cache, menurunkan performa I/O secara drastis saat proses penulisan bersamaan.
  2. Bottleneck OFFSET pada Pagination: Training worker yang mengambil batch berurutan menggunakan query berbasis offset (LIMIT n OFFSET m) memaksa engine database memindai dan membuang m baris pertama secara berulang. Kompleksitas waktunya adalah O(N).

Desain Skema Database Replay Trajectory

Tabel replay menyimpan observasi dan state transition yang disisipkan secara sekuensial berdasarkan siklus simulasi lingkungan game. Berikut skema PostgreSQL yang dirancang untuk data append-only:

CREATE TABLE trajectory_frames (
    id BIGSERIAL PRIMARY KEY,
    episode_id UUID NOT NULL,
    step_idx INT NOT NULL,
    state_embedding REAL[] NOT NULL,
    action SMALLINT NOT NULL,
    reward REAL NOT NULL,
    is_terminal BOOLEAN NOT NULL DEFAULT FALSE,
    created_at TIMESTAMPTZ NOT NULL DEFAULT CLOCK_TIMESTAMP()
);

BRIN Index: Mengapa Tepat untuk Step Trajectory

Block Range Index (BRIN) bekerja dengan membaca nilai minimum dan maksimum dari rentang halaman fisik disk (rentang default PostgreSQL adalah 128 blok buffer, atau 1 MB). BRIN mengasumsikan adanya korelasi kuat antara urutan fisik penyimpanan di disk dan urutan logis nilai kolom.

Karakteristik data environment step memenuhi asumsi ini karena frame trajectory ditulis secara append-only. Nilai id dan created_at selalu meningkat seiring bertambahnya blok data di storage.

Perbandingan B-Tree vs BRIN

  • Ukuran Penyimpanan: B-Tree menyimpan pointer untuk setiap baris, menghasilkan rasio ukuran index terhadap tabel sekitar 20-30%. BRIN hanya menyimpan rangkuman per blok halaman, memangkas ukuran index hingga lebih dari 95% dibandingkan B-Tree.
  • Write Overhead: B-Tree memerlukan rebalancing halaman secara berkala saat penyisipan data baru. BRIN hanya memperbarui nilai min/max pada range blok aktif, meminimalkan write amplification.
  • Kelemahan: BRIN tidak efisien jika data sering mengalami UPDATE atau DELETE acak, karena fragmentasi fisik akan memperlebar rentang min/max pada setiap blok, sehingga engine harus membaca lebih banyak blok data (lossy index scan).

Implementasi Index BRIN

-- Hapus index B-Tree default jika query time-series dominan
-- pages_per_range dapat disesuaikan (default 128 blok)
CREATE INDEX idx_trajectory_frames_brin_step 
ON trajectory_frames USING BRIN (id, created_at) 
WITH (pages_per_range = 64);

Pengurangan pages_per_range dari 128 ke 64 meningkatkan presisi filter range scan dengan sedikit trade-off pada ukuran index.

Mengeliminasi Latensi Tinggi dengan Keyset Pagination

Hindari pagination klasik untuk worker ingest data:

-- ANTI-PATTERN: Lambat pada dataset besar
SELECT id, episode_id, state_embedding, action 
FROM trajectory_frames 
ORDER BY id ASC 
LIMIT 512 OFFSET 1000000;

Ganti dengan Keyset Pagination (seek method). Query ini memanfaatkan predikat filter langsung pada titik data terakhir yang telah diproses oleh training worker:

-- POLA OPTIMAL: Eksekusi konstan O(1)
SELECT id, episode_id, state_embedding, action, reward, is_terminal
FROM trajectory_frames
WHERE id > 1000000
ORDER BY id ASC
LIMIT 512;

Implementasi Pipeline Fetching pada Worker Python

Potongan kode berikut menunjukkan worker sampling PyTorch/JAX yang mengekstrak replay batch secara stabil tanpa degradasi performa di step tinggi:

from typing import Iterator, List, Tuple
import psycopg2
import psycopg2.extras

def fetch_trajectory_batches(
    connection_str: str,
    batch_size: int = 512
) -> Iterator[List[Tuple]]:
    conn = psycopg2.connect(connection_str)
    last_seen_id = 0
    
    query = """
        SELECT id, episode_id, state_embedding, action, reward, is_terminal
        FROM trajectory_frames
        WHERE id > %s
        ORDER BY id ASC
        LIMIT %s;
    """
    
    try:
        with conn.cursor() as cursor:
            while True:
                cursor.execute(query, (last_seen_id, batch_size))
                rows = cursor.fetchall()
                
                if not rows:
                    break
                
                last_seen_id = rows[-1][0]
                yield rows
    finally:
        conn.close()

Pemeliharaan dan Monitoring BRIN

Jika terjadi penyisipan data di luar urutan (misalnya worker paralel memasukkan trajectory lama secara terlambat), jalankan fungsi pemeliharaan PostgreSQL untuk memastikan batas ringkasan range tetap akurat:

-- Merapikan range BRIN yang kotor atau belum terindeks
SELECT brin_desummarize_range('idx_trajectory_frames_brin_step', 0);
SELECT brin_summarize_new_values('idx_trajectory_frames_brin_step');

Kombinasi BRIN index untuk efisiensi disk/memori cache dan Keyset Pagination memastikan latensi database tetap stabil pada interval sub-milidetik, menjaga throughput tensor training worker tetap jenuh (saturating GPU memory bandwidth).