From d0a3e50469238f5ab18107adafdecfb1a53af688 Mon Sep 17 00:00:00 2001 From: Abhishek Joshi Date: Wed, 23 Oct 2024 20:54:03 -0400 Subject: [PATCH] Adding env_steps argument --- tests/test_robots/test_composite_robots.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/tests/test_robots/test_composite_robots.py b/tests/test_robots/test_composite_robots.py index 1c53a9ce32..c72892598b 100644 --- a/tests/test_robots/test_composite_robots.py +++ b/tests/test_robots/test_composite_robots.py @@ -30,8 +30,15 @@ "SpotFloating", ] - -def create_and_test_env(env: str, robots: Union[str, List[str]], controller_config: Dict, render=True): +# If you would like to visualize the scene during testing, +# set render to True and increase env_steps to a larger value. +def create_and_test_env( + env: str, + robots: Union[str, List[str]], + controller_config: Dict, + render: bool = True, + env_steps: int = 20, +): config = { "env_name": env, @@ -54,7 +61,7 @@ def create_and_test_env(env: str, robots: Union[str, List[str]], controller_conf high = np.clip(high, -1, 1) # Runs a few steps of the simulation as a sanity check - for i in range(200): + for i in range(env_steps): action = np.random.uniform(low, high) obs, reward, done, _ = env.step(action) if render: