theayos/

bt0331-save-bt-kinetix

Installation

Pre-trained checkpoints and data

Reproduce results

1. Train expert policies: uv run src/train_expert.py
   • By default, this will train 8 seeds per level for 65 million environment steps each.
   • Checkpoints, videos, and stats are written to a wandb project called rtc-kinetix-expert and the local directory ./logs-expert/<wandb-run-name>. It is recommended to control other wandb options, like the run name, using environment variables.
2. Generate data: uv run src/generate_data.py --config.run-path ./logs-expert/<wandb-run-name>
   • For each level, this will automatically load the best-performing checkpoint for each seed (discarding seeds that didn't reach a certain success threshold).
   • By default, 1 million environment steps are collected for each level using a mixture of expert policies.
   • Data is written back to ./logs-expert/<wandb-run-name>/data/.
3. Train imitation learning policies: uv run src/train_flow.py --config.run-path ./logs-expert/<wandb-run-name>
   • This will load the data from step 2 and train flow matching policies for each level.
   • Checkpoints, videos, and stats are written to a wandb project called rtc-kinetix-bc and the local directory ./logs-bc/<wandb-run-name>. It is recommended to control other wandb options, like the run name, using environment variables.
4. Evaluate imitation learning policies: uv run src/eval_flow.py --config.run-path ./logs-bc/<wandb-run-name> --output-dir <output-dir>
   • This will load the checkpoints from step 3 and evaluate them for 2048 trials per level by default.
   • Currently, the script performs an exhaustive sweep over inference delay and execution horizon for all methods.

Training-Time RTC

1. Change simulated_delay in the model config to 5.
2. Fine-tune the pre-trained checkpoint with simulated delay for 8 epochs: uv run src/train_flow.py --config.run-path <run_path> --config.load-dir bc/24 --config.num-epochs 8 where bc is the contents of gs://rtc-assets/bc/.
3. Evaluate as above.