[tab] test.ipynb (python)
import jax 
import jax.numpy as jnp
import orbax.checkpoint as ocp
from flax import nnx
from genie import Genie
import optax

from utils.lr_utils import get_lr_schedule


num_steps: int = 200_000
seed: int = 0
seq_len: int = 16
image_channels: int = 3
image_height: int = 90
image_width: int = 160
data_dir: str = ""
save_ckpt: bool = False
restore_ckpt: bool = False
# Optimization
batch_size: int = 36
init_lr: float = 0.0
max_lr: float = 3e-5
decay_end: float = 0.0
wsd_decay_steps: int = (
    10000  # NOTE: wsd_decay_steps will only be used when using a wsd-schedule
)
warmup_steps: int = 5000
lr_schedulee: str = "wsd"  # supported options: wsd, cos
# Tokenizer
tokenizer_dim: int = 512
tokenizer_ffn_dim: int = 2048
latent_patch_dim: int = 32
num_patch_latents: int = 1024
patch_size: int = 4
tokenizer_num_blocks: int = 4
tokenizer_num_heads: int = 8
tokenizer_checkpoint: str = ""
# LAM
lam_dim: int = 512
lam_ffn_dim: int = 2048
latent_action_dim: int = 32
num_latent_actions: int = 6
lam_patch_size: int = 16
lam_num_blocks: int = 4
lam_num_heads: int = 8
lam_checkpoint: str = ""
# Dynamics
dyna_dim: int = 512
dyna_ffn_dim: int = 2048
dyna_num_blocks: int = 6
dyna_num_heads: int = 8
dropout: float = 0.0
mask_limit: float = 0.5
param_dtype = jnp.float32
dtype = jnp.bfloat16
use_flash_attention: bool = True
# Logging
log: bool = False
entity: str = ""
project: str = ""
name: str = "train_dynamics"
log_interval: int = 5
log_image_interval: int = 250
ckpt_dir: str = "/home/franz.srambical/jafar/checkpoints/dynamics_openai_grain_tok_lam_restore/000030/model_state"
log_checkpoint_interval: int = 25000
log_checkpoint_keep_period: int = 20000
log_gradients: bool = False
wandb_id: str = ""

num_devices = jax.device_count()
if num_devices == 0:
    raise ValueError("No JAX devices found.")
print(f"Running on {num_devices} devices.")

rng = jax.random.PRNGKey(0)

# --- Initialize model ---
rng, _rng = jax.random.split(rng)
rngs = nnx.Rngs(_rng)
genie = Genie(
    # Tokenizer
    in_dim=image_channels,
    tokenizer_dim=tokenizer_dim,
    tokenizer_ffn_dim=tokenizer_ffn_dim,
    latent_patch_dim=latent_patch_dim,
    num_patch_latents=num_patch_latents,
    patch_size=patch_size,
    tokenizer_num_blocks=tokenizer_num_blocks,
    tokenizer_num_heads=tokenizer_num_heads,
    # LAM
    lam_dim=lam_dim,
    lam_ffn_dim=lam_ffn_dim,
    latent_action_dim=latent_action_dim,
    num_latent_actions=num_latent_actions,
    lam_patch_size=lam_patch_size,
    lam_num_blocks=lam_num_blocks,
    lam_num_heads=lam_num_heads,
    lam_co_train=not lam_checkpoint,
    # Dynamics
    dyna_dim=dyna_dim,
    dyna_ffn_dim=dyna_ffn_dim,
    dyna_num_blocks=dyna_num_blocks,
    dyna_num_heads=dyna_num_heads,
    dropout=dropout,
    mask_limit=mask_limit,
    param_dtype=param_dtype,
    dtype=dtype,
    use_flash_attention=use_flash_attention,
    rngs=rngs,
)
lr_schedule = get_lr_schedule(
    lr_schedulee,
    init_lr,
    max_lr,
    decay_end,
    num_steps,
    warmup_steps,
    wsd_decay_steps,
)
tx = optax.adamw(
    learning_rate=lr_schedule,
    b1=0.9,
    b2=0.9,
    weight_decay=1e-4,
    mu_dtype=dtype,
)
optimizer = nnx.Optimizer(genie, tx)

# optimizer.model.tokenizer.vq.drop = nnx.Dropout(0.1, rngs=rngs)
abstract_optimizer = nnx.eval_shape(lambda: optimizer)
abstract_optimizer_state = nnx.state(abstract_optimizer)

checkpointer = ocp.PyTreeCheckpointer()
restored = checkpointer.restore(ckpt_dir)