1
Fork 0

feat(hpc): rename config_path to env_config_path and support YAML configs

This commit is contained in:
Tibo De Peuter 2026-04-04 19:07:01 +02:00
parent 457a23e5cb
commit ed20325fef
5 changed files with 20 additions and 14 deletions

View file

@ -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(