[MLA-1809] catch mismatched observation sizes

com.unity.ml-agents/CHANGELOG.md

@@ -20,6 +20,7 @@ and this project adheres to
 #### com.unity.ml-agents (C#)
 #### ml-agents / ml-agents-envs / gym-unity (Python)
 - An issue that caused `GAIL` to fail for environments where agents can terminate episodes by self-sacrifice has been fixed. (#4971)
+- Made the error message when observations of different shapes are sent to the trainer clearer. (#5030)
 
 ## [1.8.0-preview] - 2021-02-17
 ### Major Changes

ml-agents-envs/mlagents_envs/rpc_utils.py

@@ -285,10 +285,22 @@ def steps_from_proto(
     ]
     decision_obs_list: List[np.ndarray] = []
     terminal_obs_list: List[np.ndarray] = []
-    for obs_index, observation_specs in enumerate(behavior_spec.observation_specs):
-        is_visual = len(observation_specs.shape) == 3
+    for obs_index, observation_spec in enumerate(behavior_spec.observation_specs):
+        # Check that all the observations match the expected size.
+        # This gives a nicer error than a cryptic numpy error later.
+        expected_obs_shape = tuple(observation_spec.shape)
+        for agent_info in agent_info_list:
+            agent_obs_shape = tuple(agent_info.observations[obs_index].shape)
+            if expected_obs_shape != agent_obs_shape:
+                raise ValueError(
+                    f"Observation at index={obs_index} for agent with "
+                    f"id={agent_info.id} didn't match the ObservationSpec. "
+                    f"Expected shape {expected_obs_shape} but got {agent_obs_shape}."
+                )
+
+        is_visual = len(observation_spec.shape) == 3
         if is_visual:
-            obs_shape = cast(Tuple[int, int, int], observation_specs.shape)
+            obs_shape = cast(Tuple[int, int, int], observation_spec.shape)
             decision_obs_list.append(
                 _process_maybe_compressed_observation(
                     obs_index, obs_shape, decision_agent_info_list
@@ -302,12 +314,12 @@ def steps_from_proto(
         else:
             decision_obs_list.append(
                 _process_rank_one_or_two_observation(
-                    obs_index, observation_specs.shape, decision_agent_info_list
+                    obs_index, observation_spec.shape, decision_agent_info_list
                 )
             )
             terminal_obs_list.append(
                 _process_rank_one_or_two_observation(
-                    obs_index, observation_specs.shape, terminal_agent_info_list
+                    obs_index, observation_spec.shape, terminal_agent_info_list
                 )
             )
     decision_rewards = np.array(

ml-agents-envs/mlagents_envs/tests/test_rpc_utils.py

@@ -369,6 +369,20 @@ def test_batched_step_result_from_proto():
     assert terminal_steps.obs[1].shape[1] == shapes[1][0]
 
 
+def test_mismatch_observations_raise_in_step_result_from_proto():
+    n_agents = 10
+    shapes = [(3,), (4,)]
+    spec = BehaviorSpec(
+        create_observation_specs_with_shapes(shapes), ActionSpec.create_continuous(3)
+    )
+    ap_list = generate_list_agent_proto(n_agents, shapes)
+    # Hack an observation to be larger, we should get an exception
+    ap_list[0].observations[0].shape[0] += 1
+    ap_list[0].observations[0].float_data.data.append(0.42)
+    with pytest.raises(ValueError):
+        steps_from_proto(ap_list, spec)
+
+
 def test_action_masking_discrete():
     n_agents = 10
     shapes = [(3,), (4,)]