fix bugs

Author: yardenas

train_brax.py

@@ -200,17 +200,17 @@ def main(cfg):
       f"\n{OmegaConf.to_yaml(cfg)}"
   )
   logger = WeightAndBiasesWriter(cfg)
-  if cfg.agent_name == "SAC":
+  if cfg.training.agent_name == "SAC":
     train_fn = get_sac_train_fn()
-  elif cfg.agent_name == "PPO":
+  elif cfg.training.agent_name == "PPO":
     train_fn = get_ppo_train_fn()
   else:
     raise NotImplementedError
   rng = jax.random.PRNGKey(cfg.training.seed)
   steps = Counter()
-  env = registry.load(cfg.task_name)
-  env_cfg = registry.get_default_config(cfg.task_name)
-  eval_env = registry.load(cfg.task_name, config=env_cfg)
+  env = registry.load(cfg.training.task_name)
+  env_cfg = registry.get_default_config(cfg.training.task_name)
+  eval_env = registry.load(cfg.training.task_name, config=env_cfg)
   with jax.disable_jit(not cfg.jit):
     make_policy, params, _ = train_fn(
         environment=env,