feat: add reference info to step/reset return

stable_gym/envs/mujoco/ant_cost/ant_cost.py

@@ -283,8 +283,6 @@ def step(self, action):
         """
         obs, _, terminated, truncated, info = super().step(action)
 
-        self.state = obs
-
         cost, cost_info = self.cost(info["x_velocity"], -info["reward_ctrl"])
 
         # Add reference, x velocity and reference error to observation.
@@ -295,8 +293,23 @@ def step(self, action):
         if not self._exclude_x_velocity_from_observation:
             obs = np.append(obs, info["x_velocity"])
 
-        # Update info.
+        self.state = obs
+
+        # Update info dictionary.
+        del (
+            info["reward_forward"],
+            info["forward_reward"],
+            info["reward_ctrl"],
+            info["reward_survive"],
+        )
         info.update(cost_info)
+        info.update(
+            {
+                "reference": self.reference_forward_velocity,
+                "state_of_interest": info["x_velocity"],
+                "reference_error": info["x_velocity"] - self.reference_forward_velocity,
+            }
+        )
 
         return obs, cost, terminated, truncated, info
 
@@ -318,6 +331,8 @@ def reset(self, seed=None, options=None):
         """
         obs, info = super().reset(seed=seed, options=options)
 
+        _, cost_info = self.cost(0.0, 0.0)
+
         # Randomize the reference forward velocity if requested.
         if self._randomise_reference_forward_velocity:
             self.reference_forward_velocity = self.np_random.uniform(
@@ -334,6 +349,16 @@ def reset(self, seed=None, options=None):
 
         self.state = obs
 
+        # Update info dictionary.
+        info.update(cost_info)
+        info.update(
+            {
+                "reference": self.reference_forward_velocity,
+                "state_of_interest": 0.0,
+                "reference_error": 0.0 - self.reference_forward_velocity,
+            }
+        )
+
         return obs, info
 
     @property

stable_gym/envs/mujoco/half_cheetah_cost/half_cheetah_cost.py

@@ -231,8 +231,6 @@ def step(self, action):
         """
         obs, _, terminated, truncated, info = super().step(action)
 
-        self.state = obs
-
         cost, cost_info = self.cost(info["x_velocity"], -info["reward_ctrl"])
 
         # Add reference, x velocity and reference error to observation.
@@ -243,9 +241,18 @@ def step(self, action):
         if not self._exclude_x_velocity_from_observation:
             obs = np.append(obs, info["x_velocity"])
 
-        # Update info.
+        self.state = obs
+
+        # Update info dictionary.
         del info["reward_run"], info["reward_ctrl"]
         info.update(cost_info)
+        info.update(
+            {
+                "reference": self.reference_forward_velocity,
+                "state_of_interest": info["x_velocity"],
+                "reference_error": info["x_velocity"] - self.reference_forward_velocity,
+            }
+        )
 
         return obs, cost, terminated, truncated, info
 
@@ -267,6 +274,8 @@ def reset(self, seed=None, options=None):
         """
         obs, info = super().reset(seed=seed, options=options)
 
+        _, cost_info = self.cost(0.0, 0.0)
+
         # Randomize the reference forward velocity if requested.
         if self._randomise_reference_forward_velocity:
             self.reference_forward_velocity = self.np_random.uniform(
@@ -283,6 +292,16 @@ def reset(self, seed=None, options=None):
 
         self.state = obs
 
+        # Update info dictionary.
+        info.update(cost_info)
+        info.update(
+            {
+                "reference": self.reference_forward_velocity,
+                "state_of_interest": 0.0,
+                "reference_error": 0.0 - self.reference_forward_velocity,
+            }
+        )
+
         return obs, info
 
     @property

stable_gym/envs/mujoco/hopper_cost/hopper_cost.py

@@ -271,8 +271,6 @@ def step(self, action):
         """
         obs, _, terminated, truncated, info = super().step(action)
 
-        self.state = obs
-
         ctrl_cost = super().control_cost(action)
         cost, cost_info = self.cost(info["x_velocity"], ctrl_cost)
 
@@ -284,8 +282,17 @@ def step(self, action):
         if not self._exclude_x_velocity_from_observation:
             obs = np.append(obs, info["x_velocity"])
 
-        # Update info.
+        self.state = obs
+
+        # Update info dictionary.
         info.update(cost_info)
+        info.update(
+            {
+                "reference": self.reference_forward_velocity,
+                "state_of_interest": info["x_velocity"],
+                "reference_error": info["x_velocity"] - self.reference_forward_velocity,
+            }
+        )
 
         return obs, cost, terminated, truncated, info
 
@@ -307,6 +314,8 @@ def reset(self, seed=None, options=None):
         """
         obs, info = super().reset(seed=seed, options=options)
 
+        _, cost_info = self.cost(0.0, 0.0)
+
         # Randomize the reference forward velocity if requested.
         if self._randomise_reference_forward_velocity:
             self.reference_forward_velocity = self.np_random.uniform(
@@ -323,6 +332,16 @@ def reset(self, seed=None, options=None):
 
         self.state = obs
 
+        # Update info dictionary.
+        info.update(cost_info)
+        info.update(
+            {
+                "reference": self.reference_forward_velocity,
+                "state_of_interest": 0.0,
+                "reference_error": 0.0 - self.reference_forward_velocity,
+            }
+        )
+
         return obs, info
 
     @property

stable_gym/envs/mujoco/humanoid_cost/humanoid_cost.py

@@ -262,8 +262,6 @@ def step(self, action):
         """
         obs, _, terminated, truncated, info = super().step(action)
 
-        self.state = obs
-
         cost, cost_info = self.cost(info["x_velocity"], -info["reward_quadctrl"])
 
         # Add reference, x velocity and reference error to observation.
@@ -274,14 +272,23 @@ def step(self, action):
         if not self._exclude_x_velocity_from_observation:
             obs = np.append(obs, info["x_velocity"])
 
-        # Update info.
+        self.state = obs
+
+        # Update info dictionary.
         del (
             info["reward_linvel"],
             info["reward_quadctrl"],
             info["reward_alive"],
             info["forward_reward"],
         )
         info.update(cost_info)
+        info.update(
+            {
+                "reference": self.reference_forward_velocity,
+                "state_of_interest": info["x_velocity"],
+                "reference_error": info["x_velocity"] - self.reference_forward_velocity,
+            }
+        )
 
         return obs, cost, terminated, truncated, info
 
@@ -303,6 +310,8 @@ def reset(self, seed=None, options=None):
         """
         obs, info = super().reset(seed=seed, options=options)
 
+        _, cost_info = self.cost(0.0, 0.0)
+
         # Randomize the reference forward velocity if requested.
         if self._randomise_reference_forward_velocity:
             self.reference_forward_velocity = self.np_random.uniform(
@@ -319,6 +328,16 @@ def reset(self, seed=None, options=None):
 
         self.state = obs
 
+        # Update info dictionary.
+        info.update(cost_info)
+        info.update(
+            {
+                "reference": self.reference_forward_velocity,
+                "state_of_interest": 0.0,
+                "reference_error": 0.0 - self.reference_forward_velocity,
+            }
+        )
+
         return obs, info
 
     @property

stable_gym/envs/mujoco/swimmer_cost/swimmer_cost.py

@@ -229,8 +229,6 @@ def step(self, action):
         """
         obs, _, terminated, truncated, info = super().step(action)
 
-        self.state = obs
-
         cost, cost_info = self.cost(info["x_velocity"], -info["reward_ctrl"])
 
         # Add reference, x velocity and reference error to observation.
@@ -241,9 +239,18 @@ def step(self, action):
         if not self._exclude_x_velocity_from_observation:
             obs = np.append(obs, info["x_velocity"])
 
-        # Update info.
+        self.state = obs
+
+        # Update info dictionary.
         del info["reward_fwd"], info["reward_ctrl"], info["forward_reward"]
         info.update(cost_info)
+        info.update(
+            {
+                "reference": self.reference_forward_velocity,
+                "state_of_interest": info["x_velocity"],
+                "reference_error": info["x_velocity"] - self.reference_forward_velocity,
+            }
+        )
 
         return obs, cost, terminated, truncated, info
 
@@ -265,6 +272,8 @@ def reset(self, seed=None, options=None):
         """
         obs, info = super().reset(seed=seed, options=options)
 
+        _, cost_info = self.cost(0.0, 0.0)
+
         # Randomize the reference forward velocity if requested.
         if self._randomise_reference_forward_velocity:
             self.reference_forward_velocity = self.np_random.uniform(
@@ -281,6 +290,16 @@ def reset(self, seed=None, options=None):
 
         self.state = obs
 
+        # Update info dictionary.
+        info.update(cost_info)
+        info.update(
+            {
+                "reference": self.reference_forward_velocity,
+                "state_of_interest": 0.0,
+                "reference_error": 0.0 - self.reference_forward_velocity,
+            }
+        )
+
         return obs, info
 
     @property

stable_gym/envs/mujoco/walker2d_cost/walker2d_cost.py

@@ -266,8 +266,6 @@ def step(self, action):
         """
         obs, _, terminated, truncated, info = super().step(action)
 
-        self.state = obs
-
         ctrl_cost = super().control_cost(action)
         cost, cost_info = self.cost(info["x_velocity"], ctrl_cost)
 
@@ -279,8 +277,17 @@ def step(self, action):
         if not self._exclude_x_velocity_from_observation:
             obs = np.append(obs, info["x_velocity"])
 
-        # Update info.
+        self.state = obs
+
+        # Update info dictionary.
         info.update(cost_info)
+        info.update(
+            {
+                "reference": self.reference_forward_velocity,
+                "state_of_interest": info["x_velocity"],
+                "reference_error": info["x_velocity"] - self.reference_forward_velocity,
+            }
+        )
 
         return obs, cost, terminated, truncated, info
 
@@ -302,6 +309,8 @@ def reset(self, seed=None, options=None):
         """
         obs, info = super().reset(seed=seed, options=options)
 
+        _, cost_info = self.cost(0.0, 0.0)
+
         # Randomize the reference forward velocity if requested.
         if self._randomise_reference_forward_velocity:
             self.reference_forward_velocity = self.np_random.uniform(
@@ -318,6 +327,16 @@ def reset(self, seed=None, options=None):
 
         self.state = obs
 
+        # Update info dictionary.
+        info.update(cost_info)
+        info.update(
+            {
+                "reference": self.reference_forward_velocity,
+                "state_of_interest": 0.0,
+                "reference_error": 0.0 - self.reference_forward_velocity,
+            }
+        )
+
         return obs, info
 
     @property