Spaces:
Running on Zero
Running on Zero
| import sys | |
| import os | |
| from utils.utils import run_command, find_latest_checkpoint, prepare_pipeline_run_root | |
| from utils.config import TrainingConfig, load_config | |
| DEFAULT_TRAINING_CONFIG = os.path.join(os.getcwd(), 'configs', 'training.yaml') | |
| def main(): | |
| # load training config | |
| argv = sys.argv | |
| try: | |
| idx = argv.index('--config') | |
| training_cfg_path = argv[idx + 1] if idx + 1 < len(argv) else DEFAULT_TRAINING_CONFIG | |
| except ValueError: | |
| print(f"No training config path provided, using default: {DEFAULT_TRAINING_CONFIG}") | |
| training_cfg_path = DEFAULT_TRAINING_CONFIG | |
| train_config: TrainingConfig = load_config(TrainingConfig, default_config_path=training_cfg_path) | |
| # torchrun if distributed, else python | |
| if train_config.nproc_per_node > 1: | |
| launcher = [ | |
| "torchrun", | |
| "--nproc_per_node", str(train_config.nproc_per_node), | |
| ] | |
| if train_config.standalone: | |
| launcher += ["--standalone"] | |
| else: | |
| launcher = [sys.executable] | |
| # top-level run root and export child processes can use | |
| run_root, run_name = prepare_pipeline_run_root(base_cwd=os.getcwd()) | |
| os.environ['NG_RUN_ROOT_DIR'] = run_root | |
| if train_config.run_video_tokenizer: | |
| v_cmd = launcher + [ | |
| "scripts/train_video_tokenizer.py", | |
| "--config", train_config.video_tokenizer_config, | |
| "--training_config", training_cfg_path, | |
| ] | |
| if not run_command(v_cmd, "Video Tokenizer Training"): | |
| return | |
| if train_config.run_latent_actions: | |
| latent_actions_cmd = launcher + [ | |
| "scripts/train_latent_actions.py", | |
| "--config", train_config.latent_actions_config, | |
| "--training_config", training_cfg_path, | |
| ] | |
| if not run_command(latent_actions_cmd, "Latent Actions Training"): | |
| return | |
| # need to get above checkpoints and pass in to dynamics | |
| video_tokenizer_checkpoint = find_latest_checkpoint(".", "video_tokenizer") | |
| latent_actions_checkpoint = find_latest_checkpoint(".", "latent_actions") | |
| if train_config.run_dynamics: | |
| dyn_cmd = launcher + [ | |
| "scripts/train_dynamics.py", | |
| "--config", train_config.dynamics_config, | |
| "--training_config", training_cfg_path, | |
| f"video_tokenizer_path={video_tokenizer_checkpoint}", | |
| f"latent_actions_path={latent_actions_checkpoint}", | |
| ] | |
| if not run_command(dyn_cmd, "Dynamics Model Training"): | |
| return | |
| dynamics_checkpoint = find_latest_checkpoint(".", "dynamics") | |
| print("\n📁 Results Summary:") | |
| print(f"Video Tokenizer: {video_tokenizer_checkpoint}") | |
| print(f"Latent Actions: {latent_actions_checkpoint}") | |
| print(f"Dynamics Model: {dynamics_checkpoint}") | |
| if __name__ == "__main__": | |
| main() |