Update lerobot/common/envs/utils.py

Co-authored-by: Remi <re.cadene@gmail.com>
This commit is contained in:
mshukor 2025-04-02 17:32:26 +02:00 committed by GitHub
parent b0511302cd
commit adb4af5d51
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 2 additions and 1 deletions

View File

@ -115,7 +115,8 @@ def check_env_attributes_and_types(env: gym.vector.VectorEnv) -> None:
) )
def infer_envs_task(env: gym.vector.VectorEnv, observation: dict[str, Any]) -> dict[str, Any]: def add_envs_task(env: gym.vector.VectorEnv, observation: dict[str, Any]) -> dict[str, Any]:
""" Adds task feature to the observation dict with respect to the first environment attribute. """
if hasattr(env.envs[0], "task_description"): if hasattr(env.envs[0], "task_description"):
observation["task"] = env.call("task_description") observation["task"] = env.call("task_description")
elif hasattr(env.envs[0], "task"): elif hasattr(env.envs[0], "task"):