Better env implementation.

google · Dec 2, 2024 · 8cd1660 · 8cd1660
1 parent a5cf070
commit 8cd1660
Showing 1 changed file with 11 additions and 7 deletions.
diff --git a/brax/envs/fast.py b/brax/envs/fast.py
@@ -34,7 +34,10 @@ def __init__(self, **kwargs):
       raise ValueError('asymmetric_obs requires use_dict_obs=True')
 
   def _get_obs(self):
-    obs = {'state': jp.zeros(2)} if self._use_dict_obs else jp.zeros(2)
+    if not self._use_dict_obs:
+      return jp.zeros(2)
+
+    obs = {'state': jp.zeros(2)}
     if self._asymmetric_obs:
       obs['privileged_state'] = jp.zeros(4)
     return obs
@@ -78,12 +81,13 @@ def step_count(self):
 
   @property
   def observation_size(self):
-    if self._use_dict_obs:
-      size = {'state': 2}
-      if self._asymmetric_obs:
-        size['privileged_state'] = 4
-      return size
-    return 2
+    if not self._use_dict_obs:
+      return 2
+
+    obs = {'state': 2}
+    if self._asymmetric_obs:
+      obs['privileged_state'] = 4
+    return obs
 
   @property
   def action_size(self):