Fix BabayAI data generation (#118)

* Fix wrong observation used in BabyAI's datasets

* Remove tasks with issues from slurm generation of BabyAI's dataset

* Some cleaning in BabyAI's dataset generation

* Handle truncated episodes for BabyAI

---------

Co-authored-by: cromac <[email protected]>

data/envs/babyai/create_babyai_dataset.py

@@ -5,12 +5,10 @@
 import gymnasium as gym
 import numpy as np
 from bot_agent import Bot
-from utils.env_wrappers import AddRGBImgPartialObsWrapper
 
 
 def create_babyai_dataset(name_env, saving_path, max_num_episodes=100000, test_set_percentage=5):
     env = gym.make(name_env)
-    env = AddRGBImgPartialObsWrapper(env)  # add rgb image to obs
 
     class BabyAIEnvDataset:
         def __init__(self):
@@ -54,8 +52,10 @@ def reset_env_and_policy(env):
         while not done:
             try:
                 action = policy.replan()
-                obs, r, done, _, infos = env.step(action)
+                _obs, r, term, trunc, infos = env.step(action)
                 dataset.add_step(obs, int(action), r)
+                obs = _obs
+                done = term or trunc
             except Exception:
                 done = True
 

data/envs/babyai/create_babyai_dataset.slurm

@@ -14,11 +14,11 @@
 #SBATCH --array=0-37
 
 ENVS=(
-    BabyAI-GoTo-v0 BabyAI-GoToDoor-v0 BabyAI-GoToImpUnlock-v0 BabyAI-GoToLocal-v0 BabyAI-GoToObj-v0 BabyAI-GoToObjDoor-v0 BabyAI-GoToRedBall-v0 BabyAI-GoToRedBallGrey-v0 BabyAI-GoToRedBallNoDists-v0 BabyAI-GoToRedBlueBall-v0 BabyAI-GoToSeq-v0 BabyAI-Open-v0 BabyAI-OpenDoor-v0 BabyAI-OpenDoorsOrderN4-v0 BabyAI-OpenRedDoor-v0 BabyAI-OpenTwoDoors-v0 BabyAI-ActionObjDoor-v0 BabyAI-FindObjS5-v0 BabyAI-KeyCorridor-v0 BabyAI-MoveTwoAcrossS8N9-v0 BabyAI-OneRoomS8-v0 BabyAI-Pickup-v0 BabyAI-PickupAbove-v0 BabyAI-PickupDist-v0 BabyAI-PickupLoc-v0 BabyAI-UnblockPickup-v0 BabyAI-BossLevel-v0 BabyAI-BossLevelNoUnlock-v0 BabyAI-MiniBossLevel-v0 BabyAI-Synth-v0 BabyAI-SynthLoc-v0 BabyAI-SynthSeq-v0 BabyAI-BlockedUnlockPickup-v0 BabyAI-KeyInBox-v0 BabyAI-Unlock-v0 BabyAI-UnlockLocal-v0 BabyAI-UnlockPickup-v0 BabyAI-UnlockToUnlock-v0
+    BabyAI-GoTo-v0 BabyAI-GoToDoor-v0 BabyAI-GoToLocal-v0 BabyAI-GoToObj-v0 BabyAI-GoToObjDoor-v0 BabyAI-GoToRedBall-v0 BabyAI-GoToRedBallGrey-v0 BabyAI-GoToRedBallNoDists-v0 BabyAI-GoToRedBlueBall-v0 BabyAI-GoToSeq-v0 BabyAI-Open-v0 BabyAI-OpenDoor-v0 BabyAI-OpenDoorsOrderN4-v0 BabyAI-OpenRedDoor-v0 BabyAI-OpenTwoDoors-v0 BabyAI-ActionObjDoor-v0 BabyAI-FindObjS5-v0 BabyAI-KeyCorridor-v0 BabyAI-MoveTwoAcrossS8N9-v0 BabyAI-OneRoomS8-v0 BabyAI-Pickup-v0 BabyAI-PickupAbove-v0 BabyAI-PickupDist-v0 BabyAI-PickupLoc-v0 BabyAI-UnblockPickup-v0 BabyAI-BossLevel-v0 BabyAI-BossLevelNoUnlock-v0 BabyAI-MiniBossLevel-v0 BabyAI-Synth-v0 BabyAI-SynthLoc-v0 BabyAI-SynthSeq-v0 BabyAI-BlockedUnlockPickup-v0 BabyAI-UnlockLocal-v0 BabyAI-UnlockPickup-v0
 )
 
 NAMES=(
-    go-to go-to-door go-to-imp-unlock go-to-local go-to-obj go-to-obj-door go-to-red-ball go-to-red-ball-grey go-to-red-ball-no-dists go-to-red-blue-ball go-to-seq open open-door open-doors-order-n4 open-red-door open-two-doors action-obj-door find-obj-s5 key-corridor move-two-across-s8n9 one-room-s8 pickup pickup-above pickup-dist pickup-loc unblock-pickup boss-level boss-level-no-unlock mini-boss-level synth synth-loc synth-seq blocked-unlock-pickup key-in-box unlock unlock-local unlock-pickup unlock-to-unlock
+    go-to go-to-door go-to-local go-to-obj go-to-obj-door go-to-red-ball go-to-red-ball-grey go-to-red-ball-no-dists go-to-red-blue-ball go-to-seq open open-door open-doors-order-n4 open-red-door open-two-doors action-obj-door find-obj-s5 key-corridor move-two-across-s8n9 one-room-s8 pickup pickup-above pickup-dist pickup-loc unblock-pickup boss-level boss-level-no-unlock mini-boss-level synth synth-loc synth-seq blocked-unlock-pickup unlock-local unlock-pickup
 )
 
 module purge