p-doom/crowd-code-dataset-0.2 huggingface dataset:
[tab] test/test_nan.ipynb (python)
# Restore a dynamics checkpoint and enable sowing
import os
from typing import Dict

import jax
import jax.numpy as jnp
import flax.nnx as nnx
import optax
import orbax.checkpoint as ocp
import grain

from utils.dataloader import get_dataloader
from models.lam import LatentActionModel

# Adjust to your checkpoint directory, dataset directory, and dynamics type
ckpt_dir = "/fast/project/HFMI_SynergyUnit/jafar_ws/checkpoints/coinrun/lam/train_lam_coinrun_reproduction_20067/100000_ckpt"
data_dir = "/fast/project/HFMI_SynergyUnit/jafar_ws/data/coinrun/array_records_10m"
num_steps: int = 200_000
seed: int = 0
seq_len: int = 16
image_channels: int = 3
image_height: int = 64
image_width: int = 64
save_ckpt: bool = False
restore_ckpt: bool = False
# Optimization
batch_size: int = 36
vq_beta: float = 0.25
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_schedule: str = "wsd"  # supported options: wsd, cos
vq_reset_thresh: int = 50
# LAM
model_dim: int = 512
ffn_dim: int = 2048
latent_dim: int = 32
num_latents: int = 6
patch_size: int = 16
num_blocks: int = 4
num_heads: int = 8
dropout: float = 0.0
codebook_dropout: float = 0.0
param_dtype = jnp.float32
dtype = jnp.bfloat16
# Logging
log_interval: int = 5
log_image_interval: int = 250
use_flash_attention: bool = True

# Build model graph matching the checkpoint
rng = jax.random.key(seed)
rng, _rng = jax.random.split(rng)
rngs = nnx.Rngs(_rng)
lam = LatentActionModel(
    in_dim=image_channels,
    model_dim=model_dim,
    ffn_dim=ffn_dim,
    latent_dim=latent_dim,
    num_latents=num_latents,
    patch_size=patch_size,
    num_blocks=num_blocks,
    num_heads=num_heads,
    dropout=dropout,
    codebook_dropout=codebook_dropout,
    param_dtype=param_dtype,
    dtype=dtype,
    use_flash_attention=use_flash_attention,
    rngs=rngs,
)

# Optimizer (matches training opt hyperparams; lr value is irrelevant for restore)
tx = optax.adamw(
    learning_rate=max_lr,
    b1=0.9,
    b2=0.9,
    weight_decay=1e-4,
    mu_dtype=param_dtype,
)
optimizer = nnx.Optimizer(lam, tx)
[tab] test/test_nan.ipynb (python)
# Restore latest checkpoint: optimizer and dataloader state, like in training
from typing import cast

handler_registry = ocp.handlers.DefaultCheckpointHandlerRegistry()
handler_registry.add(
    "model_state", ocp.args.PyTreeSave, ocp.handlers.PyTreeCheckpointHandler
)
handler_registry.add(
    "model_state", ocp.args.PyTreeRestore, ocp.handlers.PyTreeCheckpointHandler
)
handler_registry.add(
    "dataloader_state",
    grain.checkpoint.CheckpointSave,
    cast(ocp.handlers.CheckpointHandler, grain.checkpoint.CheckpointHandler),
)
handler_registry.add(
    "dataloader_state",
    grain.checkpoint.CheckpointRestore,
    cast(ocp.handlers.CheckpointHandler, grain.checkpoint.CheckpointHandler),
)

ckpt_mgr = ocp.CheckpointManager(
    directory=ckpt_dir,
    options=ocp.CheckpointManagerOptions(step_format_fixed_length=6),
    handler_registry=handler_registry,
)

# Recreate dataloader and iterator exactly like training
array_record_files = [
    os.path.join(data_dir, x)
    for x in os.listdir(data_dir)
    if x.endswith(".array_record")
]
grain_dataloader = get_dataloader(
    array_record_files,
    seq_len,
    batch_size,
    image_height,
    image_width,
    image_channels,
    num_workers=8,
    prefetch_buffer_size=1,
    seed=seed,
)
initial_state = grain_dataloader._create_initial_state()
loader_iterator = grain.DataLoaderIterator(grain_dataloader, initial_state)

# Restore optimizer and dataloader iterator states
abstract_optimizer = nnx.eval_shape(lambda: optimizer)
abstract_optimizer_state = nnx.state(abstract_optimizer)
restored = ckpt_mgr.restore(
    ckpt_mgr.latest_step(),
    args=ocp.args.Composite(
        model_state=ocp.args.PyTreeRestore(abstract_optimizer_state),  # type: ignore
        dataloader_state=grain.checkpoint.CheckpointRestore(loader_iterator),  # type: ignore
    ),
)

nnx.update(optimizer, restored["model_state"])  # type: ignore
loader_iterator = restored["dataloader_state"]
step = ckpt_mgr.latest_step() or 0
ckpt_mgr.close()

# Convenience handle
lam = optimizer.model
print(f"Restored optimizer and dataloader at step {step}.")
[content] test/test_nan.ipynb (python)

    lr_schedule = get_lr_schedule(
        args.lr_schedule,
        args.init_lr,
        args.max_lr,
        args.decay_end,
        args.num_steps,
        args.warmup_steps,
        args.wsd_decay_steps,
    )