# realwenlongwang/drama architecture

> Architecture diagram of the GitHub repository realwenlongwang/drama, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-06-14). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/realwenlongwang/drama
- Repository: https://github.com/realwenlongwang/drama

## Overview

1. **Purpose and top-level shape**
   - This repo implements **Drama**, a **model-based reinforcement learning** system built around **Mamba/Mamba2 sequence models**.
   - The architecture is layered: experiment entrypoints on top, RL agents and environment wrappers in the middle, model components below, and custom CUDA/Triton sequence operators at the bottom.

2. **Training and evaluation entrypoints**
   - `train.py` is the main orchestration script for running experiments.
   - `eval.py` handles evaluation-only runs.
   - `config_files/configure.yaml` and the training config files drive runtime behavior, letting the same codebase switch among model backbones and behavior-learning variants.

3. **RL control plane**
   - `agents.py` likely contains the high-level agent logic that ties together world modeling, action selection, and optimization.
   - `replay_buffer.py` is the storage boundary for collected experience, feeding sequences or transitions back into training.
   - `tools.py` and `utils.py` provide shared experiment utilities such as configuration handling, logging, scheduling, or metric helpers.

4. **Environment subsystem**
   - `envs/` contains task-specific environment adapters:
     - `my_atari.py`
     - `my_dmc.py`
     - `my_memory_maze.py`
   - `envs/parallel.py` suggests vectorized or subprocess-based rollout execution.
   - `env_wrapper.py` likely normalizes observation/action APIs across these environments so the agent can train against multiple benchmark families with the same interface.

5. **Model-based RL core**
   - `sub_models/world_models.py` is the likely center of the learned latent dynamics / world model.
   - `sub_models/transformer_model.py` indicates an alternative sequence backbone.
   - `sub_models/attention_blocks.py`, `sub_models/functions_losses.py`, and `sub_models/laprop.py` support model internals, training losses, and optimization behavior.
   - This split suggests a design where the world model can be paired with different sequence architectures and loss heads.

6. **Behavior learning / policy side**
   - The README states the repo supports **Actor-Critic** and **PPO** behavior models.
   - Those policy/value components likely live in `sub_models/` and are coordinated by the agent/training layer rather than being separate standalone services.
   - This means the system has two main learning loops: world-model learning and behavior optimization.

7. **Mamba/Mamba2 implementation layer**
   - The `mamba_ssm/` package is a vendored or embedded sequence-model library.
   - Key areas include:
     - `modules/` for Mamba blocks, MHA, MLP, and composite sequence blocks
     - `models/` for higher-level model assembly and configs
     - `ops/` for selective scan and Triton-based kernels
     - `distributed/` for tensor-parallel or multi-GPU support
   - This package is the main architectural dependency for the “Mamba-enabled” part of Drama.

8. **Custom sequence operators and kernel boundary**
   - `csrc/selective_scan/` contains compiled CUDA/C++ extensions for the selective scan operator used by Mamba-style recurrence.
   - The presence of forward/backward kernels for fp16, bf16, and fp32 indicates performance-critical training/inference paths are implemented below Python.
   - `mamba_ssm/ops/selective_scan_interface.py` is the Python boundary that calls into these native kernels.

9. **Distributed / performance tooling**
   - `mamba_ssm/distributed/` suggests optional tensor parallelism or distributed execution for larger models.
   - The repo also includes Triton kernels under `mamba_ssm/ops/triton/`, which likely cover layer norm, state update, chunk scan, and fused sequence operations.
   - Together, these pieces show the system is designed for GPU-heavy training with custom acceleration rather than pure PyTorch reference code.

10. **Experiment packaging and runtime**
    - `Dockerfile` supports containerized execution, mainly for reproducibility and GPU setup.
    - `requirements.txt` plus the README’s installation notes show the repo depends on a specific PyTorch and `mamba-ssm` stack.
    - `config_files/` and the README’s CLI examples imply experiments are configured primarily through YAML plus command-line overrides.

11. **Data flow through the system**
    - Environments produce trajectories or observations.
    - The agent collects them into the replay buffer.
    - Training samples sequences from the buffer to update the world model and behavior policy.
    - The world model uses Mamba/Transformer backbones to predict latent dynamics and/or future observations.
    - The policy/value side uses those learned representations to improve action selection.

12. **What to show in an architecture diagram**
    - Draw four main layers: **entrypoints/config**, **RL orchestration**, **environment + replay**, and **model backends**.
    - Under model backends, separate **Drama-specific sub_models** from the reusable **mamba_ssm** library and then from **native CUDA/Triton ops**.
    - Highlight the two external interfaces: **Gym-like environments** on one side and **GPU compiled kernels** on the other.

Main entry points: [`agents.py`](https://github.com/realwenlongwang/drama/blob/master/agents.py) (agents.py), [`envs/parallel.py`](https://github.com/realwenlongwang/drama/blob/master/envs/parallel.py) (parallel rollout), [`sub_models/world_models.py`](https://github.com/realwenlongwang/drama/blob/master/sub_models/world_models.py) (world models.py), [`mamba_ssm/models/mixer_seq_simple.py`](https://github.com/realwenlongwang/drama/blob/master/mamba_ssm/models/mixer_seq_simple.py) (mamba models)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_entry["Entrypoints"]
  node_train["train.py<br/>train entry<br/>[train.py]"]
  node_eval["eval.py<br/>eval entry<br/>[eval.py]"]
  node_config["configure.yaml<br/>runtime config<br/>[configure.yaml]"]
end

subgraph group_control["RL Control"]
  node_agents["agents.py<br/>agent control<br/>[agents.py]"]
  node_replay[("replay buffer<br/>experience store<br/>[replay_buffer.py]")]
  node_tools["tools.py<br/>experiment utils<br/>[tools.py]"]
  node_utils["utils.py<br/>shared utils<br/>[utils.py]"]
end

subgraph group_env["Environments"]
  node_wrapper["env wrapper<br/>env api<br/>[env_wrapper.py]"]
  node_parallel["parallel rollout<br/>vector env<br/>[parallel.py]"]
  node_atari["Atari env<br/>task adapter<br/>[my_atari.py]"]
  node_dmc["DMC env<br/>task adapter<br/>[my_dmc.py]"]
  node_maze["Memory maze<br/>task adapter<br/>[my_memory_maze.py]"]
end

subgraph group_models["Drama Models"]
  node_world["world models.py<br/>dynamics model<br/>[world_models.py]"]
  node_transformer["transformer model<br/>seq backbone"]
  node_blocks["attention blocks<br/>model internals"]
  node_losses["losses"]
  node_laprop["laprop<br/>optimizer helper<br/>[laprop.py]"]
end

subgraph group_mamba["Mamba Stack"]
  node_mamba_models["mamba models<br/>model assembly"]
  node_mamba_blocks["mamba blocks<br/>sequence modules<br/>[block.py]"]
  node_mamba_ops["selective scan<br/>python ops"]
end

subgraph group_native["Native Kernels"]
  node_triton["triton ops<br/>gpu kernels<br/>[ssd_combined.py]"]
  node_selective_scan["selective scan cuda<br/>cuda extension<br/>[selective_scan.cpp]"]
end

node_config -->|"configures"| node_train
node_config -->|"configures"| node_eval
node_train -->|"starts"| node_agents
node_eval -->|"runs"| node_agents
node_agents -->|"reads/writes"| node_replay
node_agents -->|"interacts"| node_wrapper
node_wrapper -->|"uses"| node_parallel
node_parallel -->|"wraps"| node_atari
node_parallel -->|"wraps"| node_dmc
node_parallel -->|"wraps"| node_maze
node_agents -->|"trains"| node_world
node_agents -->|"trains"| node_transformer
node_world -->|"uses"| node_blocks
node_world -->|"optimizes with"| node_losses
node_world -->|"steps with"| node_laprop
node_transformer -->|"uses"| node_blocks
node_blocks -->|"implements"| node_mamba_models
node_mamba_models -->|"calls"| node_mamba_ops
node_mamba_ops -->|"dispatches"| node_selective_scan
node_mamba_blocks -->|"depends on"| node_mamba_ops
node_mamba_models -->|"uses"| node_triton
node_agents -->|"uses"| node_tools
node_agents -->|"uses"| node_utils

click node_train "https://github.com/realwenlongwang/drama/blob/master/train.py"
click node_eval "https://github.com/realwenlongwang/drama/blob/master/eval.py"
click node_config "https://github.com/realwenlongwang/drama/blob/master/config_files/configure.yaml"
click node_agents "https://github.com/realwenlongwang/drama/blob/master/agents.py"
click node_replay "https://github.com/realwenlongwang/drama/blob/master/replay_buffer.py"
click node_tools "https://github.com/realwenlongwang/drama/blob/master/tools.py"
click node_utils "https://github.com/realwenlongwang/drama/blob/master/utils.py"
click node_wrapper "https://github.com/realwenlongwang/drama/blob/master/env_wrapper.py"
click node_parallel "https://github.com/realwenlongwang/drama/blob/master/envs/parallel.py"
click node_atari "https://github.com/realwenlongwang/drama/blob/master/envs/my_atari.py"
click node_dmc "https://github.com/realwenlongwang/drama/blob/master/envs/my_dmc.py"
click node_maze "https://github.com/realwenlongwang/drama/blob/master/envs/my_memory_maze.py"
click node_world "https://github.com/realwenlongwang/drama/blob/master/sub_models/world_models.py"
click node_transformer "https://github.com/realwenlongwang/drama/blob/master/sub_models/transformer_model.py"
click node_blocks "https://github.com/realwenlongwang/drama/blob/master/sub_models/attention_blocks.py"
click node_losses "https://github.com/realwenlongwang/drama/blob/master/sub_models/functions_losses.py"
click node_laprop "https://github.com/realwenlongwang/drama/blob/master/sub_models/laprop.py"
click node_mamba_models "https://github.com/realwenlongwang/drama/blob/master/mamba_ssm/models/mixer_seq_simple.py"
click node_mamba_blocks "https://github.com/realwenlongwang/drama/blob/master/mamba_ssm/modules/block.py"
click node_mamba_ops "https://github.com/realwenlongwang/drama/blob/master/mamba_ssm/ops/selective_scan_interface.py"
click node_triton "https://github.com/realwenlongwang/drama/blob/master/mamba_ssm/ops/triton/ssd_combined.py"
click node_selective_scan "https://github.com/realwenlongwang/drama/blob/master/csrc/selective_scan/selective_scan.cpp"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_train,node_eval,node_config toneBlue
class node_agents,node_replay,node_tools,node_utils toneAmber
class node_wrapper,node_parallel,node_atari,node_dmc,node_maze toneMint
class node_world,node_transformer,node_blocks,node_losses,node_laprop toneRose
class node_mamba_models,node_mamba_blocks,node_mamba_ops toneIndigo
class node_triton,node_selective_scan toneTeal
```

## Components

### Entrypoints

Experiment launch and configuration

- [train.py](https://github.com/realwenlongwang/drama/blob/master/train.py): train entry · Main experiment orchestration script · `train.py`
- [eval.py](https://github.com/realwenlongwang/drama/blob/master/eval.py): eval entry · Evaluation-only execution path · `eval.py`
- [configure.yaml](https://github.com/realwenlongwang/drama/blob/master/config_files/configure.yaml): runtime config · Primary YAML configuration for experiments · `config_files/configure.yaml`

### RL Control

Agent orchestration and shared runtime utilities

- [agents.py](https://github.com/realwenlongwang/drama/blob/master/agents.py): agent control · High-level agent logic tying together rollout, learning, and optimization · `agents.py`
- [replay buffer](https://github.com/realwenlongwang/drama/blob/master/replay_buffer.py): experience store · Experience storage and sampling boundary · `replay_buffer.py`
- [tools.py](https://github.com/realwenlongwang/drama/blob/master/tools.py): experiment utils · Shared experiment helpers · `tools.py`
- [utils.py](https://github.com/realwenlongwang/drama/blob/master/utils.py): shared utils · Common utility functions used across training and evaluation · `utils.py`

### Environments

Task adapters and rollout execution

- [env wrapper](https://github.com/realwenlongwang/drama/blob/master/env_wrapper.py): env api · Normalizes environment interfaces across tasks · `env_wrapper.py`
- [parallel rollout](https://github.com/realwenlongwang/drama/blob/master/envs/parallel.py): vector env · Parallel or subprocess-based environment execution · `envs/parallel.py`
- [Atari env](https://github.com/realwenlongwang/drama/blob/master/envs/my_atari.py): task adapter · Atari benchmark environment adapter · `envs/my_atari.py`
- [DMC env](https://github.com/realwenlongwang/drama/blob/master/envs/my_dmc.py): task adapter · DeepMind Control environment adapter · `envs/my_dmc.py`
- [Memory maze](https://github.com/realwenlongwang/drama/blob/master/envs/my_memory_maze.py): task adapter · Memory maze environment adapter · `envs/my_memory_maze.py`

### Drama Models

Task-specific learning modules

- [world models.py](https://github.com/realwenlongwang/drama/blob/master/sub_models/world_models.py): dynamics model · Latent dynamics and world model learning · `sub_models/world_models.py`
- [transformer model](https://github.com/realwenlongwang/drama/blob/master/sub_models/transformer_model.py): seq backbone · Alternative sequence-model backbone · `sub_models/transformer_model.py`
- [attention blocks](https://github.com/realwenlongwang/drama/blob/master/sub_models/attention_blocks.py): model internals · Attention and sequence building blocks · `sub_models/attention_blocks.py`
- [losses](https://github.com/realwenlongwang/drama/blob/master/sub_models/functions_losses.py): training losses · Model loss functions and optimization objectives · `sub_models/functions_losses.py`
- [laprop](https://github.com/realwenlongwang/drama/blob/master/sub_models/laprop.py): optimizer helper · Optimization support utilities · `sub_models/laprop.py`

### Mamba Stack

Reusable sequence model implementation

- [mamba models](https://github.com/realwenlongwang/drama/blob/master/mamba_ssm/models/mixer_seq_simple.py): model assembly · Higher-level Mamba model configuration and assembly · `mamba_ssm/models/mixer_seq_simple.py`
- [mamba blocks](https://github.com/realwenlongwang/drama/blob/master/mamba_ssm/modules/block.py): sequence modules · Core Mamba, Mamba2, attention, and MLP modules · `mamba_ssm/modules/block.py`
- [selective scan](https://github.com/realwenlongwang/drama/blob/master/mamba_ssm/ops/selective_scan_interface.py): python ops · Python boundary for selective-scan operators · `mamba_ssm/ops/selective_scan_interface.py`

### Native Kernels

Compiled CUDA/C++ and Triton operators

- [triton ops](https://github.com/realwenlongwang/drama/blob/master/mamba_ssm/ops/triton/ssd_combined.py): gpu kernels · Triton fused kernels for normalization and state updates · `mamba_ssm/ops/triton/ssd_combined.py`
- [selective scan cuda](https://github.com/realwenlongwang/drama/blob/master/csrc/selective_scan/selective_scan.cpp): cuda extension · Compiled CUDA/C++ selective-scan implementation · `csrc/selective_scan/selective_scan.cpp`

## Connections

- configure.yaml → train.py: configures · Runtime behavior and model variant selection
- configure.yaml → eval.py: configures · Evaluation settings and overrides
- train.py → agents.py: starts · Launches training orchestration
- eval.py → agents.py: runs · Reuses agent/evaluation logic
- agents.py → replay buffer: reads/writes · Collects and samples experience
- agents.py → env wrapper: interacts · Uses normalized environment APIs
- env wrapper → parallel rollout: uses · Executes rollouts in parallel
- parallel rollout → Atari env: wraps · Task-specific Atari environments
- parallel rollout → DMC env: wraps · Task-specific DMC environments
- parallel rollout → Memory maze: wraps · Task-specific memory-maze environments
- agents.py → world models.py: trains · Updates learned dynamics model
- agents.py → transformer model: trains · Alternative backbone path
- world models.py → attention blocks: uses · Relies on internal attention/sequence blocks
- world models.py → losses: optimizes with · Uses task and prediction losses
- world models.py → laprop: steps with · Uses optimization helpers
- transformer model → attention blocks: uses · Built from shared model blocks
- attention blocks → mamba models: implements · Sequence modules assembled into higher-level models
- mamba models → selective scan: calls · Invokes selective-scan operator boundary
- selective scan → selective scan cuda: dispatches · Calls compiled CUDA/C++ kernels
- mamba blocks → selective scan: depends on · Core modules use the operator boundary
- mamba models → triton ops: uses · Relies on Triton fused kernels for performance
- agents.py → tools.py: uses · Experiment helpers and logging utilities
- agents.py → utils.py: uses · Shared utility functions