diff --git a/babyai/levels/bonus_levels.py b/babyai/levels/bonus_levels.py index f46ae815..6d2e410b 100644 --- a/babyai/levels/bonus_levels.py +++ b/babyai/levels/bonus_levels.py @@ -7,7 +7,7 @@ class Level_GoToRedBlueBall(RoomGridLevel): """ Go to the red ball or to the blue ball. - Simple level with language and exactly two different objects. + There is only one object, and therefore language is not required. """ def __init__(self, room_size=8, seed=None): @@ -20,10 +20,8 @@ def __init__(self, room_size=8, seed=None): def gen_mission(self): self.place_agent() - red_ball, _ = self.add_object(0, 0, 'ball', 'red') - blue_ball, _ = self.add_object(0, 0, 'ball', 'blue') - - obj = self._rand_elem([red_ball, blue_ball]) + color = self._rand_elem(['red', 'blue']) + obj, _ = self.add_object(0, 0, 'ball', color) self.instrs = GoToInstr(ObjDesc(obj.type, obj.color))