sampling blocks and bowls with matching colors

app.py

@@ -73,8 +73,9 @@ class DemoRunner:
         self._model_name = model_name
 
         self._env = PickPlaceEnv(render=True, high_res=True, high_frame_rate=False)
-        block_list = np.random.choice(ALL_BLOCKS, size=n_blocks, replace=False).tolist()
-        bowl_list = np.random.choice(ALL_BOWLS, size=n_bowls, replace=False).tolist()
+        list_idxs = np.random.choice(len(ALL_BLOCKS), size=max(n_blocks, n_bowls), replace=False)
+        block_list = [ALL_BLOCKS[i] for i in list_idxs[:n_blocks]]
+        bowl_list = [ALL_BOWLS[i] for i in list_idxs[:n_bowls]]
         obj_list = block_list + bowl_list
         self._env.reset(obj_list)