Update lerobot/common/envs/utils.py
Co-authored-by: Remi <re.cadene@gmail.com>
This commit is contained in:
parent
b0511302cd
commit
adb4af5d51
|
@ -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"):
|
||||||
|
|
Loading…
Reference in New Issue