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: