#!/usr/bin/env python3 """ Submit DMC (mujoco_playground) experiments for REPPO """ import subprocess import time import argparse # List of 23 DMC tasks from experiment plan DMC_TASKS = [ "AcrobotSwingup", "CartpoleBalance", "CartpoleSwingup", "CheetahRun", "FingerSpin", "FingerTurnEasy", "FingerTurnHard", "FishUpright", "FishSwim", "HopperStand", "HopperHop", "HumanoidStand", "HumanoidWalk", "HumanoidRun", "ManipulatorBringBall", "PendulumSwingup", "PointMassEasy", "ReacherEasy", "ReacherHard", "SwimmerSwimmer6", "SwimmerSwimmer15", "WalkerStand", "WalkerWalk", "WalkerRun" ] def submit_job(env_name, seed=0): """Submit a single DMC job""" cmd = [ "sbatch", f"--job-name=reppo_dmc_{env_name}_seed{seed}", "slurm/run_reppo_dmc_prod.sh" ] env_vars = { "ENV_NAME": env_name, "SEED": str(seed) } print(f"Submitting {env_name} (seed {seed})...") try: result = subprocess.run(cmd, env={**subprocess.os.environ, **env_vars}, capture_output=True, text=True, check=True) job_id = result.stdout.strip().split()[-1] print(f" -> Job ID: {job_id}") return job_id except subprocess.CalledProcessError as e: print(f" -> Error: {e}") print(f" -> Stdout: {e.stdout}") print(f" -> Stderr: {e.stderr}") return None def main(): parser = argparse.ArgumentParser(description="Submit DMC experiments") parser.add_argument("--seeds", type=int, default=5, help="Number of seeds to run") parser.add_argument("--tasks", nargs="+", default=DMC_TASKS, help="List of tasks to run") parser.add_argument("--delay", type=float, default=1.0, help="Delay between submissions (seconds)") args = parser.parse_args() print(f"Submitting {len(args.tasks)} DMC tasks with {args.seeds} seeds each") print(f"Total jobs: {len(args.tasks) * args.seeds}") print() job_ids = [] for task in args.tasks: for seed in range(args.seeds): job_id = submit_job(task, seed) if job_id: job_ids.append(job_id) # Add delay to avoid overwhelming the scheduler time.sleep(args.delay) print(f"\nSubmitted {len(job_ids)} jobs successfully:") for i, job_id in enumerate(job_ids): print(f" {i+1}: {job_id}") print(f"\nMonitor with: squeue -u $USER") print(f"Check logs in: logs/") if __name__ == "__main__": main()