quadruped_reversible_flow_model

Model Card for quadruped_reversible_flow_model

Model Description

This repository contains the pre-trained neural network weights for the Reversible Flow Adaptation architecture, applied to the Quadruped environment using the reversible_flow policy.

These checkpoints represent the fully-converged policy at the final epoch (Epoch 600) across 3 independent random seeds (seed_1, seed_2, seed_42).

Architecture

The policy consists of a dual-objective architecture:

1. Vision Transformer (ViT): Acts as the state encoder, augmented with an Auxiliary Physics Distillation Head to predict privileged parameters (like mass or friction).
2. 1D U-Net Vector Field: A deterministic, parallelized Optimal Transport Flow Matching (Rectified Flows) model providing $O(1)$ generative inference latency.

Usage

The checkpoints are saved in the highly-optimized JAX / Orbax format.

import orbax.checkpoint as ocp

# Example: Loading seed 42
checkpoint_path = "seed_42"
mngr = ocp.CheckpointManager(os.path.abspath(checkpoint_path))

# Load the parameters
restored = mngr.restore(mngr.latest_step())
print(restored.keys())