fix(gym): return node ids and action mask for np observations

Author: justindujardin

libraries/mathy_python/mathy/envs/gym/mathy_gym_env.py

@@ -69,9 +69,10 @@ def _observe(self, state: MathyEnvState) -> Union[MathyObservation, np.ndarray]:
         self.action_space.mask = action_mask
         if self.np_observation:
             # convert mask to probabilities
+            nodes = np.array(pad_array(observation.nodes, 512, 0))
             mask = np.array(pad_array(observation.mask, 512, 0))
             mask = mask / np.sum(mask)
-            return mask
+            return np.vstack((nodes, mask))
         return observation
 
     def reset(self):