equidiff/test.py

import subprocess

n_action_steps_values = [1, 2, 4, 8, 16]
cuda_device = 7

command_template = (
    "CUDA_VISIBLE_DEVICES={cuda_device} MUJOCO_GL=osmesa PYOPENGL_PLATFORM=osmesa "
    "HYDRA_FULL_ERROR=1 python train.py --config-name=test_sq2 "
    "n_action_steps={n_action_steps}"
)

for n_action_steps in n_action_steps_values:
    command = command_template.format(cuda_device = cuda_device, n_action_steps=n_action_steps)
    
    print(f"Running command: {command}")
    
    try:
        subprocess.run(command, shell=True, check=True)
        print(f"Command for n_action_steps={n_action_steps} executed successfully!")
    except subprocess.CalledProcessError as e:
        print(f"Error occurred while executing command for n_action_steps={n_action_steps}: {e}")