From ff89addd5f556acb8054d3022d4e4c51bfadaf8d Mon Sep 17 00:00:00 2001 From: zhaojing Date: Thu, 19 Dec 2024 00:12:38 -0800 Subject: [PATCH] add timeout in env infos --- examples/locomotion/go2_env.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/examples/locomotion/go2_env.py b/examples/locomotion/go2_env.py index a23f7ed4..4368b8ab 100644 --- a/examples/locomotion/go2_env.py +++ b/examples/locomotion/go2_env.py @@ -150,6 +150,11 @@ def step(self, actions): self.reset_buf = self.episode_length_buf > self.max_episode_length self.reset_buf |= torch.abs(self.base_euler[:, 1]) > self.env_cfg["termination_if_pitch_greater_than"] self.reset_buf |= torch.abs(self.base_euler[:, 0]) > self.env_cfg["termination_if_roll_greater_than"] + + time_out_idx = (self.episode_length_buf > self.max_episode_length).nonzero(as_tuple=False).flatten() + self.extras["time_outs"] = torch.zeros_like(self.reset_buf, device=self.device, dtype=gs.tc_float) + self.extras["time_outs"][time_out_idx] = 1.0 + self.reset_idx(self.reset_buf.nonzero(as_tuple=False).flatten()) # compute reward