[tab] train_dynamics.py (python)
from dataclasses import dataclass, field
import os
import time

import einops
from flax.training import orbax_utils
from flax.training.train_state import TrainState
from jax.sharding import Mesh, PartitionSpec, NamedSharding
from jax.experimental.mesh_utils import create_device_mesh
import optax
import orbax.checkpoint as ocp
import numpy as np
import dm_pix as pix
import jax
import jax.numpy as jnp
import tyro
import wandb
from etils import epath

from models.dynamics import DynamicsModel
from utils.dataloader import get_dataloader
from utils.parameter_utils import count_parameters_by_component

ts = int(time.time())


@dataclass
class Args:
    # Experiment
    num_steps: int = 300_000
    seed: int = 0
    seq_len: int = 16
    image_channels: int = 3
    image_height: int = 90
    image_width: int = 160
    data_dir: str = ""
    checkpoint: str = ""
    # Optimization
    batch_size: int = 48
    min_lr: float = 3e-4
    max_lr: float = 3e-4
    warmup_steps: int = 10000
    # Model
    model_dim: int = 512
    num_blocks: int = 8
    num_heads: int = 8
    dropout: float = 0.0
    # Logging
    log: bool = False
    entity: str = ""
    project: str = ""
    name: str = "train_dynamics"
    tags: list[str] = field(default_factory=lambda: ["dynamics"])
    log_interval: int = 5
    log_image_interval: int = 250
    ckpt_dir: str = ""
    log_checkpoint_interval: int = 10000
    log_gradients: bool = False


args = tyro.cli(Args)


def dynamics_loss_fn(params, state, inputs):
    # --- Compute loss ---
    outputs = state.apply_fn(
        params,
        inputs,
        training=True,
        rngs={"params": inputs["rng"], "dropout": inputs["dropout_rng"]},
    )
    mse = jnp.square(inputs["videos"] - outputs["recon"]).mean()
    loss = mse

    # --- Compute validation metrics ---
    gt = inputs["videos"].clip(0, 1).reshape(-1, *inputs["videos"].shape[2:])
    recon = outputs["recon"].clip(0, 1).reshape(-1, *outputs["recon"].shape[2:])
    psnr = pix.psnr(gt, recon).mean()
    ssim = pix.ssim(gt, recon).mean()
    metrics = dict(
        loss=loss,
        mse=mse,
        psnr=psnr,
        ssim=ssim,
    )
    return loss, (outputs["recon"], metrics)


@jax.jit
def train_step(state, inputs):
    grad_fn = jax.value_and_grad(dynamics_loss_fn, has_aux=True, allow_int=True)
    (loss, (recon, metrics)), grads = grad_fn(state.params, state, inputs)
    state = state.apply_gradients(grads=grads)
    if args.log_gradients:
        metrics["gradients_std/"] = jax.tree.map(
            lambda x: x.std(), grads["params"]
        )
    return state, loss, recon, metrics


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

    if args.batch_size % num_devices != 0:
        raise ValueError(
            f"Global batch size {args.batch_size} must be divisible by "
            f"number of devices {num_devices}."
        )

    per_device_batch_size_for_init = args.batch_size // num_devices

    rng = jax.random.PRNGKey(args.seed)

    # --- Initialize model ---
    dynamics = DynamicsModel(
        in_dim=args.image_channels,
        model_dim=args.model_dim,
        num_blocks=args.num_blocks,
        num_heads=args.num_heads,
        dropout=args.dropout,
    )
    rng, _rng = jax.random.split(rng)
    image_shape = (args.image_height, args.image_width, args.image_channels)
    inputs = dict(
        videos=jnp.zeros(
            (per_device_batch_size_for_init, args.seq_len, *image_shape),
            dtype=jnp.float32,
        ),
    )
    init_params = dynamics.init(_rng, inputs)

    param_counts = count_parameters_by_component(init_params)

    if args.log and jax.process_index() == 0:
        wandb.init(
            entity=args.entity,
            project=args.project,
            name=args.name,
            tags=args.tags,
            group="debug",
            config=args,
        )
        wandb.config.update({"model_param_count": param_counts})

    print("Parameter counts:")
    print(param_counts)

    # --- Initialize optimizer ---
    lr_schedule = optax.warmup_cosine_decay_schedule(
        args.min_lr, args.max_lr, args.warmup_steps, args.num_steps
    )
    tx = optax.adamw(learning_rate=lr_schedule, b1=0.9, b2=0.9, weight_decay=1e-4)
    train_state = TrainState.create(apply_fn=dynamics.apply, params=init_params, tx=tx)

    # FIXME: switch to create_hybrid_device_mesh for runs spanning multiple nodes
    device_mesh_arr = create_device_mesh((num_devices,))
    mesh = Mesh(devices=device_mesh_arr, axis_names=("data",))

    replicated_sharding = NamedSharding(mesh, PartitionSpec())
    videos_sharding = NamedSharding(
        mesh, PartitionSpec("data", None, None, None, None)
    )
    train_state = jax.device_put(train_state, replicated_sharding)

    # --- Initialize checkpoint manager with CompositeCheckpointHandler ---
    checkpoint_dir = epath.Path(args.ckpt_dir) if args.ckpt_dir else epath.Path(f"checkpoints/dynamics_{ts}")
    checkpoint_dir.mkdir(parents=True, exist_ok=True)
    
    # Create handler registry for proper restoration
    handler_registry = ocp.handlers.DefaultCheckpointHandlerRegistry()
    handler_registry.add('model_state', ocp.args.StandardSave)
    handler_registry.add('model_state', ocp.args.StandardRestore)
    handler_registry.add('dataloader_state', ocp.args.StandardSave)
    handler_registry.add('dataloader_state', ocp.args.StandardRestore)
    
    checkpoint_options = ocp.CheckpointManagerOptions(
        save_interval_steps=args.log_checkpoint_interval,
        max_to_keep=3,
        step_format_fixed_length=6,
        cleanup_tmp_directories=True,
    )
    
    checkpoint_manager = ocp.CheckpointManager(
        checkpoint_dir,
        options=checkpoint_options,
        handler_registry=handler_registry,
    )

    # --- Load checkpoint ---
    step = 0
    grain_dataloader = None
    
    if args.checkpoint:
        # Restore both model and dataloader state using CompositeCheckpointHandler
        try:
            # Create abstract state for restoration
            abstract_train_state = jax.tree.map(ocp.tree.to_shape_dtype_struct, train_state)
            
            restored = checkpoint_manager.restore(
                checkpoint_manager.latest_step(),
                args=ocp.args.Composite(
                    model_state=ocp.args.StandardRestore(abstract_train_state),
                    dataloader_state=None,  # Will be restored separately if exists
                )
            )
            
            # Update train state with restored parameters
            train_state = restored["model_state"]
            
            # Get step from checkpoint
            step = checkpoint_manager.latest_step()
            print(f"Restored model state from step {step}")
            
        except Exception as e:
            print(f"Could not restore from checkpoint: {e}")
            print("Starting from scratch...")

    # --- TRAIN LOOP ---
    array_record_files = [
        os.path.join(args.data_dir, x)
        for x in os.listdir(args.data_dir)
        if x.endswith(".array_record")
    ]
    
    # Create Grain dataloader
    grain_dataloader = get_dataloader(
        # NOTE: We deliberately pass the global batch size
        # The dataloader shards the dataset across all processes
        array_record_files,
        args.seq_len,
        args.batch_size,
        *image_shape,
        num_workers=8,
        prefetch_buffer_size=1,
        seed=args.seed,
    )
    
    # Check if dataloader state was also restored
    if "dataloader_state" in restored and restored["dataloader_state"] is not None:
        grain_dataloader = restored["dataloader_state"]
        print(f"Restored dataloader state from step {step}")
    else:
        print("No dataloader state found in checkpoint, using fresh dataloader")
    
    # Convert to JAX arrays
    dataloader = (jax.make_array_from_process_local_data(videos_sharding, elem) for elem in grain_dataloader) # type: ignore
    
    print(f"Starting training from step {step}...")
    while step < args.num_steps:
        for videos in dataloader:
            # --- Train step ---
            rng, _rng, _rng_dropout = jax.random.split(rng, 3)

            inputs = dict(videos=videos, rng=_rng, dropout_rng=_rng_dropout)
            train_state, loss, recon, metrics = train_step(train_state, inputs)
            print(f"Step {step}, loss: {loss}")
            step += 1

            # --- Logging ---
            if args.log:
                if step % args.log_interval == 0 and jax.process_index() == 0:
                    wandb.log(
                        {
                            "loss": loss,
                            "step": step,
                            **metrics,
                        }
                    )
                if step % args.log_image_interval == 0:
                    gt_seq = inputs["videos"][0]
                    recon_seq = recon[0].clip(0, 1)
                    comparison_seq = jnp.concatenate((gt_seq, recon_seq), axis=1)
                    comparison_seq = einops.rearrange(
                        comparison_seq * 255, "t h w c -> h (t w) c"
                    )
                    # NOTE: Process-dependent control flow deliberately happens
                    # after indexing operation since it must not contain code
                    # sections that lead to cross-accelerator communication.
                    if jax.process_index() == 0:
                        log_images = dict(
                            image=wandb.Image(np.asarray(gt_seq[0])),
                            recon=wandb.Image(np.asarray(recon_seq[0])),
                            true_vs_recon=wandb.Image(
                                np.asarray(comparison_seq.astype(np.uint8))
                            ),
                        )
                        wandb.log(log_images)
            
            # --- Checkpointing ---
            if step % args.log_checkpoint_interval == 0:
                # Save both model and dataloader state using CompositeCheckpointHandler
                checkpoint_manager.save(
                    step,
                    args=ocp.args.Composite(
                        model_state=ocp.args.StandardSave(train_state),
                        dataloader_state=ocp.args.StandardSave(grain_dataloader),
                    )
                )
                checkpoint_manager.wait_until_finished()
                print(f"Saved checkpoint at step {step}")
            
            if step >= args.num_steps:
                break
    
    checkpoint_manager.close()