feat(hpc): rename config_path to env_config_path and support YAML configs
This commit is contained in:
parent
457a23e5cb
commit
ed20325fef
5 changed files with 20 additions and 14 deletions
|
|
@ -31,11 +31,11 @@ def convert_obs_dict_to_array(obs_dict: dict) -> jnp.ndarray:
|
|||
)
|
||||
|
||||
|
||||
def make_env(config_path: str | None, num_envs: int) -> Callable:
|
||||
def make_env(env_config_path: str | None, num_envs: int) -> Callable:
|
||||
def thunk():
|
||||
if config_path is None:
|
||||
if env_config_path is None:
|
||||
return BrittleStarJaxEnvWrapper.default(num_envs=num_envs)
|
||||
return BrittleStarJaxEnvWrapper.from_config(config_path, num_envs=num_envs)
|
||||
return BrittleStarJaxEnvWrapper.from_config(env_config_path, num_envs=num_envs)
|
||||
|
||||
return thunk
|
||||
|
||||
|
|
@ -88,7 +88,7 @@ def train(args: PPOArgs):
|
|||
print(f"Running on device: {device}")
|
||||
|
||||
print("Creating the environment...")
|
||||
env = make_env(config_path=args.config_path, num_envs=args.num_envs)()
|
||||
env = make_env(env_config_path=args.env_config_path, num_envs=args.num_envs)()
|
||||
print(f"Environment: {env}")
|
||||
|
||||
episode_stats = EpisodeStatistics(
|
||||
|
|
|
|||
Reference in a new issue