add reset method to vecenv wrapper and fix test

Author: ClemensSchwarke

source/isaaclab_rl/isaaclab_rl/rsl_rl/vecenv_wrapper.py

@@ -136,6 +136,11 @@ def episode_length_buf(self, value: torch.Tensor):
     def seed(self, seed: int = -1) -> int:  # noqa: D102
         return self.unwrapped.seed(seed)
 
+    def reset(self) -> tuple[TensorDict, dict]:  # noqa: D102
+        # reset the environment
+        obs_dict, extras = self.env.reset()
+        return TensorDict(obs_dict, batch_size=[self.num_envs]), extras
+
     def get_observations(self) -> TensorDict:
         """Returns the current observations of the environment."""
         if hasattr(self.unwrapped, "observation_manager"):

source/isaaclab_rl/test/test_rsl_rl_wrapper.py

@@ -16,6 +16,7 @@
 
 import gymnasium as gym
 import torch
+from tensordict import TensorDict
 
 import carb
 import omni.usd
@@ -161,6 +162,8 @@ def _check_valid_tensor(data: torch.Tensor | dict) -> bool:
     """
     if isinstance(data, torch.Tensor):
         return not torch.any(torch.isnan(data))
+    elif isinstance(data, TensorDict):
+        return not data.isnan().any()
     elif isinstance(data, dict):
         valid_tensor = True
         for value in data.values():