[Question] Why is episode_length_buf not initialized for use in observation functions (for a manager-based environment), but can be used to determine time outs (a termination function)?

[TheIndoorDad]

Hello all,

I was writing an observation function to pass the time remaining in an episode (in seconds) to an observation term in a Manager-based environment. To do so, I figured I could use episode_length_buf from the ManagerBasedRLEnv class, like is done for time out terminations (source/isaaclab/isaaclab/envs/mdp/terminations.py line 32). My function looks like:

import torch
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from isaaclab.envs import ManagerBasedRLEnv

def time_remaining_s(env: ManagerBasedRLEnv) -> torch.Tensor:
    """Compute the remaining time in the episode."""
    current_time_s = torch.reshape(env.episode_length_buf, [env.num_envs,-1]) * env.step_dt
    return env.max_episode_length_s - current_time_s

However, when attempting to train a policy calling this observation function (using rsl_rl) I would get an error when the Observation Manager was initialized:

[INFO] Action Manager:  <ActionManager> contains 1 active terms.
+------------------------------------+
|  Active Action Terms (shape: 14)   |
+--------+-------------+-------------+
| Index  | Name        |   Dimension |
+--------+-------------+-------------+
|   0    | joint_pos   |          14 |
+--------+-------------+-------------+

Error executing job with overrides: []
Traceback (most recent call last):
  File "/home/trans/ContinuO_IsaacLab/IsaacLab-ContinuO/source/isaaclab_tasks/isaaclab_tasks/utils/hydra.py", line 101, in hydra_main
    func(env_cfg, agent_cfg, *args, **kwargs)
  File "/home/trans/ContinuO_IsaacLab/ContinuO_IsaacLab_Ext/scripts/rsl_rl/train.py", line 105, in main
    env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None)
  File "/home/trans/isaacsim/exts/omni.isaac.ml_archive/pip_prebundle/gymnasium/envs/registration.py", line 802, in make
    env = env_creator(**env_spec_kwargs)
  File "/home/trans/ContinuO_IsaacLab/IsaacLab-ContinuO/source/isaaclab/isaaclab/envs/manager_based_rl_env.py", line 79, in __init__
    super().__init__(cfg=cfg)
  File "/home/trans/ContinuO_IsaacLab/IsaacLab-ContinuO/source/isaaclab/isaaclab/envs/manager_based_env.py", line 152, in __init__
    self.load_managers()
  File "/home/trans/ContinuO_IsaacLab/IsaacLab-ContinuO/source/isaaclab/isaaclab/envs/manager_based_rl_env.py", line 117, in load_managers
    super().load_managers()
  File "/home/trans/ContinuO_IsaacLab/IsaacLab-ContinuO/source/isaaclab/isaaclab/envs/manager_based_env.py", line 237, in load_managers
    self.observation_manager = ObservationManager(self.cfg.observations, self)
  File "/home/trans/ContinuO_IsaacLab/IsaacLab-ContinuO/source/isaaclab/isaaclab/managers/observation_manager.py", line 82, in __init__
    super().__init__(cfg, env)
  File "/home/trans/ContinuO_IsaacLab/IsaacLab-ContinuO/source/isaaclab/isaaclab/managers/manager_base.py", line 131, in __init__
    self._prepare_terms()
  File "/home/trans/ContinuO_IsaacLab/IsaacLab-ContinuO/source/isaaclab/isaaclab/managers/observation_manager.py", line 411, in _prepare_terms
    obs_dims = tuple(term_cfg.func(self._env, **term_cfg.params).shape)
  File "/home/trans/ContinuO_IsaacLab/ContinuO_IsaacLab_Ext/exts/ContinuO_IsaacLab_Ext/ContinuO_IsaacLab_Ext/tasks/locomotion/position_nav/mdp/observations.py", line 18, in time_remaining_s
    current_time_s = torch.reshape(env.episode_length_buf, [env.num_envs,-1]) * env.step_dt
AttributeError: 'ManagerBasedRLEnv' object has no attribute 'episode_length_buf'

I think this was happening because the Observation Manager is loaded in the course of the initialization of the ManagerBasedEnv class, before episode_length_buf is initialized in ManagerBasedRLEnv, but I can't figure out why it would work for the time_out termination function in that case.

I managed to get it working by initializing episode_length_buf before load_managers() is called in ManagerBasedEnv (i.e., by inserting self.episode_length_buf = torch.zeros(self.num_envs, device=self.device, dtype=torch.long) at line 162). I'd rather not have to do that, however, as otherwise my project is entirely contained in an external extension. I also suppose it would probably be easier in a direct workflow, but I'd also rather not switch just for this.

Does anyone know why episode_length_buf works for the time_out termination function (source/isaaclab/isaaclab/envs/mdp/terminations.py line 32), but not for my observation function without modifying ManagerBasedEnv?

Thank you in advance for your assistance, and for your patience with anything I'm likely to have misunderstood - I'm relatively new to the field and my background is not in software dev.

[RandomOakForest]

Thank you for posting this, great question. The team is aware of this, would you like to submit a PR that would add this functionality? Thanks!

[TheIndoorDad]

Sure I can do that. I can likely get it started early this week.

[TheIndoorDad]

#2332