update experiment train evals

BricksRL · Sep 27, 2024 · 52ef271 · 52ef271
1 parent 64328b7
commit 52ef271
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/experiments/2wheeler/eval.py b/experiments/2wheeler/eval.py
@@ -17,7 +17,7 @@
 
 from environments import make_env
 from src.agents import get_agent
-from src.utils import login, setup_check , logout
+from src.utils import login, logout, setup_check
 
 
 @hydra.main(version_base=None, config_path=project_root + "/conf", config_name="config")

diff --git a/experiments/roboarm/eval.py b/experiments/roboarm/eval.py
@@ -42,7 +42,7 @@ def run(cfg: DictConfig) -> None:
     _ = input("Press Enter to start evaluation...")
     try:
         for e in tqdm(range(eval_episodes), desc="Evaluation"):
-            td = env.reset(env.get_reset_tensordict())
+            td = env.reset()
             done = td.get("done", False)
             truncated = td.get("truncated", False)
             ep_return = 0

diff --git a/experiments/roboarm/train.py b/experiments/roboarm/train.py
@@ -56,7 +56,7 @@ def run(cfg: DictConfig) -> None:
     train_episodes = cfg.episodes
     max_episode_steps = cfg.env.max_episode_steps
 
-    print("Start training...")  
+    print("Start training...")
     quit = False
     try:
         for e in tqdm(range(train_episodes), desc="Training"):