Bc fix

brax/envs/fast.py

@@ -31,12 +31,14 @@ class ObservationMode(enum.Enum):
     DICT_STATE: Dictionary of state info.
     DICT_PIXELS: Dictionary of pixel observations.
     DICT_PIXELS_STATE: Dictionary of both state and pixel info.
+    DICT_LATENT_STATE: Dictionary with latent pixel and state info.
   """
 
   NDARRAY = 'ndarray'
   DICT_STATE = 'dict_state'
   DICT_PIXELS = 'dict_pixels'
   DICT_PIXELS_STATE = 'dict_pixels_state'
+  DICT_LATENT_STATE = 'dict_latent_state'
 
 
 class Fast(PipelineEnv):
@@ -73,11 +75,17 @@ def reset(self, rng: jax.Array) -> State:
         'pixels/view_0': jp.zeros((4, 4, 3)),
         'pixels/view_1': jp.zeros((4, 4, 3)),
     }
+    latent_pixels = {
+        'latent_pixels/view_0': jp.zeros(12),
+        'latent_pixels/view_1': jp.zeros(12),
+    }
 
     if self._obs_mode == ObservationMode.DICT_PIXELS:
       obs = pixels
     elif self._obs_mode == ObservationMode.DICT_PIXELS_STATE:
       obs = {**obs, **pixels}
+    elif self._obs_mode == ObservationMode.DICT_LATENT_STATE:
+      obs = {**obs, **latent_pixels}
     elif self._obs_mode == ObservationMode.NDARRAY:
       obs = obs['state']
 
@@ -101,11 +109,17 @@ def step(self, state: State, action: jax.Array) -> State:
         'pixels/view_0': jp.zeros((4, 4, 3)),
         'pixels/view_1': jp.zeros((4, 4, 3)),
     }
+    latent_pixels = {
+        'latent_pixels/view_0': jp.zeros(12),
+        'latent_pixels/view_1': jp.zeros(12),
+    }
 
     if self._obs_mode == ObservationMode.DICT_PIXELS:
       obs = pixels
     elif self._obs_mode == ObservationMode.DICT_PIXELS_STATE:
       obs = {**obs, **pixels}
+    elif self._obs_mode == ObservationMode.DICT_LATENT_STATE:
+      obs = {**obs, **latent_pixels}
     elif self._obs_mode == ObservationMode.NDARRAY:
       obs = obs['state']
 

brax/training/agents/bc/networks.py

@@ -78,18 +78,21 @@ def make_bc_networks(
     policy_hidden_layer_sizes: Sequence[int] = (32,) * 4,
     activation: networks.ActivationFn = linen.swish,
     policy_obs_key: str = '',
-    vision: bool = False
+    vision: bool = False,
+    latent_vision: bool = False,
 ) -> BCNetworks:
   """Make BC networks with preprocessor.
   Note that vision = True assumes Frozen Encoder."""
   parametric_action_distribution = distribution.NormalTanhDistribution(
       event_size=action_size
   )
-  make_policy_network = (
-      networks.make_policy_network_latents
-      if vision
-      else networks.make_policy_network
-  )
+
+  if latent_vision:
+    make_policy_network = networks.make_policy_network_latents
+  elif vision:
+    make_policy_network = networks.make_policy_network_vision
+  else:
+    make_policy_network = networks.make_policy_network
 
   policy_network = make_policy_network(
       parametric_action_distribution.param_size,

brax/training/agents/bc/train_test.py

@@ -127,13 +127,13 @@ def testNetworkEncoding(self):
   def testPixelsBC(self):
     """Test BC with pixel observations."""
     teacher_policy = self._create_teacher_policy(num_envs=50)
-    fast = envs.get_environment('fast', obs_mode='dict_pixels_state')
+    fast = envs.get_environment('fast', obs_mode='dict_latent_state')
     fast = envs.training.wrap(fast, episode_length=128, action_repeat=1)
     network_factory = functools.partial(
         bc_networks.make_bc_networks,
         policy_hidden_layer_sizes=(32,),
         policy_obs_key='state',
-        vision=True,
+        latent_vision=True,
     )
     _, _, metrics = bc.train(
         env=fast,