p-doom/crowd-code-dataset-0.2 huggingface dataset:
[tab] requirements.txt (pip-requirements)
# core requirements
dm_pix>=0.4.3
einops>=0.8.0
flax>=0.10.7
jax[cuda12]>=0.6.2
optax>=0.2.3
tyro>=0.8.5
wandb>=0.17.4
grain>=0.2.10
array-record>=0.7.2

# data pipeline
gsutil>=5.35
ffmpeg-python==0.2.0
hf-transfer==0.1.9
huggingface-hub[cli]>=0.34.3
procgen>=0.10.7
tqdm>=4.67.1
tensorflow>=2.19.0
tensorflow-datasets>=4.12.0
rlds>=0.1.1

# dev
pre-commit>=4.2.0
[content] train_dynamics.py (python)
            if dataloader_val and step % args.val_interval == 0:
                rng, _rng_mask_val = jax.random.split(rng, 2)
                print("Calculating validation metrics...")
                val_metrics, val_gt_batch, val_recon, val_full_frame = calculate_validation_metrics(dataloader_val, optimizer.model, _rng_mask_val)
                print(f"Step {step}, validation loss: {val_metrics['val_loss']}")
                val_results = {
                    "metrics": val_metrics,
                    "gt_batch": val_gt_batch,
                    "recon": val_recon,
                    "full_frame": val_full_frame,
                }

            # --- Logging ---
            if args.log:
                if step % args.log_interval == 0 and jax.process_index() == 0:
                    log_dict = {
                        "loss": loss,
                        "step": step,
                        **metrics
                        }
                    if val_results:
                        log_dict.update(val_results["metrics"])
                    wandb.log(log_dict)
                if step % args.log_image_interval == 0:
                    gt_seq = inputs["videos"][0].astype(jnp.float32) / 255.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"
                    )
                    if val_results:
                        val_results["gt_seq_val"] = val_results["gt_batch"]["videos"][0].astype(jnp.float32) / 255.0
                        val_results["recon_seq_val"] = val_results["recon"].clip(0, 1)
                        val_comparison_seq = jnp.concatenate((val_results["gt_seq_val"], val_results["recon_seq_val"]), axis=1)
                        val_results["val_comparison_seq"] = einops.rearrange(
                            val_comparison_seq * 255, "t h w c -> h (t w) c"
                        )
                        val_results["full_frame_seq_val"] = val_results["full_frame"][0].clip(0, 1)
                        val_results["val_full_frame_comparison_seq"] = jnp.concatenate((val_results["gt_seq_val"], val_results["full_frame_seq_val"]), axis=1)
                        val_results["val_full_frame_comparison_seq"] = einops.rearrange(
                            val_results["val_full_frame_comparison_seq"] * 255, "t h w c -> h (t w) c"