def build_and_train(slot_affinity_code, log_dir, run_ID, config_key):
if slot_affinity_code is 'None':
# affinity = affinity_from_code(run_slot_affinity_code)
slot_affinity_code = prepend_run_slot(0, affinity_code)
affinity = affinity_from_code(slot_affinity_code)
else:
affinity = affinity_from_code(slot_affinity_code)
config = configs[config_key]
# load variant of experiment (there may not be a variant, though)
variant = load_variant(log_dir)
config = update_config(config, variant)
sampler = CpuSampler(
EnvCls=make_env,
env_kwargs={},
CollectorCls=CpuResetCollector,
**config["sampler"]
)
algo = PPO(optim_kwargs=config["optim"], **config["algo"])
agent = MujocoFfAgent(model_kwargs=config["model"], **config["agent"])
runner = MinibatchRl(
algo=algo,
agent=agent,
sampler=sampler,
affinity=affinity,
**config["runner"]
)
name = config["env"]["id"]
with logger_context(log_dir, run_ID, name, config):
runner.train()
def launch_experiment(script, run_slot, affinity_code, log_dir, variant,
run_ID, args):
"""Launches one learning run using ``subprocess.Popen()`` to call the
python script. Calls the script as:
``python {script} {slot_affinity_code} {log_dir} {run_ID} {*args}``
If ``affinity_code["all_cpus"]`` is provided, then the call is prepended
with ``tasket -c ..`` and the listed cpus (this is the most sure way to
keep the run limited to these CPU cores). Also saves the `variant` file.
Returns the process handle, which can be monitored.
"""
slot_affinity_code = prepend_run_slot(run_slot, affinity_code)
affinity = affinity_from_code(slot_affinity_code)
call_list = list()
if isinstance(affinity, dict) and affinity.get("all_cpus", False):
cpus = ",".join(str(c) for c in affinity["all_cpus"])
elif isinstance(affinity, list) and affinity[0].get("all_cpus", False):
cpus = ",".join(str(c) for aff in affinity for c in aff["all_cpus"])
else:
cpus = ()
if cpus:
call_list += ["taskset", "-c",
cpus] # PyTorch obeys better than just psutil.
call_list += ["python", script, slot_affinity_code, log_dir, str(run_ID)]
call_list += [str(a) for a in args]
save_variant(variant, log_dir)
print("\ncall string:\n", " ".join(call_list))
p = subprocess.Popen(call_list)
return p
def launch_experiment(
script,
run_slot,
affinity_code,
log_dir,
variant,
run_ID,
args,
python_executable=None,
set_egl_device=False,
):
"""Launches one learning run using ``subprocess.Popen()`` to call the
python script. Calls the script as:
``python {script} {slot_affinity_code} {log_dir} {run_ID} {*args}``
If ``affinity_code["all_cpus"]`` is provided, then the call is prepended
with ``tasket -c ..`` and the listed cpus (this is the most sure way to
keep the run limited to these CPU cores). Also saves the `variant` file.
Returns the process handle, which can be monitored.
Use ``set_egl_device=True`` to set an environment variable
``EGL_DEVICE_ID`` equal to the same value as the cuda index for the
algorithm. For example, can use with DMControl environment modified
to look for this environment variable when selecting a GPU for headless
rendering.
"""
slot_affinity_code = prepend_run_slot(run_slot, affinity_code)
affinity = affinity_from_code(slot_affinity_code)
call_list = list()
if isinstance(affinity, dict) and affinity.get("all_cpus", False):
cpus = ",".join(str(c) for c in affinity["all_cpus"])
elif isinstance(affinity, list) and affinity[0].get("all_cpus", False):
cpus = ",".join(str(c) for aff in affinity for c in aff["all_cpus"])
else:
cpus = ()
if cpus:
call_list += ["taskset", "-c",
cpus] # PyTorch obeys better than just psutil.
py = python_executable if python_executable else sys.executable or "python"
call_list += [py, script, slot_affinity_code, log_dir, str(run_ID)]
call_list += [str(a) for a in args]
save_variant(variant, log_dir)
print("\ncall string:\n", " ".join(call_list))
print(os.getcwd())
if set_egl_device and affinity.get("cuda_idx", None) is not None:
egl_device_id = str(affinity["cuda_idx"])
egl_env = os.environ.copy()
egl_env["EGL_DEVICE_ID"] = egl_device_id
print(f"Assigning EGL_DEVICE_ID={egl_device_id}")
p = subprocess.Popen(call_list, env=egl_env)
else:
p = subprocess.Popen(call_list)
return p
def launch_experiment(script, run_slot, affinity_code, log_dir, variant,
run_ID, args):
slot_affinity_code = prepend_run_slot(run_slot, affinity_code)
affinity = affinity_from_code(slot_affinity_code)
call_list = list()
if isinstance(affinity, dict) and affinity.get("all_cpus", False):
cpus = ",".join(str(c) for c in affinity["all_cpus"])
elif isinstance(affinity, list) and affinity[0].get("all_cpus", False):
cpus = ",".join(str(c) for aff in affinity for c in aff["all_cpus"])
else:
cpus = ()
if cpus:
call_list += ["taskset", "-c",
cpus] # PyTorch obeys better than just psutil.
call_list += ["python", script, slot_affinity_code, log_dir, str(run_ID)]
call_list += [str(a) for a in args]
save_variant(variant, log_dir)
print("\ncall string:\n", " ".join(call_list))
p = subprocess.Popen(call_list)
return p
def build_and_train(log_dir, run_ID, config_key):
# affinity = affinity_from_code(run_slot_affinity_code)
slot_affinity_code = prepend_run_slot(0, affinity_code)
affinity = affinity_from_code(slot_affinity_code)
config = configs[config_key]
sampler = CpuSampler(EnvCls=make_env,
env_kwargs={},
CollectorCls=CpuResetCollector,
**config["sampler"])
algo = PPO(optim_kwargs=config["optim"], **config["algo"])
agent = MultiFfAgent(model_kwargs=config["model"], **config["agent"])
runner = MinibatchRl(algo=algo,
agent=agent,
sampler=sampler,
affinity=affinity,
**config["runner"])
name = config["env"]["id"]
with logger_context(log_dir, run_ID, name, config):
runner.train()
def build_and_train(game="pong", run_ID=0):
# Seems like we should be able to skip the intermediate step of the code,
# but so far have just always run that way.
# Change these inputs to match local machine and desired parallelism.
affinity_code = encode_affinity(
n_cpu_cores=16, # Use 16 cores across all experiments.
n_gpu=8, # Use 8 gpus across all experiments.
hyperthread_offset=24, # If machine has 24 cores.
n_socket=2, # Presume CPU socket affinity to lower/upper half GPUs.
gpu_per_run=2, # How many GPUs to parallelize one run across.
# cpu_per_run=1,
)
slot_affinity_code = prepend_run_slot(run_slot=0, affinity_code=affinity_code)
affinity = get_affinity(slot_affinity_code)
breakpoint()
sampler = GpuParallelSampler(
EnvCls=AtariEnv,
env_kwargs=dict(game=game),
CollectorCls=WaitResetCollector,
batch_T=5,
batch_B=16,
max_decorrelation_steps=400,
)
algo = A2C() # Run with defaults.
agent = AtariFfAgent()
runner = MultiGpuRl(
algo=algo,
agent=agent,
sampler=sampler,
n_steps=50e6,
log_interval_steps=1e5,
affinity=affinity,
)
config = dict(game=game)
name = "a2c_" + game
log_dir = "example_7"
with logger_context(log_dir, run_ID, name, config):
runner.train()
parser.add_argument("--savepath", type=str, default="./vae_data/experiment/")
args = parser.parse_args()
os.makedirs(args.savepath, exist_ok=True)
EmptyAgentInfo = namedarraytuple("EmptyAgentInfo", [])
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
affinity_code = encode_affinity(
n_cpu_core=4,
n_gpu=1,
# hyperthread_offset=20,
n_socket=2
# cpu_per_run=2,
)
affinity = affinity_from_code(prepend_run_slot(0, affinity_code))
# Get Params
config = configs["pretrain"]
with open('%s/params.json' % args.savepath, 'w') as fp:
json.dump(config, fp, indent=4, sort_keys=True)
# Setup Data Collection
# The below are util classes
class RandomAgent(BaseAgent):
def step(self, observation, prev_action, prev_reward):
model_inputs = buffer_to((observation, prev_action, prev_reward),
device=self.device)
action = self.model(*model_inputs)
agent_info = EmptyAgentInfo()
action, agent_info = buffer_to((action, agent_info), device="cpu")
