Fixed wrappers.vector.RecordEpisodeStatistics episode length computation from new autoreset api

[TimSchneider42]

Description

Fixed wrappers.vector.RecordEpisodeStatistics episode length computation.

Fixes #1017

Type of change

Please delete options that are not relevant.

• Documentation only change (no code changed)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• This change requires a documentation update

Checklist:

• I have run the pre-commit checks with pre-commit run --all-files (see CONTRIBUTING.md instructions to set it up)
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[pseudo-rnd-thoughts]

I believe this isn't possible due to the episode time taken will vary between the two

This is the reason for the sequent code

A comment would be helpful

[TimSchneider42]

Strangely, this test did not fail for me. I will move it back to where it was originally. Not sure why I moved it in the first place.

[pseudo-rnd-thoughts]

Once we pop t, why can we not just do data_equivalence for the two infos?
Why did this test fail previously?

[pseudo-rnd-thoughts]

I messed up on the testing of this, it should be assert data_equivalence(...)

The new testing code should be

        assert data_equivalence(wrapper_vector_obs, vector_wrapper_obs)
        assert data_equivalence(wrapper_vector_reward, vector_wrapper_reward)
        assert data_equivalence(wrapper_vector_terminated, vector_wrapper_terminated)
        assert data_equivalence(wrapper_vector_truncated, vector_wrapper_truncated)

        if "episode" in wrapper_vector_info:
            assert "episode" in vector_wrapper_info

            wrapper_vector_time_taken = wrapper_vector_info["episode"].pop("t")
            vector_wrapper_time_taken = vector_wrapper_info["episode"].pop("t")

            assert wrapper_vector_time_taken.shape == vector_wrapper_time_taken.shape
            assert wrapper_vector_time_taken.dtype == vector_wrapper_time_taken.dtype

            vector_wrapper_info["episode"].pop("_l")
            vector_wrapper_info["episode"].pop("_r")
            vector_wrapper_info["episode"].pop("_t")

        assert data_equivalence(wrapper_vector_info, vector_wrapper_info)

This should now test it correctly.

I added max_episode_steps to the make_vec such that episode ends for this to actually be tested (update to main, this was previously broken)

[TimSchneider42]

Oh, that is why it did not fail when I moved the info dict comparison! Do you want me to fix the test, or are you going to do it?

[pseudo-rnd-thoughts]

Could you update the test to my suggested code above, I think it is simplier than your code and achieves the same thing.

[TimSchneider42]

I totally agree that your version is simpler. For some reason, I thought the test did not consider the episode lengths at all, but after looking at it more carefully, I now understand what it actually does. I just changed the code accordingly.

[pseudo-rnd-thoughts]

Yeah, it does that implicitly with the info check at the end

[TimSchneider42]

I moved the info dict comparison back to its original place

[pseudo-rnd-thoughts]

An alternative is self.episode_returns = np.where(self.prev_dones, self.episode_returns + rewards, 0)
However as this recreates the array each time I don't think this is a good idea

[pseudo-rnd-thoughts]

Yeah, it does that implicitly with the info check at the end