Fixes in record_dataset and import gym_hil

This commit is contained in:
Michel Aractingi
2025-05-09 12:00:21 +02:00
parent e22411ff22
commit fb9bb89cb4

View File

@@ -1794,7 +1794,7 @@ def make_robot_env(cfg) -> gym.vector.VectorEnv:
A vectorized gym environment with all necessary wrappers applied. A vectorized gym environment with all necessary wrappers applied.
""" """
if cfg.type == "hil": if cfg.type == "hil":
import gymnasium as gym import gym_hil # noqa: F401
# TODO (azouitine) # TODO (azouitine)
env = gym.make( env = gym.make(
@@ -1923,7 +1923,7 @@ def init_reward_classifier(cfg):
########################################################### ###########################################################
def record_dataset(env, policy, cfg, success_collection_steps=15): def record_dataset(env, policy, cfg, success_collection_steps=0):
""" """
Record a dataset of robot interactions using either a policy or teleop. Record a dataset of robot interactions using either a policy or teleop.
@@ -2076,7 +2076,7 @@ def record_dataset(env, policy, cfg, success_collection_steps=15):
logging.info(f"Re-recording episode {episode_index}") logging.info(f"Re-recording episode {episode_index}")
continue continue
dataset.save_episode(cfg.task) dataset.save_episode()
episode_index += 1 episode_index += 1
# Finalize dataset # Finalize dataset
@@ -2144,7 +2144,7 @@ def main(cfg: EnvConfig):
env, env,
policy=policy, policy=policy,
cfg=cfg, cfg=cfg,
success_collection_steps=15, success_collection_steps=0,
) )
exit() exit()