Does GAE skip NaN values for auto-resetting environments?

[ErcBunny]

I've created an auto-resetting transformed environment where observation at done steps are marked as NaN. Then I tried to use the bundled GAE with the rollout data of my env, I found NaNs in the output entries of the GAE. Here is my transformed env.

env = AutoResetEnv(
    original_env,
    Compose(
        AutoResetTransform(fill_bool="nan", fill_int="nan", fill_float="nan"),
        InitTracker(),
        lstm.make_tensordict_primer(),
    ),
)

And GAE is used like this: (btw am I using the right selection of data?)

for i, tensordict_data in enumerate(self.data_collector):
    for epoch in range(self.num_epochs):

        gae_io = tensordict_data.select(
            "obs",
            ("next", "obs"),
            ("next", "done"),
            ("next", "reward"),
        ).to(device="cuda")
        self.gae(gae_io)

        tensordict_data.update(
            gae_io.select(
                "advantage",
                "state_value",
                "value_target",
                ("next", "state_value"),
            ).to(device=tensordict_data.device)
        )

Then by printing "advantage", "state_value", and "value_target", I found NaNs in the tensor. So my question is: does the GAE implementation consider autoresetting envs? If yes, then there shouldn't be NaNs or I am using it wrongly. Could you please help me diagnose the problem?

[vmoens]

We don't look at NaN afair.
That would induce some overhead. But GAE usually works fine with autoresetting envs and contiguous trajectories, can you elaborate on why the NaN are needed?

[ErcBunny]

Thanks for pointing out that the GAE works fine with autoresetting envs. I asked this question because I found network weights turned to NaN after an optimization step. I might have made a mistake somewhere else in my code causing the NaNs.