Initial Commit

.gitattributes

@@ -0,0 +1,40 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text
+*.jpeg filter=lfs diff=lfs merge=lfs -text
+*.mp3 filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,38 @@
+*.idea/
+*.vscode/
+
+*__pycache__/
+*.pyc
+*.pyo
+*.pyd
+
+*.pt
+*.pth
+*.tar.gz
+*.zip
+
+!**/train/
+**/train/*
+!**/train/saved
+
+!**/inference/
+**/inference/*
+!**/inference/saved
+
+!**/maps/
+**/maps/*
+!**/maps/example
+!**/maps/gpt4o
+!**/maps/lisa
+
+# For taxabind_avs
+**/dataset/
+**/checkpoints/
+
+!**/lightning_logs/
+**/lightning_logs/*
+!**/lightning_logs/saved
+
+# Saved weights & logs
+**avs_rl_policy.pth
+**/avs_rl_policy_21.5k/*

.vscode/launch.json

@@ -0,0 +1,15 @@
+{
+    "version": "0.2.0",
+    "configurations": [
+        {
+            "name": "Debug app.py",
+            "type": "debugpy",
+            "request": "launch",
+            "program": "${workspaceFolder}/app.py",
+            "cwd": "${workspaceFolder}",
+            "console": "integratedTerminal",
+            "justMyCode": false,
+            "python": "/home/user/anaconda3/envs/vlm-search/bin/python3"
+        }
+    ]
+}

README.md

@@ -0,0 +1,13 @@
+---
+title: Search-TTA
+emoji: 🦁
+colorFrom: green
+colorTo: gray
+sdk: gradio
+sdk_version: 5.31.0
+app_file: app.py
+pinned: false
+short_description: Multimodal Test-time Adaptation Framework for Visual Search
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,528 @@
+"""
+Simplified Gradio demo for Search-TTA evaluation.
+"""
+
+# ────────────────────────── imports ───────────────────────────────────
+from pathlib import Path
+import matplotlib
+matplotlib.use("Agg", force=True)
+
+import gradio as gr
+import ctypes  # for safely stopping background threads
+import os, glob, threading, time
+import torch
+from PIL import Image
+import json
+import shutil
+import spaces   # integration with ZeroGPU on hf
+from planner.test_parameter import *         
+from planner.model import PolicyNet           
+from planner.test_worker import TestWorker 
+from taxabind_avs.satbind.clip_seg_tta import ClipSegTTA  
+
+
+# Helper to kill a Python thread by injecting SystemExit
+def _stop_thread(thread: threading.Thread):
+    """Forcefully raise SystemExit in the given thread (best-effort)."""
+    if thread is None or not thread.is_alive():
+        return
+    tid = thread.ident
+    if tid is None:
+        return
+    # Ask CPython to raise SystemExit in the thread context
+    res = ctypes.pythonapi.PyThreadState_SetAsyncExc(ctypes.c_long(tid), ctypes.py_object(SystemExit))
+    if res > 1:
+        # If it returned >1, cleanup and fail safe
+        ctypes.pythonapi.PyThreadState_SetAsyncExc(ctypes.c_long(tid), None)
+
+# ──────────── Thread Registry for Cleanup on Tab Switch ─────────────
+_running_threads: list[threading.Thread] = []
+_running_threads_lock = threading.Lock()
+
+# Map worker threads to their ClipSegTTA instance so UI can read executing_tta flag
+_thread_clip_map: dict[threading.Thread, ClipSegTTA] = {}
+
+# ──────────── Run directory rotation ─────────────
+RUN_HISTORY_LIMIT = 30  # keep at most this many timestamped run directories per instance
+
+def _prune_old_run_dirs(base_dir: str, limit: int = RUN_HISTORY_LIMIT):
+    """Delete oldest timestamp-named run directories leaving only *limit* of the newest ones."""
+    try:
+        dirs = [d for d in os.listdir(base_dir) if os.path.isdir(os.path.join(base_dir, d))]
+        dirs.sort()
+        if len(dirs) > limit:
+            for obsolete in dirs[:-limit]:
+                shutil.rmtree(os.path.join(base_dir, obsolete), ignore_errors=True)
+    except Exception:
+        pass
+
+
+# CHANGE ME!
+POLL_INTERVAL = 1.0 # For visualization
+
+# Prepare the model
+device = torch.device('cuda') if USE_GPU and torch.cuda.is_available() else torch.device('cpu')
+policy_net = PolicyNet(INPUT_DIM, EMBEDDING_DIM).to(device)
+script_dir = Path(__file__).resolve().parent
+print("real_script_dir: ", script_dir)
+checkpoint = torch.load(f'{MODEL_PATH}/{MODEL_NAME}')
+policy_net.load_state_dict(checkpoint['policy_model'])
+print('Model loaded!')    
+
+# Load metadata json
+tgts_metadata_json_path = os.path.join(script_dir, "examples/metadata.json")
+tgts_metadata = json.load(open(tgts_metadata_json_path))
+
+
+# ────────────────────────── Gradio process fn ─────────────────────────
+
+### integration with ZeroGPU on hf
+# @spaces.GPU
+def process_search_tta(
+    sat_path: str | None,
+    ground_path: str | None,
+    taxonomy: str | None = None,
+    session_threads: list[threading.Thread] | None = None,
+):
+    """Run both TTA and non-TTA search episodes concurrently and stream both heat-maps."""
+
+    if session_threads is None:
+        session_threads = []
+
+    # Disable Run button and clear image/status outputs, hide sliders, clear frame states
+    yield (
+        gr.update(interactive=False),
+        gr.update(value=None),
+        gr.update(value=None),
+        gr.update(value="Initializing model…", visible=True),
+        gr.update(value="Initializing model…", visible=True),
+        gr.update(visible=False),
+        gr.update(visible=False),
+        [],
+        [],
+        session_threads,
+    )
+
+    # Bail early if satellite image missing
+    if sat_path is None:
+        yield (
+            gr.update(interactive=True),
+            gr.update(value=None),
+            gr.update(value=None),
+            gr.update(value="No satellite image provided.", visible=True),
+            gr.update(value="", visible=True),
+            gr.update(visible=False),
+            gr.update(visible=False),
+            [],
+            [],
+            session_threads,
+        )
+        return
+
+    # Prepare PIL images
+    sat_img = Image.open(sat_path).convert("RGB")
+    ground_img_pil = Image.open(ground_path).convert("RGB") if ground_path else None
+
+    # Lookup target positions metadata (may be empty)
+    tgt_positions = []
+    if taxonomy and taxonomy in tgts_metadata:
+        tgt_positions = [tuple(t) for t in tgts_metadata[taxonomy]["target_positions"]]
+
+    # Helper to build a TestWorker with/without TTA
+    def build_planner(enable_tta: bool, save_dir: str, clip_obj):
+        # Lazily (re)create a ClipSegTTA instance per thread if not provided
+        local_clip = clip_obj
+        if LOAD_AVS_BENCH and local_clip is None:
+            local_clip = ClipSegTTA(
+                img_dir=AVS_IMG_DIR,
+                imo_dir=AVS_IMO_DIR,
+                json_path=AVS_INAT_JSON_PATH,
+                sat_to_img_ids_path=AVS_SAT_TO_IMG_IDS_PATH,
+                sat_checkpoint_path=AVS_SAT_CHECKPOINT_PATH,
+                load_pretrained_hf_ckpt=AVS_LOAD_PRETRAINED_HF_CHECKPOINT,
+                blur_kernel = AVS_GAUSSIAN_BLUR_KERNEL,
+                sample_index=-1,    
+                device=device,
+                sat_to_img_ids_json_is_train_dict=False, 
+                tax_to_filter_val=QUERY_TAX,
+                load_model=USE_CLIP_PREDS,
+                query_modality=QUERY_MODALITY,
+                sound_dir = AVS_SOUND_DIR,
+                sound_checkpoint_path=AVS_SOUND_CHECKPOINT_PATH,
+            )
+
+        if local_clip is not None:
+            # Feed inputs to ClipSegTTA copy
+            local_clip.img_paths = [ground_path] if ground_path else []
+            local_clip.imo_path = sat_path
+            local_clip.imgs = ([local_clip.dataset.img_transform(ground_img_pil).to(device)] if ground_img_pil else [])
+            local_clip.imo = local_clip.dataset.imo_transform(sat_img).to(device)
+            local_clip.sounds = []
+            local_clip.sound_ids = []
+            local_clip.species_name = taxonomy or ""
+            local_clip.gt_mask_name = taxonomy.replace(" ", "_") if taxonomy else ""
+            local_clip.target_positions = tgt_positions if tgt_positions else [(0, 0)]
+
+        planner = TestWorker(
+            meta_agent_id=0,
+            n_agent=1,
+            policy_net=policy_net,
+            global_step=-1,
+            device=device,
+            greedy=True,
+            save_image=SAVE_GIFS,
+            clip_seg_tta=local_clip,
+        )
+        planner.execute_tta = enable_tta
+        planner.gifs_path = save_dir
+        return planner
+
+    # ────────────── Per-run output directories ──────────────
+    # Ensure base directory exists
+    os.makedirs(GIFS_PATH, exist_ok=True)
+
+    run_id = time.strftime("%Y%m%d_%H%M%S")  # unique timestamp
+    run_root = os.path.join(GIFS_PATH, run_id)
+    gifs_dir_tta = os.path.join(run_root, "with_tta")
+    gifs_dir_no  = os.path.join(run_root, "no_tta")
+
+    os.makedirs(gifs_dir_tta, exist_ok=True)
+    os.makedirs(gifs_dir_no,  exist_ok=True)
+
+    # House-keep old runs so we never keep more than RUN_HISTORY_LIMIT
+    _prune_old_run_dirs(GIFS_PATH, RUN_HISTORY_LIMIT)
+
+    # Shared dict to record if a thread hit an exception
+    error_flags = {"tta": False, "no": False}
+
+    def _planner_thread(enable_tta: bool, save_dir: str, clip_obj, key: str):
+        """Prepare directory, build planner, run an episode, record errors."""
+        try:
+            planner = build_planner(enable_tta, save_dir, clip_obj)
+            _thread_clip_map[threading.current_thread()] = planner.clip_seg_tta
+            planner.run_episode(0)
+        except Exception as exc:
+            # Mark that this planner crashed so UI can show an error status
+            error_flags[key] = True
+            # Log full traceback so developers can debug via console logs
+            import traceback, sys
+            traceback.print_exc()
+            # Still exit the thread
+            return
+
+    # Launch both planners in background threads – preparation included
+    thread_tta = threading.Thread(
+        target=_planner_thread,
+        args=(True, gifs_dir_tta, None, "tta"),
+        daemon=True,
+    )
+    thread_no = threading.Thread(
+        target=_planner_thread,
+        args=(False, gifs_dir_no, None, "no"),
+        daemon=True,
+    )
+    # Track threads for this user session
+    session_threads.extend([thread_tta, thread_no])
+    thread_tta.start()
+    thread_no.start()
+
+
+    sent_tta: set[str] = set()
+    sent_no:  set[str] = set()
+    last_tta = None
+    last_no  = None
+    # Track previous status strings so we can emit updates when only the
+    # status (Running…/Done.) changes even if no new frame was produced.
+    # Previous status values so we can detect changes and yield updates
+    prev_status_tta = "Initializing model…"
+    prev_status_no  = "Initializing model…"
+
+    try:
+        while thread_tta.is_alive() or thread_no.is_alive():
+            updated = False
+            # Collect new frames from TTA dir
+            pngs = glob.glob(os.path.join(gifs_dir_tta, "*.png"))
+            pngs.sort(key=lambda p: int(os.path.splitext(os.path.basename(p))[0]))
+            for fp in pngs:
+                if fp not in sent_tta:
+                    # Ensure file is fully written (non-empty & readable)
+                    try:
+                        if os.path.getsize(fp) == 0:
+                            continue
+                        with open(fp, "rb") as fh:
+                            fh.read(1)
+                    except Exception:
+                        # Skip this round; we'll retry next poll
+                        continue
+                    sent_tta.add(fp)
+                    last_tta = fp
+                    updated = True
+            # Collect new frames from no-TTA dir
+            pngs = glob.glob(os.path.join(gifs_dir_no, "*.png"))
+            pngs.sort(key=lambda p: int(os.path.splitext(os.path.basename(p))[0]))
+            for fp in pngs:
+                if fp not in sent_no:
+                    try:
+                        if os.path.getsize(fp) == 0:
+                            continue
+                        with open(fp, "rb") as fh:
+                            fh.read(1)
+                    except Exception:
+                        continue
+                    sent_no.add(fp)
+                    last_no = fp
+                    updated = True
+
+            # Determine status based on whether we already have a frame and whether
+            # the corresponding thread is still alive.
+            def _mk_status(last_frame, thread_alive, errored: bool, running_tta: bool=False):
+                if errored:
+                    return "Error!"
+                if last_frame is None:
+                    return "Initializing model…"
+                if not thread_alive:
+                    return "Done."
+                return "Executing TTA (Scheduling GPUs)…" if running_tta else "Executing Planner…"
+
+            exec_tta_flag = False
+            if thread_tta.is_alive():
+                clip_obj = _thread_clip_map.get(thread_tta)
+                if clip_obj is not None and getattr(clip_obj, "executing_tta", False):
+                    exec_tta_flag = True
+
+            status_tta = _mk_status(last_tta, thread_tta.is_alive(), error_flags["tta"], exec_tta_flag)
+            status_no  = _mk_status(last_no,  thread_no.is_alive(), error_flags["no"], False)
+
+            # Determine if we should reveal sliders (once corresponding thread has finished)
+            show_slider_tta = (not thread_tta.is_alive()) and (last_tta is not None)
+            show_slider_no  = (not thread_no.is_alive()) and (last_no  is not None)
+
+            # Build slider updates
+            slider_tta_upd = gr.update()
+            slider_no_upd  = gr.update()
+            frames_tta_upd = gr.update()
+            frames_no_upd  = gr.update()
+
+            if show_slider_tta:
+                n_tta_frames = max(len(sent_tta), 1)
+                slider_tta_upd = gr.update(visible=True, minimum=1, maximum=n_tta_frames, value=n_tta_frames)
+                frames_tta_upd = sorted(sent_tta, key=lambda p: int(os.path.splitext(os.path.basename(p))[0]))
+            if show_slider_no:
+                n_no_frames = max(len(sent_no), 1)
+                slider_no_upd = gr.update(visible=True, minimum=1, maximum=n_no_frames, value=n_no_frames)
+                frames_no_upd = sorted(sent_no, key=lambda p: int(os.path.splitext(os.path.basename(p))[0]))
+
+            # Emit update if we have a new frame OR status changed OR slider visibility changed
+            if (
+                updated
+                or status_tta != prev_status_tta
+                or status_no != prev_status_no
+                or show_slider_tta
+                or show_slider_no
+            ):
+                yield (
+                    gr.update(interactive=False),
+                    last_tta,
+                    last_no,
+                    gr.update(value=status_tta, visible=True),
+                    gr.update(value=status_no, visible=True),
+                    slider_tta_upd,
+                    slider_no_upd,
+                    frames_tta_upd,
+                    frames_no_upd,
+                    session_threads,
+                )
+
+                prev_status_tta = status_tta
+                prev_status_no  = status_no
+
+            time.sleep(POLL_INTERVAL)
+    finally:
+        # Ensure background threads are stopped on cancel
+        for th in (thread_tta, thread_no):
+            if th.is_alive():
+                _stop_thread(th)
+                th.join(timeout=1)
+
+    # Remove finished threads from global registry
+    with _running_threads_lock:
+        # Clear session thread list
+        session_threads.clear()
+
+    # Small delay to ensure last frame files are fully flushed
+    time.sleep(0.2)
+    # One last scan after both threads have finished to catch any frame
+    # that may have been written just before termination but after the last
+    # polling iteration.
+    for fp in sorted(glob.glob(os.path.join(gifs_dir_tta, "*.png")), key=lambda p: int(os.path.splitext(os.path.basename(p))[0])):
+        if fp not in sent_tta:
+            sent_tta.add(fp)
+            last_tta = fp
+    for fp in sorted(glob.glob(os.path.join(gifs_dir_no, "*.png")), key=lambda p: int(os.path.splitext(os.path.basename(p))[0])):
+        if fp not in sent_no:
+            sent_no.add(fp)
+            last_no = fp
+
+    # Prepare frames list and slider configs
+    frames_tta = sorted(glob.glob(os.path.join(gifs_dir_tta, "*.png")), key=lambda p: int(os.path.splitext(os.path.basename(p))[0]))
+    frames_no = sorted(glob.glob(os.path.join(gifs_dir_no, "*.png")), key=lambda p: int(os.path.splitext(os.path.basename(p))[0]))
+    if last_tta is None and frames_tta:
+        last_tta = frames_tta[-1]
+    if last_no is None and frames_no:
+        last_no = frames_no[-1]
+    n_tta = len(frames_tta) or 1  # prevent zero-range slider
+    n_no  = len(frames_no) or 1
+
+    # Final emit: re-enable button, hide statuses, show sliders set to last frame
+    yield (
+        gr.update(interactive=True),
+        last_tta,
+        last_no,
+        gr.update(visible=False),
+        gr.update(visible=False),
+        gr.update(visible=True, minimum=1, maximum=n_tta, value=n_tta),
+        gr.update(visible=True, minimum=1, maximum=n_no,  value=n_no),
+        frames_tta,
+        frames_no,
+        session_threads,
+    )
+
+
+# ────────────────────────── Gradio UI ─────────────────────────────────
+with gr.Blocks(title="Search-TTA (Simplified)", theme=gr.themes.Base()) as demo:
+
+    gr.Markdown(
+        """
+        # Search-TTA: A Multimodal Test-Time Adaptation Framework for Visual Search in the Wild Demo
+        Click on any of the <b>examples below</b> and run the <b>TTA demo</b>. Check out the <b>multimodal heatmap generation feature</b> by switching to the other tab above. <br>
+        Note that the model initialization, RL planner, and TTA updates are not fully optimized on GPU for this huggingface demo, and hence may experience some lag during execution. <br>
+        If you encounter an 'Error' status, refresh the browser and rerun the demo, or try again the next day. We will improve this in the future.  <br>
+        <a href="https://search-tta.github.io">Project Website</a> 
+        """
+    )
+
+    with gr.Row(variant="panel"):
+        with gr.Column():
+            gr.Markdown("### Model Inputs")
+            sat_input = gr.Image(
+                label="Satellite Image",
+                sources=["upload"],
+                type="filepath",
+                height=320,
+            )
+            taxonomy_input = gr.Textbox(
+                label="Full Taxonomy Name (optional)",
+                placeholder="e.g. Animalia Chordata Mammalia Rodentia Sciuridae Marmota marmota",
+            )
+            ground_input = gr.Image(
+                label="Ground-level Image (optional)",
+                sources=["upload"],
+                type="filepath",
+                height=320,
+            )
+            run_btn = gr.Button("Run Search-TTA", variant="primary")
+
+        with gr.Column():
+            gr.Markdown("### Live Heatmap Output")
+            display_img_tta = gr.Image(label="Heatmap (TTA per 20 steps)", type="filepath", height=400)  # 512
+            status_tta = gr.Markdown("")
+            slider_tta = gr.Slider(label="TTA Frame", minimum=1, maximum=1, step=1, value=1, visible=False)
+
+            display_img_no_tta = gr.Image(label="Heatmap (no TTA)", type="filepath", height=400)  # 512
+            status_no_tta = gr.Markdown("")
+            slider_no = gr.Slider(label="No-TTA Frame", minimum=1, maximum=1, step=1, value=1, visible=False)
+
+            frames_state_tta = gr.State([])
+            frames_state_no = gr.State([])
+            session_threads_state = gr.State([])
+
+    # Slider callbacks (updates image when user drags slider)
+    def _show_frame(idx: int, frames: list[str]):
+        # Slider is 1-indexed; convert to 0-indexed list access
+        if 1 <= idx <= len(frames):
+            return frames[idx - 1]
+        return gr.update()
+
+    slider_tta.change(_show_frame, inputs=[slider_tta, frames_state_tta], outputs=display_img_tta)
+    slider_no.change(_show_frame, inputs=[slider_no, frames_state_no], outputs=display_img_no_tta)
+
+    # EXAMPLES
+    with gr.Row():
+        gr.Markdown("### Taxonomy")
+    with gr.Row():
+        gr.Examples(
+            examples=[
+                [
+                    "examples/Animalia_Chordata_Reptilia_Squamata_Varanidae_Varanus_salvator/410613_5.35573_100.28948.jpg",
+                    "examples/Animalia_Chordata_Reptilia_Squamata_Varanidae_Varanus_salvator/461d8e6c-0e66-4acc-8ecd-bfd9c218bc14.jpg",
+                    "Animalia Chordata Reptilia Squamata Varanidae Varanus salvator",
+                ],
+                [
+                    "examples/Animalia_Chordata_Mammalia_Carnivora_Canidae_Canis_aureus/1528408_13.00422_80.23033.jpg",
+                    "examples/Animalia_Chordata_Mammalia_Carnivora_Canidae_Canis_aureus/37faabd2-a613-4461-b27e-82fe5955ecaf.jpg",
+                    "Animalia Chordata Mammalia Carnivora Canidae Canis aureus",
+                ],
+                [
+                    "examples/Animalia_Chordata_Reptilia_Crocodylia_Alligatoridae_Caiman_crocodilus/340271_10.52832_-83.49678.jpg",
+                    "examples/Animalia_Chordata_Reptilia_Crocodylia_Alligatoridae_Caiman_crocodilus/938aab7b-4509-4de7-afad-2c8ea51f4799.jpg",
+                    "Animalia Chordata Reptilia Crocodylia Alligatoridae Caiman crocodilus",
+                ],
+                [
+                    "examples/Animalia_Chordata_Mammalia_Carnivora_Canidae_Urocyon_littoralis/304160_34.0144_-119.54417.jpg",
+                    "examples/Animalia_Chordata_Mammalia_Carnivora_Canidae_Urocyon_littoralis/0cbdfbf2-6cfe-4d61-9602-c949f24d0293.jpg",
+                    "Animalia Chordata Mammalia Carnivora Canidae Urocyon littoralis",
+                ],
+            ],
+            inputs=[sat_input, ground_input, taxonomy_input],
+            outputs=[run_btn, display_img_tta, display_img_no_tta, status_tta, status_no_tta, slider_tta, slider_no, frames_state_tta, frames_state_no],
+            fn=process_search_tta,
+            cache_examples=False,
+        )
+
+    run_btn.click(
+        fn=process_search_tta,
+        inputs=[sat_input, ground_input, taxonomy_input, session_threads_state],
+        outputs=[run_btn, display_img_tta, display_img_no_tta, status_tta, status_no_tta, slider_tta, slider_no, frames_state_tta, frames_state_no, session_threads_state],
+    )
+
+    # Footer to point out to model and data from app page.
+    gr.Markdown(
+        """
+        The satellite image CLIP encoder is fine-tuned using [Sentinel-2 Level 2A](https://docs.sentinel-hub.com/api/latest/data/sentinel-2-l2a/) satellite image and taxonomy images (with GPS locations) from [iNaturalist](https://inaturalist.org/). The sound CLIP encoder is fine-tuned with a subset of the same taxonomy images and their corresponding sounds from [iNaturalist](https://inaturalist.org/). Some of these iNaturalist data are also used in [Taxabind](https://arxiv.org/abs/2411.00683). Note that while some of the examples above result in poor probability distributions, they will be improved using our test-time adaptation framework during the search process.
+        """
+    )
+
+
+if __name__ == "__main__":
+
+    # Build UI with explicit Tabs so we can detect tab selection and clean up
+    from app_multimodal_inference import demo as multimodal_demo
+
+    with gr.Blocks() as root:
+        with gr.Tabs() as tabs:
+            with gr.TabItem("Multimodal Inference"):
+                multimodal_demo.render()
+            with gr.TabItem("Search-TTA"):
+                demo.render()
+
+        # Hidden textbox purely to satisfy Gradio's need for an output component.
+        _cleanup_status = gr.Textbox(visible=False)
+
+        outputs_on_tab = [_cleanup_status]
+
+        def _on_tab_change(evt: gr.SelectData, session_threads: list[threading.Thread]):
+            # evt.value contains the name of the newly-selected tab.
+            if evt.value == "Multimodal Inference":
+                # Stop only threads started in this session
+                for th in list(session_threads):
+                    if th is not None and th.is_alive():
+                        _stop_thread(th)
+                        th.join(timeout=1)
+                session_threads.clear()
+                return "Stopped running Search-TTA threads."
+            return ""
+
+        tabs.select(_on_tab_change, inputs=[session_threads_state], outputs=outputs_on_tab)
+
+    root.queue(max_size=15)
+    root.launch(share=True)

app_multimodal_inference.py

@@ -0,0 +1,350 @@
+"""
+Search-TTA multimodal heatmap generation demo
+"""
+
+# ────────────────────────── imports ───────────────────────────────────
+import cv2
+import gradio as gr
+import torch
+import numpy as np
+from PIL import Image
+import matplotlib.pyplot as plt
+import io
+import torchaudio
+import spaces   # integration with ZeroGPU on hf
+
+from torchvision import transforms
+import open_clip
+from taxabind_avs.satbind.clip_vision_per_patch_model import CLIPVisionPerPatchModel
+from transformers import ClapAudioModelWithProjection
+from transformers import ClapProcessor
+
+# ────────────────────────── global config & models ────────────────────
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# BioCLIP (ground-image & text encoder)
+bio_model, _, _ = open_clip.create_model_and_transforms("hf-hub:imageomics/bioclip")
+bio_model = bio_model.to(device).eval()
+bio_tokenizer = open_clip.get_tokenizer("hf-hub:imageomics/bioclip")
+
+# Satellite patch encoder CLIP-L-336 per-patch)
+sat_model: CLIPVisionPerPatchModel = (
+    CLIPVisionPerPatchModel.from_pretrained("derektan95/search-tta-sat")
+    .to(device)
+    .eval()
+)
+
+# Sound CLAP model
+sound_model: ClapAudioModelWithProjection = (
+    ClapAudioModelWithProjection.from_pretrained("derektan95/search-tta-sound")
+    .to(device)
+    .eval()
+)
+sound_processor: ClapProcessor = ClapProcessor.from_pretrained("derektan95/search-tta-sound")
+SAMPLE_RATE = 48000
+
+logit_scale = torch.nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+logit_scale = logit_scale.exp()
+blur_kernel = (5,5)
+
+# ────────────────────────── transforms (exact spec) ───────────────────
+img_transform = transforms.Compose(
+    [
+        transforms.Resize((256, 256)),
+        transforms.CenterCrop((224, 224)),
+        transforms.ToTensor(),
+        transforms.Normalize(
+            mean=[0.485, 0.456, 0.406],
+            std=[0.229, 0.224, 0.225],
+        ),
+    ]
+)
+
+imo_transform = transforms.Compose(
+    [
+        transforms.Resize((336, 336)),
+        transforms.ToTensor(),
+        transforms.Normalize(
+            mean=[0.485, 0.456, 0.406],
+            std=[0.229, 0.224, 0.225],
+        ),
+    ]
+)
+
+def get_audio_clap(path_to_audio,format="mp3",padding="repeatpad",truncation="fusion"):
+    track, sr = torchaudio.load(path_to_audio, format=format)  # torchaudio.load(path_to_audio)
+    track = track.mean(axis=0)
+    track = torchaudio.functional.resample(track, orig_freq=sr, new_freq=SAMPLE_RATE)
+    output = sound_processor(audios=track, sampling_rate=SAMPLE_RATE, max_length_s=10, return_tensors="pt",padding=padding,truncation=truncation)
+    return output
+
+# ────────────────────────── helpers ───────────────────────────────────
+
+@torch.no_grad()
+def _encode_ground(img_pil: Image.Image) -> torch.Tensor:
+    img = img_transform(img_pil).unsqueeze(0).to(device)
+    img_embeds, *_ = bio_model(img) 
+    return img_embeds
+
+
+@torch.no_grad()
+def _encode_text(text: str) -> torch.Tensor:
+    toks = bio_tokenizer(text).to(device)
+    _, txt_embeds, _ = bio_model(text=toks)
+    return txt_embeds
+
+
+@torch.no_grad()
+def _encode_sat(img_pil: Image.Image) -> torch.Tensor:
+    imo = imo_transform(img_pil).unsqueeze(0).to(device)
+    imo_embeds = sat_model(imo)
+    return imo_embeds
+
+
+@torch.no_grad()
+def _encode_sound(sound) -> torch.Tensor:
+    processed_sound = get_audio_clap(sound)
+    for k in processed_sound.keys():
+        processed_sound[k] = processed_sound[k].to(device)
+    unnormalized_audio_embeds = sound_model(**processed_sound).audio_embeds
+    sound_embeds = torch.nn.functional.normalize(unnormalized_audio_embeds, dim=-1)
+    return sound_embeds
+
+
+def _similarity_heatmap(query: torch.Tensor, patches: torch.Tensor) -> np.ndarray:
+    sims = torch.matmul(query, patches.t()) * logit_scale
+    sims = sims.t().sigmoid()
+    sims = sims[1:].squeeze()  # drop CLS token
+    side = int(np.sqrt(len(sims)))
+    sims = sims.reshape(side, side)
+    return sims.cpu().detach().numpy()
+
+
+def _array_to_pil(arr: np.ndarray) -> Image.Image:
+    """
+    Render arr with viridis, automatically stretching its own min→max to 0→1
+    so that the most-similar patches appear yellow.
+    """
+
+    # Gausian Smoothing
+    if blur_kernel != (0,0):
+        arr = cv2.GaussianBlur(arr, blur_kernel, 0)
+
+    # --- contrast-stretch to local 0-1 range --------------------------
+    arr_min, arr_max = float(arr.min()), float(arr.max())
+    if arr_max - arr_min < 1e-6:        # avoid /0 when the heat-map is flat
+        arr_scaled = np.zeros_like(arr)
+    else:
+        arr_scaled = (arr - arr_min) / (arr_max - arr_min)
+    # ------------------------------------------------------------------
+    fig, ax = plt.subplots(figsize=(2.6, 2.6), dpi=96)
+    ax.imshow(arr_scaled, cmap="viridis", vmin=0.0, vmax=1.0)
+    ax.axis("off")
+    buf = io.BytesIO()
+    plt.tight_layout(pad=0)
+    fig.savefig(buf, format="png", bbox_inches="tight", pad_inches=0)
+    plt.close(fig)
+    buf.seek(0)
+    return Image.open(buf)
+
+# ────────────────────────── main inference ────────────────────────────
+# integration with ZeroGPU on hf
+@spaces.GPU(duration=5)
+def process(
+    sat_img: Image.Image,
+    taxonomy: str,
+    ground_img: Image.Image | None,
+    sound: torch.Tensor | None,
+):
+    if sat_img is None:
+        return None, None
+
+    patches = _encode_sat(sat_img)
+
+    heat_ground, heat_text, heat_sound = None, None, None
+
+    if ground_img is not None:
+        q_img = _encode_ground(ground_img)
+        heat_ground = _array_to_pil(_similarity_heatmap(q_img, patches))
+
+    if taxonomy.strip():
+        q_txt = _encode_text(taxonomy.strip())
+        heat_text = _array_to_pil(_similarity_heatmap(q_txt, patches))
+
+    if sound is not None:
+        q_sound = _encode_sound(sound)
+        heat_sound = _array_to_pil(_similarity_heatmap(q_sound, patches))
+
+    return heat_ground, heat_text, heat_sound
+
+
+# ────────────────────────── Gradio UI ─────────────────────────────────
+with gr.Blocks(title="Search-TTA", theme=gr.themes.Base()) as demo:
+
+    gr.Markdown(
+        """
+        # Search-TTA: A Multimodal Test-Time Adaptation Framework for Visual Search in the Wild Demo
+        Click on any of the <b>examples below</b> and run the <b>multimodal inference demo</b>. Check out the <b>test-time adaptation feature</b> by switching to the other tab above. <br>
+        If you encounter any errors, refresh the browser and rerun the demo, or try again the next day. We will improve this in the future. <br>
+        <a href="https://search-tta.github.io">Project Website</a> 
+        """
+    )
+
+    with gr.Row(variant="panel"):
+
+        # LEFT COLUMN  (satellite, taxonomy, run)
+        with gr.Column():
+            sat_input = gr.Image(
+                label="Satellite Image",
+                sources=["upload"],
+                type="pil",
+                height=320,
+            )
+            taxonomy_input = gr.Textbox(
+                label="Full Taxonomy Name (optional)",
+                placeholder="e.g. Animalia Chordata Mammalia Rodentia Sciuridae Marmota marmota",
+            )
+
+            # ─── NEW: sound input ───────────────────────────
+            sound_input = gr.Audio(
+                label="Sound Input (optional)",
+                sources=["upload"],    
+                type="filepath",     
+            )
+            run_btn = gr.Button("Run", variant="primary")
+
+        # RIGHT COLUMN  (ground image + two heat-maps)
+        with gr.Column():
+            ground_input = gr.Image(
+                label="Ground-level Image (optional)",
+                sources=["upload"],
+                type="pil",
+                height=320,
+            )
+            gr.Markdown("### Heat-map Results")
+            with gr.Row():
+                # Separate label and image to avoid overlap
+                with gr.Column(scale=1, min_width=100):
+                    gr.Markdown("**Ground Image Query**", elem_id="label-ground")
+                    heat_ground_out = gr.Image(
+                        show_label=False,
+                        height=160,
+                    )
+                with gr.Column(scale=1, min_width=100):
+                    gr.Markdown("**Text Query**", elem_id="label-text")
+                    heat_text_out = gr.Image(
+                        show_label=False,
+                        height=160,
+                    )
+                with gr.Column(scale=1, min_width=100):
+                    gr.Markdown("**Sound Query**", elem_id="label-sound")
+                    heat_sound_out = gr.Image(
+                        show_label=False,
+                        height=160,
+                    )
+            
+
+    # EXAMPLES
+    with gr.Row():
+        gr.Markdown("### In-Domain Taxonomy")
+    with gr.Row():
+        gr.Examples(
+            examples=[
+                [
+                    "examples/Animalia_Chordata_Aves_Charadriiformes_Laridae_Larus_marinus/80645_39.76079_-74.10316.jpg",
+                    "examples/Animalia_Chordata_Aves_Charadriiformes_Laridae_Larus_marinus/cc1ebaf9-899d-49f2-81c8-d452249a8087.jpg",
+                    "Animalia Chordata Aves Charadriiformes Laridae Larus marinus",
+                    "examples/Animalia_Chordata_Aves_Charadriiformes_Laridae_Larus_marinus/89758229.mp3"
+                ],
+                [
+                    "examples/Animalia_Chordata_Mammalia_Rodentia_Caviidae_Hydrochoerus_hydrochaeris/28871_-12.80255_-69.29999.jpg",
+                    "examples/Animalia_Chordata_Mammalia_Rodentia_Caviidae_Hydrochoerus_hydrochaeris/1b8064f8-7deb-4b30-98cd-69da98ba6a3d.jpg",
+                    "Animalia Chordata Mammalia Rodentia Caviidae Hydrochoerus hydrochaeris",
+                    "examples/Animalia_Chordata_Mammalia_Rodentia_Caviidae_Hydrochoerus_hydrochaeris/166631961.mp3"
+                ],
+                [
+                    "examples/Animalia_Arthropoda_Malacostraca_Decapoda_Ocypodidae_Ocypode_quadrata/277303_38.72364_-75.07749.jpg",
+                    "examples/Animalia_Arthropoda_Malacostraca_Decapoda_Ocypodidae_Ocypode_quadrata/0b9cc264-a2ba-44bd-8e41-0d01a6edd1e8.jpg",
+                    "Animalia Arthropoda Malacostraca Decapoda Ocypodidae Ocypode quadrata",
+                    "examples/Animalia_Arthropoda_Malacostraca_Decapoda_Ocypodidae_Ocypode_quadrata/12372063.mp3"
+                ],
+                [
+                    "examples/Animalia_Chordata_Mammalia_Rodentia_Sciuridae_Marmota_marmota/388246_45.49036_7.14796.jpg",
+                    "examples/Animalia_Chordata_Mammalia_Rodentia_Sciuridae_Marmota_marmota/327e1f07-692b-4140-8a3e-bd098bc064ff.jpg",
+                    "Animalia Chordata Mammalia Rodentia Sciuridae Marmota marmota",
+                    "examples/Animalia_Chordata_Mammalia_Rodentia_Sciuridae_Marmota_marmota/59677071.mp3"
+                ],
+                [
+                    "examples/Animalia_Chordata_Reptilia_Squamata_Varanidae_Varanus_salvator/410613_5.35573_100.28948.jpg",
+                    "examples/Animalia_Chordata_Reptilia_Squamata_Varanidae_Varanus_salvator/461d8e6c-0e66-4acc-8ecd-bfd9c218bc14.jpg",
+                    "Animalia Chordata Reptilia Squamata Varanidae Varanus salvator",
+                    None
+                ],
+            ],
+            inputs=[sat_input, ground_input, taxonomy_input, sound_input],
+            outputs=[heat_ground_out, heat_text_out, heat_sound_out],
+            fn=process,
+            cache_examples=False,
+        )
+
+    # EXAMPLES
+    with gr.Row():
+        gr.Markdown("### Out-Domain Taxonomy")
+    with gr.Row():
+        gr.Examples(
+            examples=[
+                [
+                    "examples/Animalia_Chordata_Mammalia_Carnivora_Phocidae_Mirounga_angustirostris/27423_35.64005_-121.17595.jpg",
+                    "examples/Animalia_Chordata_Mammalia_Carnivora_Phocidae_Mirounga_angustirostris/3aac526d-c921-452a-af6a-cb4f2f52e2c4.jpg",
+                    "Animalia Chordata Mammalia Carnivora Phocidae Mirounga angustirostris",
+                    "examples/Animalia_Chordata_Mammalia_Carnivora_Phocidae_Mirounga_angustirostris/3123948.mp3"
+                ],
+                [
+                    "examples/Animalia_Chordata_Mammalia_Carnivora_Canidae_Canis_aureus/1528408_13.00422_80.23033.jpg",
+                    "examples/Animalia_Chordata_Mammalia_Carnivora_Canidae_Canis_aureus/37faabd2-a613-4461-b27e-82fe5955ecaf.jpg",
+                    "Animalia Chordata Mammalia Carnivora Canidae Canis aureus",
+                    "examples/Animalia_Chordata_Mammalia_Carnivora_Canidae_Canis_aureus/189318716.mp3"
+                ],
+                [
+                    "examples/Animalia_Chordata_Reptilia_Crocodylia_Alligatoridae_Caiman_crocodilus/340271_10.52832_-83.49678.jpg",
+                    "examples/Animalia_Chordata_Reptilia_Crocodylia_Alligatoridae_Caiman_crocodilus/938aab7b-4509-4de7-afad-2c8ea51f4799.jpg",
+                    "Animalia Chordata Reptilia Crocodylia Alligatoridae Caiman crocodilus",
+                    "examples/Animalia_Chordata_Reptilia_Crocodylia_Alligatoridae_Caiman_crocodilus/45193295.mp3"
+                ],
+                [
+                    "examples/Animalia_Chordata_Mammalia_Carnivora_Canidae_Urocyon_littoralis/304160_34.0144_-119.54417.jpg",
+                    "examples/Animalia_Chordata_Mammalia_Carnivora_Canidae_Urocyon_littoralis/0cbdfbf2-6cfe-4d61-9602-c949f24d0293.jpg",
+                    "Animalia Chordata Mammalia Carnivora Canidae Urocyon littoralis",
+                    None
+                ],
+                [
+                    "examples/Animalia_Chordata_Elasmobranchii_Carcharhiniformes_Carcharhinidae_Triaenodon_obesus/5041_-0.28573_-90.54837.jpg",
+                    "examples/Animalia_Chordata_Elasmobranchii_Carcharhiniformes_Carcharhinidae_Triaenodon_obesus/c834edf7-b073-4cd5-8726-9c88ebe943c0.jpg",
+                    "Animalia Chordata Elasmobranchii Carcharhiniformes Carcharhinidae Triaenodon obesus",
+                    None
+                ],
+            ],
+            inputs=[sat_input, ground_input, taxonomy_input, sound_input],
+            outputs=[heat_ground_out, heat_text_out, heat_sound_out],
+            fn=process,
+            cache_examples=False,
+        )
+
+    # CALLBACK
+    run_btn.click(
+        fn=process,
+        inputs=[sat_input, taxonomy_input, ground_input, sound_input],
+        outputs=[heat_ground_out, heat_text_out, heat_sound_out],
+    )
+
+    # Footer to point out to model and data from app page.
+    gr.Markdown(
+        """
+        The satellite image CLIP encoder is fine-tuned using [Sentinel-2 Level 2A](https://docs.sentinel-hub.com/api/latest/data/sentinel-2-l2a/) satellite image and taxonomy images (with GPS locations) from [iNaturalist](https://inaturalist.org/). The sound CLIP encoder is fine-tuned with a subset of the same taxonomy images and their corresponding sounds from [iNaturalist](https://inaturalist.org/). Some of these iNaturalist data are also used in [Taxabind](https://arxiv.org/abs/2411.00683). Note that while some of the examples above result in poor probability distributions, they will be improved using our test-time adaptation framework during the search process.
+        """
+    )
+
+# LAUNCH
+if __name__ == "__main__":
+    demo.queue(max_size=15)
+    demo.launch(share=True)

examples/Animalia_Arthropoda_Malacostraca_Decapoda_Ocypodidae_Ocypode_quadrata/12372063.mp3

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:575959883981159f2e40593bf5be87be006026c41da36a34d1e40783de648116
+size 54027

examples/Animalia_Chordata_Aves_Charadriiformes_Laridae_Larus_marinus/89758229.mp3

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e2c7ad6df49668d29f9b7f9f9f0739b97ef4edc5219413a41d01983a9863cccc
+size 2601487

examples/Animalia_Chordata_Mammalia_Carnivora_Canidae_Canis_aureus/189318716.mp3

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4639c226ea5a0464b98e89b33a1f821b6625c6637d206d3d355e05bc7c89c641
+size 148019

examples/Animalia_Chordata_Mammalia_Carnivora_Phocidae_Mirounga_angustirostris/3123948.mp3

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:96ca3a92e6f614cce82972dacb04f5c0c170c1aea3d70d15778af56820ed02c9
+size 276768

examples/Animalia_Chordata_Mammalia_Rodentia_Caviidae_Hydrochoerus_hydrochaeris/166631961.mp3

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cc02eca19d0c408d038e205d82f6624c0515858ac374cf7298161a14e169e6a9
+size 266258

examples/Animalia_Chordata_Mammalia_Rodentia_Sciuridae_Marmota_marmota/59677071.mp3

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bb043991fe851d6a1e12f32c5a9277dad5a77a939cf15ccb4afcb215b4bc08e3
+size 92876

examples/Animalia_Chordata_Reptilia_Crocodylia_Alligatoridae_Caiman_crocodilus/45193295.mp3

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c4cd2e4fd7094a07d79da7fd54788705e8ce7567e65911d87edfd23ff1c0e484
+size 247762

examples/metadata.json

@@ -0,0 +1,173 @@
+{
+  "Animalia Chordata Reptilia Squamata Varanidae Varanus salvator": {
+    "id": 410613,
+    "sat_key": "410613_5.35573_100.28948",
+    "sat_path": "410613_5.35573_100.28948.jpg",
+    "taxonomy": "Animalia Chordata Reptilia Squamata Varanidae Varanus salvator",
+    "count": 6,
+    "spread": 58.00460580210422,
+    "sat_bounds": {
+      "min_lat": 5.344155081363914,
+      "max_lat": 5.367304914271601,
+      "min_lon": 100.27793148340874,
+      "max_lon": 100.30102851659126
+    },
+    "img_ids": [
+      707815,
+      411949,
+      701168,
+      1619682,
+      2100008,
+      1548498
+    ],
+    "target_positions": [
+      [
+        225,
+        240
+      ],
+      [
+        232,
+        275
+      ],
+      [
+        277,
+        449
+      ],
+      [
+        220,
+        369
+      ],
+      [
+        180,
+        393
+      ],
+      [
+        294,
+        478
+      ]
+    ],
+    "num_landmarks": 2
+  },
+  "Animalia Chordata Mammalia Carnivora Canidae Canis aureus": {
+    "id": 1528408,
+    "sat_key": "1528408_13.00422_80.23033",
+    "sat_path": "1528408_13.00422_80.23033.jpg",
+    "taxonomy": "Animalia Chordata Mammalia Carnivora Canidae Canis aureus",
+    "count": 3,
+    "spread": 58.14007011752667,
+    "sat_bounds": {
+      "min_lat": 12.992649951077192,
+      "max_lat": 13.015790038631529,
+      "min_lon": 80.21853090802841,
+      "max_lon": 80.24212909197156
+    },
+    "img_ids": [
+      1528479,
+      2555188,
+      2555189
+    ],
+    "target_positions": [
+      [
+        309,
+        128
+      ],
+      [
+        239,
+        428
+      ],
+      [
+        240,
+        419
+      ]
+    ],
+    "num_landmarks": 3
+  },
+  "Animalia Chordata Reptilia Crocodylia Alligatoridae Caiman crocodilus": {
+    "id": 340271,
+    "sat_key": "340271_10.52832_-83.49678",
+    "sat_path": "340271_10.52832_-83.49678.jpg",
+    "taxonomy": "Animalia Chordata Reptilia Crocodylia Alligatoridae Caiman crocodilus",
+    "count": 7,
+    "spread": 40.13902957324975,
+    "sat_bounds": {
+      "min_lat": 10.516747947357544,
+      "max_lat": 10.53989204420829,
+      "min_lon": -83.50847402265151,
+      "max_lon": -83.48508597734848
+    },
+    "img_ids": [
+      1683531,
+      1281855,
+      223089,
+      688111,
+      330757,
+      2408375,
+      1955359
+    ],
+    "target_positions": [
+      [
+        347,
+        75
+      ],
+      [
+        47,
+        22
+      ],
+      [
+        111,
+        43
+      ],
+      [
+        116,
+        51
+      ],
+      [
+        86,
+        108
+      ],
+      [
+        31,
+        62
+      ],
+      [
+        4,
+        78
+      ]
+    ],
+    "num_landmarks": 3
+  },
+  "Animalia Chordata Mammalia Carnivora Canidae Urocyon littoralis": {
+    "id": 304160,
+    "sat_key": "304160_34.0144_-119.54417",
+    "sat_path": "304160_34.0144_-119.54417.jpg",
+    "taxonomy": "Animalia Chordata Mammalia Carnivora Canidae Urocyon littoralis",
+    "count": 3,
+    "spread": 237.64152837579553,
+    "sat_bounds": {
+      "min_lat": 34.00286041606169,
+      "max_lat": 34.02593956225012,
+      "min_lon": -119.55802743361286,
+      "max_lon": -119.53031256638712
+    },
+    "img_ids": [
+      304160,
+      1473173,
+      384867
+    ],
+    "target_positions": [
+      [
+        255,
+        256
+      ],
+      [
+        19,
+        22
+      ],
+      [
+        29,
+        274
+      ]
+    ],
+    "num_landmarks": 3
+  }
+}

inference/model/avs_rl_policy.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:44e642df9aaa2847ba44dd4707985c67ef712f5264272ef7993aeb7805c80f5a
+size 52167246

planner/env.py

@@ -0,0 +1,610 @@
+#######################################################################
+# Name: env.py
+#
+# - Reads and processes training and test maps 
+# - Processes rewards, new frontiers given action
+# - Updates a graph representation of environment for input into network
+#######################################################################
+
+import sys
+if sys.modules['TRAINING']:
+    from .parameter import *
+else:
+    from .test_parameter import *
+
+import os
+import cv2
+import copy
+import matplotlib.image as mpimg
+import matplotlib.pyplot as plt
+from skimage import io
+from skimage.measure import block_reduce
+from scipy.ndimage import label, find_objects
+from .sensor import *
+from .graph_generator import *
+from .node import *
+
+
+class Env():
+    def __init__(self, map_index, n_agent, k_size=20, plot=False, test=False, mask_index=None):
+        self.n_agent = n_agent
+        self.test = test
+        self.map_dir = GRIDMAP_SET_DIR 
+
+        # Import environment gridmap
+        self.map_list = os.listdir(self.map_dir)
+        self.map_list.sort(reverse=True)
+
+        # NEW: Import segmentation utility map
+        self.seg_dir = MASK_SET_DIR    
+        self.segmentation_mask, self.target_positions, self.target_found_idxs = None, [], []
+        self.segmentation_mask_list = os.listdir(self.seg_dir)
+        self.segmentation_mask_list.sort(reverse=True)
+
+        # # NEW: Find common files in both directories
+        self.map_index = map_index % len(self.map_list)
+        if mask_index is not None:
+            self.mask_index = mask_index % len(self.segmentation_mask_list)
+        else:
+            self.mask_index = map_index % len(self.segmentation_mask_list)
+
+        # Import ground truth and segmentation mask
+        self.ground_truth, self.map_start_position = self.import_ground_truth(
+            os.path.join(self.map_dir, self.map_list[self.map_index]))
+        self.ground_truth_size = np.shape(self.ground_truth)  
+        self.robot_belief = np.ones(self.ground_truth_size) * 127  # unexplored 127
+        self.downsampled_belief = None
+        self.old_robot_belief = copy.deepcopy(self.robot_belief)
+        self.coverage_belief = np.ones(self.ground_truth_size) * 127  # unexplored 127
+
+        # Import segmentation mask
+        mask_filename = self.segmentation_mask_list[self.mask_index]
+        self.segmentation_mask = self.import_segmentation_mask(
+            os.path.join(self.seg_dir, mask_filename))
+        
+        # Overwrite target positions if directory specified
+        if self.test and TARGETS_SET_DIR != "":
+            self.target_positions = self.import_targets(
+                os.path.join(TARGETS_SET_DIR, self.map_list[self.map_index])) 
+        
+        self.segmentation_info_mask = None
+        self.segmentation_info_mask_unnormalized = None
+        self.filtered_seg_info_mask = None
+        self.num_targets_found = 0
+        self.num_new_targets_found = 0
+        self.resolution = 4
+        self.sensor_range = SENSOR_RANGE
+        self.explored_rate = 0
+        self.targets_found_rate = 0
+        self.frontiers = None
+        self.start_positions = []
+        self.plot = plot
+        self.frame_files = []
+        self.graph_generator = Graph_generator(map_size=self.ground_truth_size, sensor_range=self.sensor_range, k_size=k_size, plot=plot)
+        self.node_coords, self.graph, self.node_utility, self.guidepost = None, None, None, None
+
+        self.begin(self.map_start_position)
+
+
+    def find_index_from_coords(self, position):
+        index = np.argmin(np.linalg.norm(self.node_coords - position, axis=1))
+        return index
+
+    def begin(self, start_position):
+        self.robot_belief = self.ground_truth   
+        self.downsampled_belief = block_reduce(self.robot_belief.copy(), block_size=(self.resolution, self.resolution), func=np.min)
+        self.frontiers = self.find_frontier()
+        self.old_robot_belief = copy.deepcopy(self.robot_belief)
+
+        self.node_coords, self.graph, self.node_utility, self.guidepost = self.graph_generator.generate_graph(
+                self.robot_belief, self.frontiers)
+        
+        # Define start positions
+        if FIX_START_POSITION:
+            coords_res_row = int(self.robot_belief.shape[0]/NUM_COORDS_HEIGHT)
+            coords_res_col = int(self.robot_belief.shape[1]/NUM_COORDS_WIDTH)
+            self.start_positions = [(int(self.robot_belief.shape[1]/2)-coords_res_col/2,int(self.robot_belief.shape[0]/2)-coords_res_row/2)  for _ in range(self.n_agent)]   
+        else:
+            nearby_coords = self.graph_generator.get_neighbors_grid_coords(start_position)
+            itr = 0
+            for i in range(self.n_agent):
+                if i == 0 or len(nearby_coords) == 0:
+                    self.start_positions.append(start_position)
+                else:
+                    idx = min(itr, len(nearby_coords)-1)
+                    self.start_positions.append(nearby_coords[idx])
+                    itr += 1
+
+        for i in range(len(self.start_positions)):
+            self.start_positions[i] = self.node_coords[self.find_index_from_coords(self.start_positions[i])]
+            self.coverage_belief = self.update_robot_belief(self.start_positions[i], self.sensor_range, self.coverage_belief,
+                                                        self.ground_truth)
+
+        for start_position in self.start_positions:
+            self.graph_generator.route_node.append(start_position)
+
+        # Info map from ground truth
+        rng_x = 0.5 * (self.ground_truth.shape[1] / NUM_COORDS_WIDTH)
+        rng_y = 0.5 * (self.ground_truth.shape[0] / NUM_COORDS_HEIGHT)
+        self.segmentation_info_mask = np.zeros((len(self.node_coords), 1))
+        for i, node_coord in enumerate(self.node_coords):
+            max_x = min(node_coord[0] + int(math.ceil(rng_x)), self.ground_truth.shape[1])
+            min_x = max(node_coord[0] - int(math.ceil(rng_x)), 0)
+            max_y = min(node_coord[1] + int(math.ceil(rng_y)), self.ground_truth.shape[0])
+            min_y = max(node_coord[1] - int(math.ceil(rng_y)), 0)
+
+            if TARGETS_SET_DIR == "":   
+                exclude = {208} # Exclude target positions 
+            else:
+                exclude = {}
+            self.segmentation_info_mask[i] = max(x for x in self.segmentation_mask[min_y:max_y, min_x:max_x].flatten() if x not in exclude) / 100.0
+
+        self.filtered_seg_info_mask = copy.deepcopy(self.segmentation_info_mask)
+        done, num_targets_found = self.check_done()
+        self.num_targets_found = num_targets_found
+
+
+    def multi_robot_step(self, next_position_list, dist_list, travel_dist_list):
+        reward_list = []
+        for dist, robot_position in zip(dist_list, next_position_list):
+            self.graph_generator.route_node.append(robot_position)
+            next_node_index = self.find_index_from_coords(robot_position)
+            self.graph_generator.nodes_list[next_node_index].set_visited()
+            self.coverage_belief = self.update_robot_belief(robot_position, self.sensor_range, self.coverage_belief,
+                                                         self.ground_truth)
+            self.robot_belief = self.ground_truth   
+            self.downsampled_belief = block_reduce(self.robot_belief.copy(),
+                                                   block_size=(self.resolution, self.resolution),
+                                                   func=np.min)
+
+            frontiers = self.find_frontier()
+            individual_reward = -dist / 32 
+
+            info_gain_reward = 0
+            robot_position_idx = self.find_index_from_coords(robot_position)
+            info_gain_reward = self.filtered_seg_info_mask[robot_position_idx][0]  * 1.5
+            if self.guidepost[robot_position_idx] == 0.0:
+                info_gain_reward += 0.2
+            individual_reward += info_gain_reward
+
+            reward_list.append(individual_reward)
+
+        self.node_coords, self.graph, self.node_utility, self.guidepost = self.graph_generator.update_graph(self.robot_belief, self.old_robot_belief, frontiers, self.frontiers)
+        self.old_robot_belief = copy.deepcopy(self.robot_belief)
+
+        self.filtered_seg_info_mask = [info[0] if self.guidepost[i] == 0.0 else 0.0 for i, info in enumerate(self.segmentation_info_mask)]
+        self.filtered_seg_info_mask = np.expand_dims(np.array(self.filtered_seg_info_mask), axis=1)
+
+        self.frontiers = frontiers
+        self.explored_rate = self.evaluate_exploration_rate()
+
+        done, num_targets_found = self.check_done()
+        self.num_new_targets_found = num_targets_found - self.num_targets_found
+        team_reward = 0.0
+
+        self.num_targets_found = num_targets_found
+        self.targets_found_rate = self.evaluate_targets_found_rate()
+
+        if done:
+            team_reward += 40 
+        for i in range(len(reward_list)):
+            reward_list[i] += team_reward
+
+        return reward_list, done
+
+
+    def import_ground_truth(self, map_index):
+        # occupied 1, free 255, unexplored 127
+
+        try:
+            ground_truth = (io.imread(map_index, 1)).astype(int)
+            if np.all(ground_truth == 0):
+                ground_truth = (io.imread(map_index, 1) * 255).astype(int)
+        except:
+            new_map_index = self.map_dir + '/' + self.map_list[0]
+            ground_truth = (io.imread(new_map_index, 1)).astype(int)
+            print('could not read the map_path ({}), hence skipping it and using ({}).'.format(map_index, new_map_index))
+
+        robot_location = np.nonzero(ground_truth == 208)
+        robot_location = np.array([np.array(robot_location)[1, 127], np.array(robot_location)[0, 127]])
+        ground_truth = (ground_truth > 150)
+        ground_truth = ground_truth * 254 + 1
+        return ground_truth, robot_location
+
+
+    def import_segmentation_mask(self, map_index):
+        mask = cv2.imread(map_index).astype(int)
+        return mask 
+
+    def import_targets(self, map_index):
+        # occupied 1, free 255, unexplored 127, target 208
+        mask = cv2.imread(map_index).astype(int)
+        target_positions = self.find_target_locations(mask)
+        return target_positions
+
+
+    def find_target_locations(self, image_array, grey_value=208):
+
+        grey_pixels = np.where(image_array == grey_value)
+        binary_array = np.zeros_like(image_array, dtype=bool)
+        binary_array[grey_pixels] = True
+        labeled_array, num_features = label(binary_array)
+        slices = find_objects(labeled_array)
+
+        # Calculate the center of each box
+        centers = []
+        for slice in slices:
+            row_center = (slice[0].start + slice[0].stop - 1) // 2
+            col_center = (slice[1].start + slice[1].stop - 1) // 2
+            centers.append((col_center, row_center))    # (y,x)
+
+        return centers
+
+    def free_cells(self):
+        index = np.where(self.ground_truth == 255)
+        free = np.asarray([index[1], index[0]]).T
+        return free
+
+    def update_robot_belief(self, robot_position, sensor_range, robot_belief, ground_truth):
+        robot_belief = sensor_work(robot_position, sensor_range, robot_belief, ground_truth)
+        return robot_belief
+
+
+    def check_done(self):
+        done = False
+        num_targets_found = 0
+        self.target_found_idxs = []
+        for i, target in enumerate(self.target_positions):
+            if self.coverage_belief[target[1], target[0]] == 255: 
+                num_targets_found += 1
+                self.target_found_idxs.append(i)
+
+        if TERMINATE_ON_TGTS_FOUND and num_targets_found >= len(self.target_positions):
+            done = True
+        if not TERMINATE_ON_TGTS_FOUND and np.sum(self.coverage_belief == 255) / np.sum(self.ground_truth == 255) >= 0.99:
+            done = True
+        
+        return done, num_targets_found
+
+
+    def calculate_num_observed_frontiers(self, old_frontiers, frontiers):
+        frontiers_to_check = frontiers[:, 0] + frontiers[:, 1] * 1j
+        pre_frontiers_to_check = old_frontiers[:, 0] + old_frontiers[:, 1] * 1j
+        frontiers_num = np.intersect1d(frontiers_to_check, pre_frontiers_to_check).shape[0]
+        pre_frontiers_num = pre_frontiers_to_check.shape[0]
+        delta_num = pre_frontiers_num - frontiers_num
+
+        return delta_num
+
+    def calculate_reward(self, dist, frontiers):
+        reward = 0
+        reward -= dist / 64
+
+        frontiers_to_check = frontiers[:, 0] + frontiers[:, 1] * 1j
+        pre_frontiers_to_check = self.frontiers[:, 0] + self.frontiers[:, 1] * 1j
+        frontiers_num = np.intersect1d(frontiers_to_check, pre_frontiers_to_check).shape[0]
+        pre_frontiers_num = pre_frontiers_to_check.shape[0]
+        delta_num = pre_frontiers_num - frontiers_num
+
+        reward += delta_num / 50
+
+        return reward
+
+    def evaluate_exploration_rate(self):
+        rate = np.sum(self.coverage_belief == 255) / np.sum(self.ground_truth == 255)
+        return rate
+
+    def evaluate_targets_found_rate(self):
+        if len(self.target_positions) == 0:
+            return 0
+        else:
+            rate = self.num_targets_found / len(self.target_positions)
+            return rate
+
+    def calculate_new_free_area(self):
+        old_free_area = self.old_robot_belief == 255
+        current_free_area = self.robot_belief == 255
+
+        new_free_area = (current_free_area.astype(np.int) - old_free_area.astype(np.int)) * 255
+
+        return new_free_area, np.sum(old_free_area)
+
+    def calculate_dist_path(self, path):
+        dist = 0
+        start = path[0]
+        end = path[-1]
+        for index in path:
+            if index == end:
+                break
+            dist += np.linalg.norm(self.node_coords[start] - self.node_coords[index])
+            start = index
+        return dist
+
+    def find_frontier(self):
+        y_len = self.downsampled_belief.shape[0]
+        x_len = self.downsampled_belief.shape[1]
+        mapping = self.downsampled_belief.copy()
+        belief = self.downsampled_belief.copy()
+        # 0-1 unknown area map
+        mapping = (mapping == 127) * 1
+        mapping = np.lib.pad(mapping, ((1, 1), (1, 1)), 'constant', constant_values=0)
+        fro_map = mapping[2:][:, 1:x_len + 1] + mapping[:y_len][:, 1:x_len + 1] + mapping[1:y_len + 1][:, 2:] + \
+                  mapping[1:y_len + 1][:, :x_len] + mapping[:y_len][:, 2:] + mapping[2:][:, :x_len] + mapping[2:][:,
+                                                                                                      2:] + \
+                  mapping[:y_len][:, :x_len]
+        ind_free = np.where(belief.ravel(order='F') == 255)[0]
+        ind_fron_1 = np.where(1 < fro_map.ravel(order='F'))[0]
+        ind_fron_2 = np.where(fro_map.ravel(order='F') < 8)[0]
+        ind_fron = np.intersect1d(ind_fron_1, ind_fron_2)
+        ind_to = np.intersect1d(ind_free, ind_fron)
+
+        map_x = x_len
+        map_y = y_len
+        x = np.linspace(0, map_x - 1, map_x)
+        y = np.linspace(0, map_y - 1, map_y)
+        t1, t2 = np.meshgrid(x, y)
+        points = np.vstack([t1.T.ravel(), t2.T.ravel()]).T
+
+        f = points[ind_to]
+        f = f.astype(int)
+
+        f = f * self.resolution
+
+        return f
+
+
+
+    def plot_env(self, n, path, step, travel_dist, robots_route, img_path_override=None, sat_path_override=None, msk_name_override=None, sound_id_override=None):
+
+        plt.switch_backend('agg')
+        plt.cla()
+        color_list = ["r", "g", "c", "m", "y", "k"]
+
+        if not LOAD_AVS_BENCH:
+            fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 5))
+        else:
+            fig, (ax1, ax2, ax3, ax4) = plt.subplots(1, 4, figsize=(20, 5.5))
+
+        ### Fig: Segmentation Mask ###
+        if LOAD_AVS_BENCH:
+            ax = ax1
+            image = mpimg.imread(img_path_override)
+            ax.imshow(image)
+            ax.set_title("Ground Image")
+            ax.axis("off")
+
+        ### Fig: Environment ###
+        msk_name = ""
+        if LOAD_AVS_BENCH:
+            image = mpimg.imread(sat_path_override)
+            msk_name = msk_name_override
+
+            ### Fig1: Environment ###
+            ax = ax2
+            ax.imshow(image)
+            ax.axis((0, self.ground_truth_size[1], self.ground_truth_size[0], 0))
+            ax.set_title("Image")
+            for i, route in enumerate(robots_route):
+                robot_marker_color = color_list[i % len(color_list)]
+                xPoints = route[0]
+                yPoints = route[1]
+                ax.plot(xPoints, yPoints, c=robot_marker_color, linewidth=2)
+                ax.plot(xPoints[-1], yPoints[-1], markersize=12, zorder=99, marker="^", ls="-", c=robot_marker_color, mec="black")
+                ax.plot(xPoints[0], yPoints[0], 'co', c=robot_marker_color, markersize=8, zorder=5)
+
+                # Sensor range
+                rng_x = 0.5 * (self.ground_truth.shape[1] / NUM_COORDS_WIDTH)
+                rng_y = 0.5 * (self.ground_truth.shape[0] / NUM_COORDS_HEIGHT)
+                max_x = min(xPoints[-1] + int(math.ceil(rng_x)), self.ground_truth.shape[1])
+                min_x = max(xPoints[-1] - int(math.ceil(rng_x)), 0)
+                max_y = min(yPoints[-1] + int(math.ceil(rng_y)), self.ground_truth.shape[0])
+                min_y = max(yPoints[-1] - int(math.ceil(rng_y)), 0)
+                ax.plot((min_x, min_x), (min_y, max_y), c=robot_marker_color, linewidth=1)
+                ax.plot((min_x, max_x), (max_y, max_y), c=robot_marker_color, linewidth=1)
+                ax.plot((max_x, max_x), (max_y, min_y), c=robot_marker_color, linewidth=1)
+                ax.plot((max_x, min_x), (min_y, min_y), c=robot_marker_color, linewidth=1)
+
+
+        ### Fig: Graph  ###
+        ax = ax3 if LOAD_AVS_BENCH else ax1
+        ax.imshow(self.coverage_belief, cmap='gray')
+        ax.axis((0, self.ground_truth_size[1], self.ground_truth_size[0], 0))
+        ax.set_title("Information Graph")
+        if VIZ_GRAPH_EDGES:
+            for i in range(len(self.graph_generator.x)):
+                ax.plot(self.graph_generator.x[i], self.graph_generator.y[i], 'tan', zorder=1)
+        ax.scatter(self.node_coords[:, 0], self.node_coords[:, 1], c=self.filtered_seg_info_mask, zorder=5, s=8)
+
+        for i, route in enumerate(robots_route):
+            robot_marker_color = color_list[i % len(color_list)]
+            xPoints = route[0]
+            yPoints = route[1]
+            ax.plot(xPoints, yPoints, c=robot_marker_color, linewidth=2)
+            ax.plot(xPoints[-1], yPoints[-1], markersize=12, zorder=99, marker="^", ls="-", c=robot_marker_color, mec="black")
+            ax.plot(xPoints[0], yPoints[0], 'co', c=robot_marker_color, markersize=8, zorder=5)
+
+            # Sensor range
+            rng_x = 0.5 * (self.ground_truth.shape[1] / NUM_COORDS_WIDTH)
+            rng_y = 0.5 * (self.ground_truth.shape[0] / NUM_COORDS_HEIGHT)
+            max_x = min(xPoints[-1] + int(math.ceil(rng_x)), self.ground_truth.shape[1])
+            min_x = max(xPoints[-1] - int(math.ceil(rng_x)), 0)
+            max_y = min(yPoints[-1] + int(math.ceil(rng_y)), self.ground_truth.shape[0])
+            min_y = max(yPoints[-1] - int(math.ceil(rng_y)), 0)
+            ax.plot((min_x, min_x), (min_y, max_y), c=robot_marker_color, linewidth=1)
+            ax.plot((min_x, max_x), (max_y, max_y), c=robot_marker_color, linewidth=1)
+            ax.plot((max_x, max_x), (max_y, min_y), c=robot_marker_color, linewidth=1)
+            ax.plot((max_x, min_x), (min_y, min_y), c=robot_marker_color, linewidth=1)
+
+        # Plot target positions
+        for target in self.target_positions:
+            if self.coverage_belief[target[1], target[0]] == 255:
+                ax.plot(target[0], target[1], color='g', marker='x', linestyle='-', markersize=12, markeredgewidth=4, zorder=99)
+            else:
+                ax.plot(target[0], target[1], color='r', marker='x', linestyle='-', markersize=12, markeredgewidth=4, zorder=99)
+
+        ### Fig: Segmentation Mask ###
+        ax = ax4 if LOAD_AVS_BENCH else ax2
+        if LOAD_AVS_BENCH and USE_CLIP_PREDS:
+            H, W = self.ground_truth_size  
+            mask_viz = self.segmentation_info_mask.squeeze().reshape((NUM_COORDS_WIDTH, NUM_COORDS_HEIGHT)).T
+            im = ax.imshow(
+                mask_viz,
+                cmap="viridis",
+                origin="upper",
+                extent=[0, W, H, 0],  
+                interpolation="nearest",  
+                zorder=0,
+            )
+            ax.set_xlim(0, W)
+            ax.set_ylim(H, 0)
+            ax.set_axis_off()  
+        else:
+            im = ax.imshow(self.segmentation_mask.mean(axis=-1), cmap='viridis', vmin=0, vmax=100)  # cmap='gray'
+            ax.axis((0, self.ground_truth_size[1], self.ground_truth_size[0], 0))
+        ax.set_title(f"Predicted Mask (Normalized)")
+        for i, route in enumerate(robots_route):
+            robot_marker_color = color_list[i % len(color_list)]
+            xPoints = route[0]
+            yPoints = route[1]
+            ax.plot(xPoints, yPoints, c=robot_marker_color, linewidth=2)
+            ax.plot(xPoints[-1], yPoints[-1], markersize=12, zorder=99, marker="^", ls="-", c=robot_marker_color, mec="black")
+            ax.plot(xPoints[0], yPoints[0], 'co', c=robot_marker_color, markersize=8, zorder=5)
+
+            # Sensor range
+            rng_x = 0.5 * (self.ground_truth.shape[1] / NUM_COORDS_WIDTH)
+            rng_y = 0.5 * (self.ground_truth.shape[0] / NUM_COORDS_HEIGHT)
+            max_x = min(xPoints[-1] + int(math.ceil(rng_x)), self.ground_truth.shape[1])
+            min_x = max(xPoints[-1] - int(math.ceil(rng_x)), 0)
+            max_y = min(yPoints[-1] + int(math.ceil(rng_y)), self.ground_truth.shape[0])
+            min_y = max(yPoints[-1] - int(math.ceil(rng_y)), 0)
+            ax.plot((min_x, min_x), (min_y, max_y), c=robot_marker_color, linewidth=1)
+            ax.plot((min_x, max_x), (max_y, max_y), c=robot_marker_color, linewidth=1)
+            ax.plot((max_x, max_x), (max_y, min_y), c=robot_marker_color, linewidth=1)
+            ax.plot((max_x, min_x), (min_y, min_y), c=robot_marker_color, linewidth=1)
+
+        # Add a colorbar 
+        cbar = fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
+        cbar.set_label("Normalized Probs")
+                    
+        if sound_id_override is not None:
+            plt.suptitle('Targets Found: {}/{}  Coverage ratio: {:.4g}  Travel Dist: {:.4g} \n ({}) \n (Sound ID: {})'.format(self.num_targets_found, len(self.target_positions), self.explored_rate, travel_dist, msk_name, sound_id_override))
+        elif msk_name != "":
+            plt.suptitle('Targets Found: {}/{}  Coverage ratio: {:.4g}  Travel Dist: {:.4g} \n ({})'.format(self.num_targets_found, len(self.target_positions), self.explored_rate, travel_dist, msk_name))
+        else:
+            plt.suptitle('Targets Found: {}/{}  Coverage ratio: {:.4g}  Travel Dist: {:.4g}'.format(self.num_targets_found, len(self.target_positions), self.explored_rate, travel_dist))
+
+        plt.tight_layout()
+        plt.savefig('{}/{}_{}_samples.png'.format(path, n, step, dpi=100))
+        frame = '{}/{}_{}_samples.png'.format(path, n, step)
+        self.frame_files.append(frame)
+        plt.close()
+
+
+    ####################
+    # ADDED: For app.py
+    ####################
+
+    def plot_heatmap(self, save_dir, step, travel_dist, robots_route=None):
+        """Plot only the segmentation heatmap and save it as ``{step}.png`` in
+        ``save_dir``. This lightweight helper is meant for asynchronous
+        streaming in the Gradio demo when full `plot_env` is too heavy.
+
+        Parameters
+        ----------
+        save_dir : str
+            Directory to save the generated PNG file.
+        step : int
+            Current timestep; becomes the filename ``{step}.png``.
+        robots_route : list | None
+            Optional list of routes (xPoints, yPoints) to overlay.
+        Returns
+        -------
+        str
+            Full path to the generated PNG file.
+        """
+        import os
+        plt.switch_backend('agg')
+        # Do not clear the global figure state in case it interferes with
+        # the current figure. Each call creates its own Figure object that
+        # we close explicitly at the end, so a global clear is unnecessary
+        # and may break concurrent drawing.
+        # plt.cla()
+
+        color_list = ["r", "g", "c", "m", "y", "k"]
+        fig, ax = plt.subplots(1, 1, figsize=(6, 6))
+
+        # Select the mask to visualise
+        # if TAXABIND_TTA and USE_CLIP_PREDS:
+        side_dim = int(np.sqrt(self.segmentation_info_mask.shape[0]))
+        mask_viz = self.segmentation_info_mask.squeeze().reshape((side_dim, side_dim)).T
+
+        # Properly map image to pixel coordinates and keep limits fixed
+        H, W = self.ground_truth_size  # rows (y), cols (x)
+        im = ax.imshow(
+            mask_viz,
+            cmap="viridis",
+            origin="upper",
+            extent=[0, W, H, 0],  # x: 0..W, y: H..0 (origin at top-left)
+            interpolation="nearest",  # keep cell edges sharp & aligned
+            zorder=0,
+        )
+        ax.set_xlim(0, W)
+        ax.set_ylim(H, 0)
+        ax.set_axis_off()  # hide ticks but keep limits
+        # else:
+        #     im = ax.imshow(self.segmentation_mask.mean(axis=-1), cmap='viridis', vmin=0, vmax=100)
+        #     ax.axis((0, self.ground_truth_size[1], self.ground_truth_size[0], 0))
+
+        # Optionally overlay robot paths
+        if robots_route is not None:
+            for i, route in enumerate(robots_route):
+                robot_marker_color = color_list[i % len(color_list)]
+                xPoints, yPoints = route
+                ax.plot(xPoints, yPoints, c=robot_marker_color, linewidth=2)
+                ax.plot(xPoints[-1], yPoints[-1], markersize=12, zorder=99, marker="^", ls="-", c=robot_marker_color, mec="black")
+                ax.plot(xPoints[0], yPoints[0], 'co', c=robot_marker_color, markersize=8, zorder=5)
+
+        # Plot target positions
+        for target in self.target_positions:
+            if self.coverage_belief[target[1], target[0]] == 255:
+                # ax.plot(target[0], target[1], 'go', markersize=8, zorder=99)
+                ax.plot(target[0], target[1], color='g', marker='x', linestyle='-', markersize=12, markeredgewidth=4, zorder=99)
+            else:
+                # ax.plot(target[0], target[1], 'ro', markersize=8, zorder=99)
+                ax.plot(target[0], target[1], color='r', marker='x', linestyle='-', markersize=12, markeredgewidth=4, zorder=99)
+
+        # Sensor range
+        rng_x = 0.5 * (self.ground_truth.shape[1] / NUM_COORDS_WIDTH)
+        rng_y = 0.5 * (self.ground_truth.shape[0] / NUM_COORDS_HEIGHT)
+        max_x = min(xPoints[-1] + int(math.ceil(rng_x)), self.ground_truth.shape[1])
+        min_x = max(xPoints[-1] - int(math.ceil(rng_x)), 0)
+        max_y = min(yPoints[-1] + int(math.ceil(rng_y)), self.ground_truth.shape[0])
+        min_y = max(yPoints[-1] - int(math.ceil(rng_y)), 0)
+        ax.plot((min_x, min_x), (min_y, max_y), c=robot_marker_color, linewidth=1)
+        ax.plot((min_x, max_x), (max_y, max_y), c=robot_marker_color, linewidth=1)
+        ax.plot((max_x, max_x), (max_y, min_y), c=robot_marker_color, linewidth=1)
+        ax.plot((max_x, min_x), (min_y, min_y), c=robot_marker_color, linewidth=1)
+
+        # Color bar
+        cbar = fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
+        cbar.set_label("Normalized Probs")
+
+        # Change coverage to 1dp
+        plt.suptitle('Targets Found: {}/{}  Coverage: {:.1f}%  Steps: {}/{}'.format(
+            self.num_targets_found, \
+            len(self.target_positions), 
+            self.explored_rate*100, 
+            step+1, 
+            NUM_EPS_STEPS), 
+            y=0.94,     # Closer to plot
+        )
+        
+        plt.tight_layout()
+        os.makedirs(save_dir, exist_ok=True)
+        out_path = os.path.join(save_dir, f"{step}.png")
+        # Save atomically: write to temp file then move into place so the poller never sees a partial file.
+        tmp_path = out_path + ".tmp"
+        fig.savefig(tmp_path, dpi=100, format='png')
+        os.replace(tmp_path, out_path)  # atomic on same filesystem
+        plt.close(fig)
+        return out_path

planner/graph.py

@@ -0,0 +1,167 @@
+#######################################################################
+# Name: env.py
+#
+# - Adapted from https://gist.github.com/betandr/541a1f6466b6855471de5ca30b74cb31
+# - Simple graph class to perform distance calculations (E.g. A-Star, Djikstra)
+#######################################################################
+
+
+class Edge:
+    def __init__(self, to_node, length):
+        self.to_node = to_node
+        self.length = length
+
+
+class Graph:
+    def __init__(self):
+        self.nodes = set()
+        self.edges = dict()
+
+    def add_node(self, node):
+        self.nodes.add(node)
+
+    def add_edge(self, from_node, to_node, length):
+        edge = Edge(to_node, length)
+        if from_node in self.edges:
+            from_node_edges = self.edges[from_node]
+        else:
+            self.edges[from_node] = dict()
+            from_node_edges = self.edges[from_node]
+        from_node_edges[to_node] = edge
+
+    def clear_edge(self, from_node):
+        if from_node in self.edges:
+            self.edges[from_node] = dict()
+
+def min_dist(q, dist):
+    """
+    Returns the node with the smallest distance in q.
+    Implemented to keep the main algorithm clean.
+    """
+    min_node = None
+    for node in q:
+        if min_node == None:
+            min_node = node
+        elif dist[node] < dist[min_node]:
+            min_node = node
+
+    return min_node
+
+
+INFINITY = float('Infinity')
+
+
+def dijkstra(graph, source):
+    q = set()
+    dist = {}
+    prev = {}
+
+    for v in graph.nodes:       
+        dist[v] = INFINITY      # unknown distance from source to v
+        prev[v] = INFINITY      # previous node in optimal path from source
+        q.add(v)                # all nodes initially in q (unvisited nodes)
+
+    # distance from source to source
+    dist[source] = 0
+
+    while q:
+        # node with the least distance selected first
+        u = min_dist(q, dist)
+
+        q.remove(u)
+
+        try:
+            if u in graph.edges:
+                for _, v in graph.edges[u].items():
+                    alt = dist[u] + v.length
+                    if alt < dist[v.to_node]:
+                        # a shorter path to v has been found
+                        dist[v.to_node] = alt
+                        prev[v.to_node] = u
+        except:
+            pass
+
+    return dist, prev
+
+
+def to_array(prev, from_node):
+    """Creates an ordered list of labels as a route."""
+    previous_node = prev[from_node]
+    route = [from_node]
+
+    while previous_node != INFINITY:
+        route.append(previous_node)
+        temp = previous_node
+        previous_node = prev[temp]
+
+    route.reverse()
+    return route
+
+
+def h(index, destination, node_coords):
+    current = node_coords[index]
+    end = node_coords[destination]
+    h = abs(end[0] - current[0]) + abs(end[1] - current[1])
+    return h
+
+
+def a_star(start, destination, node_coords, graph):
+    if start == destination:
+        return [], 0
+    if str(destination) in graph.edges[str(start)].keys():
+        cost = graph.edges[str(start)][str(destination)].length
+        return [start, destination], cost
+    open_list = {start}
+    closed_list = set([])
+
+    g = {start: 0}
+    parents = {start: start}
+
+    while len(open_list) > 0:
+        n = None
+        h_n = 1e5
+        for v in open_list:
+            h_v = h(v, destination, node_coords)
+            if n is not None:
+                h_n = h(n, destination, node_coords)
+            if n is None or g[v] + h_v < g[n] + h_n:
+                n = v
+
+        if n is None:
+            print('Path does not exist!')
+            return None, 1e5
+
+        if n == destination:
+            reconst_path = []
+            while parents[n] != n:
+                reconst_path.append(n)
+                n = parents[n]
+            reconst_path.append(start)
+            reconst_path.reverse()
+            return reconst_path, g[destination]
+
+        for edge in graph.edges[str(n)].values():
+            m = int(edge.to_node)
+            cost = edge.length
+            if m not in open_list and m not in closed_list:
+                open_list.add(m)
+                parents[m] = n
+                g[m] = g[n] + cost
+
+            else:
+                if g[m] > g[n] + cost:
+                    g[m] = g[n] + cost
+                    parents[m] = n
+
+                    if m in closed_list:
+                        closed_list.remove(m)
+                        open_list.add(m)
+
+        open_list.remove(n)
+        closed_list.add(n)
+
+    print('Path does not exist!')
+    return None, 1e5
+
+
+

planner/graph_generator.py

@@ -0,0 +1,300 @@
+#######################################################################
+# Name: graph_generator.py
+#
+# - Wrapper for graph.py
+# - Sends the formatted inputs into graph.py to get useful info 
+#######################################################################
+
+import sys
+if sys.modules['TRAINING']:
+    from .parameter import *
+else:
+    from .test_parameter import *
+
+import numpy as np
+import shapely.geometry
+from sklearn.neighbors import NearestNeighbors
+from .node import Node
+from .graph import Graph, a_star
+
+
+class Graph_generator:
+    def __init__(self, map_size, k_size, sensor_range, plot=False):
+        self.k_size = k_size
+        self.graph = Graph()
+        self.node_coords = None
+        self.plot = plot
+        self.x = []
+        self.y = []
+        self.map_x = map_size[1]
+        self.map_y = map_size[0]
+        self.uniform_points, self.grid_coords = self.generate_uniform_points()
+        self.sensor_range = sensor_range
+        self.route_node = []
+        self.nodes_list = []
+        self.node_utility = None
+        self.guidepost = None
+
+
+    def edge_clear_all_nodes(self):
+        self.graph = Graph()
+        self.x = []
+        self.y = []
+
+
+    def edge_clear(self, coords):
+        node_index = str(self.find_index_from_coords(self.node_coords, coords))
+        self.graph.clear_edge(node_index)
+
+
+    def generate_graph(self, robot_belief, frontiers):
+        self.edge_clear_all_nodes()
+        free_area = self.free_area(robot_belief)
+
+        free_area_to_check = free_area[:, 0] + free_area[:, 1] * 1j
+        uniform_points_to_check = self.uniform_points[:, 0] + self.uniform_points[:, 1] * 1j
+        _, _, candidate_indices = np.intersect1d(free_area_to_check, uniform_points_to_check, return_indices=True)
+        node_coords = self.uniform_points[candidate_indices]
+
+        self.node_coords = self.unique_coords(node_coords).reshape(-1, 2)
+        self.find_nearest_neighbor_all_nodes(self.node_coords, robot_belief)
+
+        self.node_utility = []
+        for coords in self.node_coords:
+            node = Node(coords, frontiers, robot_belief)
+            self.nodes_list.append(node)
+            utility = node.utility
+            self.node_utility.append(utility)
+        self.node_utility = np.array(self.node_utility)
+
+        self.guidepost = np.zeros((self.node_coords.shape[0], 1))
+        x = self.node_coords[:,0] + self.node_coords[:,1]*1j
+        for node in self.route_node:
+            index = np.argwhere(x.reshape(-1) == node[0]+node[1]*1j)[0]
+            self.guidepost[index] = 1
+
+        return self.node_coords, self.graph.edges, self.node_utility, self.guidepost
+
+
+    def update_graph(self, robot_belief, old_robot_belief, frontiers, old_frontiers):
+        new_free_area = self.free_area((robot_belief - old_robot_belief > 0) * 255)
+        free_area_to_check = new_free_area[:, 0] + new_free_area[:, 1] * 1j
+        uniform_points_to_check = self.uniform_points[:, 0] + self.uniform_points[:, 1] * 1j
+        _, _, candidate_indices = np.intersect1d(free_area_to_check, uniform_points_to_check, return_indices=True)
+        new_node_coords = self.uniform_points[candidate_indices]
+        self.node_coords = np.concatenate((self.node_coords, new_node_coords))
+
+        old_node_to_update = []
+        for coords in new_node_coords:
+            neighbor_indices = self.find_k_neighbor(coords, self.node_coords, robot_belief)
+            old_node_to_update += neighbor_indices
+        old_node_to_update = set(old_node_to_update)
+        for index in old_node_to_update:
+            coords = self.node_coords[index]
+            self.edge_clear(coords)
+            self.find_k_neighbor(coords, self.node_coords, robot_belief)
+
+        old_frontiers_to_check = old_frontiers[:, 0] + old_frontiers[:, 1] * 1j
+        new_frontiers_to_check = frontiers[:, 0] + frontiers[:, 1] * 1j
+        observed_frontiers_index = np.where(
+            np.isin(old_frontiers_to_check, new_frontiers_to_check, assume_unique=True) == False)
+        new_frontiers_index = np.where(
+            np.isin(new_frontiers_to_check, old_frontiers_to_check, assume_unique=True) == False)
+        observed_frontiers = old_frontiers[observed_frontiers_index]
+        new_frontiers = frontiers[new_frontiers_index]
+        for node in self.nodes_list:
+            if node.zero_utility_node is True:
+                pass
+            else:
+                node.update_observable_frontiers(observed_frontiers, new_frontiers, robot_belief)
+
+        for new_coords in new_node_coords:
+            node = Node(new_coords, frontiers, robot_belief)
+            self.nodes_list.append(node)
+
+        self.node_utility = []
+        for i, coords in enumerate(self.node_coords):
+            utility = self.nodes_list[i].utility
+            self.node_utility.append(utility)
+        self.node_utility = np.array(self.node_utility)
+
+        self.guidepost = np.zeros((self.node_coords.shape[0], 1))
+        x = self.node_coords[:, 0] + self.node_coords[:, 1] * 1j
+        for node in self.route_node:
+            index = np.argwhere(x.reshape(-1) == node[0] + node[1] * 1j)
+            self.guidepost[index] = 1
+
+        return self.node_coords, self.graph.edges, self.node_utility, self.guidepost
+
+
+    def generate_uniform_points(self):
+        padding_x = 0.5 * (self.map_x / NUM_COORDS_WIDTH)
+        padding_y = 0.5 * (self.map_y / NUM_COORDS_HEIGHT)
+        x = np.linspace(padding_x, self.map_x - padding_x - 1, NUM_COORDS_WIDTH).round().astype(int)
+        y = np.linspace(padding_y, self.map_y - padding_y - 1, NUM_COORDS_HEIGHT).round().astype(int)
+
+        t1, t2 = np.meshgrid(x, y)
+        points = np.vstack([t1.T.ravel(), t2.T.ravel()]).T
+        matrix = np.stack((t1, t2), axis=-1)
+        return points, matrix
+
+
+    def free_area(self, robot_belief):
+        index = np.where(robot_belief == 255)
+        free = np.asarray([index[1], index[0]]).T
+        return free
+
+
+    def unique_coords(self, coords):
+        x = coords[:, 0] + coords[:, 1] * 1j
+        indices = np.unique(x, return_index=True)[1]
+        coords = np.array([coords[idx] for idx in sorted(indices)])
+        return coords
+
+
+    def find_k_neighbor(self, coords, node_coords, robot_belief):
+        dist_list = np.linalg.norm((coords-node_coords), axis=-1)
+        sorted_index = np.argsort(dist_list)
+        k = 0
+        neighbor_index_list = []
+        while k < self.k_size and k< node_coords.shape[0]:
+            neighbor_index = sorted_index[k]
+            neighbor_index_list.append(neighbor_index)
+            dist = dist_list[k]
+            start = coords
+            end = node_coords[neighbor_index]
+            if not self.check_collision(start, end, robot_belief):
+                a = str(self.find_index_from_coords(node_coords, start))
+                b = str(neighbor_index)
+                self.graph.add_node(a)
+                self.graph.add_edge(a, b, dist)
+
+                if self.plot:
+                    self.x.append([start[0], end[0]])
+                    self.y.append([start[1], end[1]])
+            k += 1
+        return neighbor_index_list
+
+
+    def find_k_neighbor_all_nodes(self, node_coords, robot_belief):
+        X = node_coords
+        if len(node_coords) >= self.k_size:
+            knn = NearestNeighbors(n_neighbors=self.k_size)
+        else:
+            knn = NearestNeighbors(n_neighbors=len(node_coords))
+        knn.fit(X)
+        distances, indices = knn.kneighbors(X)
+
+        for i, p in enumerate(X):
+            for j, neighbour in enumerate(X[indices[i][:]]):
+                start = p
+                end = neighbour
+                if not self.check_collision(start, end, robot_belief):
+                    a = str(self.find_index_from_coords(node_coords, p))
+                    b = str(self.find_index_from_coords(node_coords, neighbour))
+                    self.graph.add_node(a)
+                    self.graph.add_edge(a, b, distances[i, j])
+
+                    if self.plot:
+                        self.x.append([p[0], neighbour[0]])
+                        self.y.append([p[1], neighbour[1]])
+
+
+    def find_nearest_neighbor_all_nodes(self, node_coords, robot_belief):
+        for i, p in enumerate(node_coords):
+            filtered_coords = self.get_neighbors_grid_coords(p)
+
+            for j, neighbour in enumerate(filtered_coords):
+                start = p
+                end = neighbour
+                if not self.check_collision(start, end, robot_belief):
+                    a = str(self.find_index_from_coords(node_coords, p))
+                    b = str(self.find_index_from_coords(node_coords, neighbour))
+                    self.graph.add_node(a)
+                    self.graph.add_edge(a, b, np.linalg.norm(start-end))
+
+                    if self.plot:
+                        self.x.append([p[0], neighbour[0]])
+                        self.y.append([p[1], neighbour[1]])
+
+
+    def find_index_from_coords(self, node_coords, p):
+        return np.where(np.linalg.norm(node_coords - p, axis=1) < 1e-5)[0][0]
+
+
+    def find_closest_index_from_coords(self, node_coords, p):
+        return np.argmin(np.linalg.norm(node_coords - p, axis=1))
+
+
+    def find_index_from_grid_coords_2d(self, p):
+        diffs = np.linalg.norm(self.grid_coords - p, axis=2)  
+        indices = np.where(diffs < 1e-5)
+        
+        if indices[0].size > 0:
+            return indices[0][0], indices[1][0]
+        else:
+            raise ValueError(f"Coordinate {p} not found in self.grid_coords.")
+
+
+    def find_closest_index_from_grid_coords_2d(self, p):
+        distances = np.linalg.norm(self.grid_coords - p, axis=2)  
+        flat_index = np.argmin(distances)
+        return np.unravel_index(flat_index, distances.shape)
+    
+
+    def check_collision(self, start, end, robot_belief):
+        collision = False
+        line = shapely.geometry.LineString([start, end])
+
+        sortx = np.sort([start[0], end[0]])
+        sorty = np.sort([start[1], end[1]])
+
+        robot_belief = robot_belief[sorty[0]:sorty[1]+1, sortx[0]:sortx[1]+1]
+
+        occupied_area_index = np.where(robot_belief == 1)
+        occupied_area_coords = np.asarray([occupied_area_index[1]+sortx[0], occupied_area_index[0]+sorty[0]]).T
+        unexplored_area_index = np.where(robot_belief == 127)
+        unexplored_area_coords = np.asarray([unexplored_area_index[1]+sortx[0], unexplored_area_index[0]+sorty[0]]).T
+        unfree_area_coords = np.concatenate((occupied_area_coords, unexplored_area_coords))
+
+        for i in range(unfree_area_coords.shape[0]):
+            coords = ([(unfree_area_coords[i][0], unfree_area_coords[i][1]),
+                       (unfree_area_coords[i][0] + 1, unfree_area_coords[i][1]),
+                       (unfree_area_coords[i][0], unfree_area_coords[i][1] + 1),
+                       (unfree_area_coords[i][0] + 1, unfree_area_coords[i][1] + 1)])
+            obstacle = shapely.geometry.Polygon(coords)
+            collision = line.intersects(obstacle)
+            if collision:
+                break
+
+        return collision
+
+
+    def find_shortest_path(self, current, destination, node_coords):
+        start_node = str(self.find_index_from_coords(node_coords, current))
+        end_node = str(self.find_index_from_coords(node_coords, destination))
+        route, dist = a_star(int(start_node), int(end_node), self.node_coords, self.graph)
+        if start_node != end_node:
+            assert route != []
+        route = list(map(str, route))
+        return dist, route
+
+    def get_neighbors_grid_coords(self, coord):
+        # Return the 8 closest neighbors of a given coordinate
+        
+        nearest_coord = self.node_coords[self.find_closest_index_from_coords(self.node_coords, coord)]
+        rows, cols = self.grid_coords.shape[:2]
+        neighbors = []
+        i, j  = self.find_index_from_grid_coords_2d(nearest_coord)
+
+        # Create a range of indices for rows and columns
+        row_range = np.clip([i - 1, i, i + 1], 0, rows - 1)
+        col_range = np.clip([j - 1, j, j + 1], 0, cols - 1)
+
+        # Iterate over the valid indices
+        for ni in row_range:
+            for nj in col_range:
+                if (ni, nj) != (i, j):  # Skip the center point
+                    neighbors.append(tuple(self.grid_coords[ni, nj]))
+
+        return neighbors

planner/model.py

@@ -0,0 +1,312 @@
+#######################################################################
+# Name: model.py
+#
+# - Attention-based encoders & decoders
+# - Policy Net: Input = Augmented Graph, Output = Node to go to
+# - Critic Net: Input = Augmented Graph + Action, Output = Q_Value
+#######################################################################
+
+import torch
+import torch.nn as nn
+import math
+
+
+class SingleHeadAttention(nn.Module):
+    def __init__(self, embedding_dim):
+        super(SingleHeadAttention, self).__init__()
+        self.input_dim = embedding_dim
+        self.embedding_dim = embedding_dim
+        self.value_dim = embedding_dim
+        self.key_dim = self.value_dim
+        self.tanh_clipping = 10
+        self.norm_factor = 1 / math.sqrt(self.key_dim)
+
+        self.w_query = nn.Parameter(torch.Tensor(self.input_dim, self.key_dim))
+        self.w_key = nn.Parameter(torch.Tensor(self.input_dim, self.key_dim))
+
+        self.init_parameters()
+
+    def init_parameters(self):
+        for param in self.parameters():
+            stdv = 1. / math.sqrt(param.size(-1))
+            param.data.uniform_(-stdv, stdv)
+
+    def forward(self, q, k, mask=None):
+
+        n_batch, n_key, n_dim = k.size()
+        n_query = q.size(1)
+
+        k_flat = k.reshape(-1, n_dim)
+        q_flat = q.reshape(-1, n_dim)
+
+        shape_k = (n_batch, n_key, -1)
+        shape_q = (n_batch, n_query, -1)
+
+        Q = torch.matmul(q_flat, self.w_query).view(shape_q)
+        K = torch.matmul(k_flat, self.w_key).view(shape_k)
+
+        U = self.norm_factor * torch.matmul(Q, K.transpose(1, 2))
+        U = self.tanh_clipping * torch.tanh(U)
+
+        if mask is not None:
+            U = U.masked_fill(mask == 1, -1e8)
+        attention = torch.log_softmax(U, dim=-1)  # n_batch*n_query*n_key
+
+        return attention
+
+
+class MultiHeadAttention(nn.Module):
+    def __init__(self, embedding_dim, n_heads=8):
+        super(MultiHeadAttention, self).__init__()
+        self.n_heads = n_heads
+        self.input_dim = embedding_dim
+        self.embedding_dim = embedding_dim
+        self.value_dim = self.embedding_dim // self.n_heads
+        self.key_dim = self.value_dim
+        self.norm_factor = 1 / math.sqrt(self.key_dim)
+
+        self.w_query = nn.Parameter(torch.Tensor(self.n_heads, self.input_dim, self.key_dim))
+        self.w_key = nn.Parameter(torch.Tensor(self.n_heads, self.input_dim, self.key_dim))
+        self.w_value = nn.Parameter(torch.Tensor(self.n_heads, self.input_dim, self.value_dim))
+        self.w_out = nn.Parameter(torch.Tensor(self.n_heads, self.value_dim, self.embedding_dim))
+
+        self.init_parameters()
+
+    def init_parameters(self):
+        for param in self.parameters():
+            stdv = 1. / math.sqrt(param.size(-1))
+            param.data.uniform_(-stdv, stdv)
+
+    def forward(self, q, k=None, v=None, key_padding_mask=None, attn_mask=None):
+        if k is None:
+            k = q
+        if v is None:
+            v = q
+
+        n_batch, n_key, n_dim = k.size()
+        n_query = q.size(1)
+        n_value = v.size(1)
+
+        k_flat = k.contiguous().view(-1, n_dim)
+        v_flat = v.contiguous().view(-1, n_dim)
+        q_flat = q.contiguous().view(-1, n_dim)
+        shape_v = (self.n_heads, n_batch, n_value, -1)
+        shape_k = (self.n_heads, n_batch, n_key, -1)
+        shape_q = (self.n_heads, n_batch, n_query, -1)
+
+        Q = torch.matmul(q_flat, self.w_query).view(shape_q)  # n_heads*batch_size*n_query*key_dim
+        K = torch.matmul(k_flat, self.w_key).view(shape_k)  # n_heads*batch_size*targets_size*key_dim
+        V = torch.matmul(v_flat, self.w_value).view(shape_v)  # n_heads*batch_size*targets_size*value_dim
+
+        U = self.norm_factor * torch.matmul(Q, K.transpose(2, 3))  # n_heads*batch_size*n_query*targets_size
+
+        if attn_mask is not None:
+            attn_mask = attn_mask.view(1, n_batch, n_query, n_key).expand_as(U)
+
+        if key_padding_mask is not None:
+            key_padding_mask = key_padding_mask.repeat(1, n_query, 1)
+            key_padding_mask = key_padding_mask.view(1, n_batch, n_query, n_key).expand_as(U)  # copy for n_heads times
+
+        if attn_mask is not None and key_padding_mask is not None:
+            mask = (attn_mask + key_padding_mask)
+        elif attn_mask is not None:
+            mask = attn_mask
+        elif key_padding_mask is not None:
+            mask = key_padding_mask
+        else:
+            mask = None
+
+        if mask is not None:
+            U = U.masked_fill(mask > 0, -1e8)
+
+        attention = torch.softmax(U, dim=-1)  # n_heads*batch_size*n_query*targets_size
+        heads = torch.matmul(attention, V)  # n_heads*batch_size*n_query*value_dim
+        out = torch.mm(
+            heads.permute(1, 2, 0, 3).reshape(-1, self.n_heads * self.value_dim),
+            # batch_size*n_query*n_heads*value_dim
+            self.w_out.view(-1, self.embedding_dim)
+            # n_heads*value_dim*embedding_dim
+        ).view(-1, n_query, self.embedding_dim)
+
+
+        return out, attention  # batch_size*n_query*embedding_dim
+
+
+class Normalization(nn.Module):
+    def __init__(self, embedding_dim):
+        super(Normalization, self).__init__()
+        self.normalizer = nn.LayerNorm(embedding_dim)
+
+    def forward(self, input):
+        return self.normalizer(input.view(-1, input.size(-1))).view(*input.size())
+
+
+class EncoderLayer(nn.Module):
+    def __init__(self, embedding_dim, n_head):
+        super(EncoderLayer, self).__init__()
+        self.multiHeadAttention = MultiHeadAttention(embedding_dim, n_head)
+        self.normalization1 = Normalization(embedding_dim)
+        self.feedForward = nn.Sequential(nn.Linear(embedding_dim, 512), nn.ReLU(inplace=True),
+                                         nn.Linear(512, embedding_dim))
+        self.normalization2 = Normalization(embedding_dim)
+
+    def forward(self, src, key_padding_mask=None, attn_mask=None):
+        h0 = src
+        h = self.normalization1(src)
+        h, _ = self.multiHeadAttention(q=h, key_padding_mask=key_padding_mask, attn_mask=attn_mask)
+        h = h + h0
+        h1 = h
+        h = self.normalization2(h)
+        h = self.feedForward(h)
+        h2 = h + h1
+        return h2
+
+
+class DecoderLayer(nn.Module):
+    def __init__(self, embedding_dim, n_head):
+        super(DecoderLayer, self).__init__()
+        self.multiHeadAttention = MultiHeadAttention(embedding_dim, n_head)
+        self.normalization1 = Normalization(embedding_dim)
+        self.feedForward = nn.Sequential(nn.Linear(embedding_dim, 512),
+                                         nn.ReLU(inplace=True),
+                                         nn.Linear(512, embedding_dim))
+        self.normalization2 = Normalization(embedding_dim)
+
+    def forward(self, tgt, memory, key_padding_mask=None, attn_mask=None):
+        h0 = tgt
+        tgt = self.normalization1(tgt)
+        memory = self.normalization1(memory)
+        h, w = self.multiHeadAttention(q=tgt, k=memory, v=memory, key_padding_mask=key_padding_mask, attn_mask=attn_mask)
+        h = h + h0
+        h1 = h
+        h = self.normalization2(h)
+        h = self.feedForward(h)
+        h2 = h + h1
+        return h2, w
+
+
+class Encoder(nn.Module):
+    def __init__(self, embedding_dim=128, n_head=8, n_layer=1):
+        super(Encoder, self).__init__()
+        self.layers = nn.ModuleList(EncoderLayer(embedding_dim, n_head) for i in range(n_layer))
+
+    def forward(self, src, key_padding_mask=None, attn_mask=None):
+        for layer in self.layers:
+            src = layer(src, key_padding_mask=key_padding_mask, attn_mask=attn_mask)
+        return src
+
+
+class Decoder(nn.Module):
+    def __init__(self, embedding_dim=128, n_head=8, n_layer=1):
+        super(Decoder, self).__init__()
+        self.layers = nn.ModuleList([DecoderLayer(embedding_dim, n_head) for i in range(n_layer)])
+
+    def forward(self, tgt, memory, key_padding_mask=None, attn_mask=None):
+        for layer in self.layers:
+            tgt, w = layer(tgt, memory, key_padding_mask=key_padding_mask, attn_mask=attn_mask)
+        return tgt, w
+
+
+class PolicyNet(nn.Module):
+    def __init__(self, input_dim, embedding_dim):
+        super(PolicyNet, self).__init__()
+        self.initial_embedding = nn.Linear(input_dim, embedding_dim) # layer for non-end position
+
+        self.current_embedding = nn.Linear(embedding_dim * 2, embedding_dim)
+
+        self.encoder = Encoder(embedding_dim=embedding_dim, n_head=8, n_layer=6)
+        self.decoder = Decoder(embedding_dim=embedding_dim, n_head=8, n_layer=1)
+        self.pointer = SingleHeadAttention(embedding_dim)
+
+    def encode_graph(self, node_inputs, node_padding_mask, edge_mask):
+        node_feature = self.initial_embedding(node_inputs)
+        enhanced_node_feature = self.encoder(src=node_feature, key_padding_mask=node_padding_mask, attn_mask=edge_mask)
+
+        return enhanced_node_feature
+
+    def output_policy(self, enhanced_node_feature, edge_inputs, current_index, edge_padding_mask, node_padding_mask):
+        k_size = edge_inputs.size()[2]
+        current_edge = torch.gather(edge_inputs, 1, current_index.repeat(1, 1, k_size))
+        current_edge = current_edge.permute(0, 2, 1)
+        embedding_dim = enhanced_node_feature.size()[2]
+
+        neigboring_feature = torch.gather(enhanced_node_feature, 1, current_edge.repeat(1, 1, embedding_dim))
+
+        current_node_feature = torch.gather(enhanced_node_feature, 1, current_index.repeat(1, 1, embedding_dim))
+
+        if edge_padding_mask is not None:
+            current_mask = torch.gather(edge_padding_mask, 1, current_index.repeat(1,1,k_size)).to(enhanced_node_feature.device)
+        else:
+            current_mask = None
+        current_mask[:,:,0] = 1 # don't stay at current position
+
+        if not 0 in current_mask:
+            current_mask[:,:,0] = 0
+
+        enhanced_current_node_feature, _ = self.decoder(current_node_feature, enhanced_node_feature, node_padding_mask)
+        enhanced_current_node_feature = self.current_embedding(torch.cat((enhanced_current_node_feature, current_node_feature), dim=-1))
+        logp = self.pointer(enhanced_current_node_feature, neigboring_feature, current_mask)
+        logp= logp.squeeze(1) # batch_size*k_size
+
+        return logp
+
+    def forward(self, node_inputs, edge_inputs, current_index, node_padding_mask=None, edge_padding_mask=None, edge_mask=None):
+        enhanced_node_feature = self.encode_graph(node_inputs, node_padding_mask, edge_mask)
+        logp = self.output_policy(enhanced_node_feature, edge_inputs, current_index, edge_padding_mask, node_padding_mask)
+        return logp
+
+
+class QNet(nn.Module):
+    def __init__(self, input_dim, embedding_dim):
+        super(QNet, self).__init__()
+        self.initial_embedding = nn.Linear(input_dim, embedding_dim) # layer for non-end position
+        self.action_embedding = nn.Linear(embedding_dim*3, embedding_dim)
+
+        self.encoder = Encoder(embedding_dim=embedding_dim, n_head=8, n_layer=6)
+        self.decoder = Decoder(embedding_dim=embedding_dim, n_head=8, n_layer=1)
+
+        self.q_values_layer = nn.Linear(embedding_dim, 1)
+
+    def encode_graph(self, node_inputs, node_padding_mask, edge_mask):
+        embedding_feature = self.initial_embedding(node_inputs)
+        embedding_feature = self.encoder(src=embedding_feature, key_padding_mask=node_padding_mask, attn_mask=edge_mask)
+
+        return embedding_feature
+
+    def output_q_values(self, enhanced_node_feature, edge_inputs, current_index, edge_padding_mask, node_padding_mask):
+        k_size = edge_inputs.size()[2]
+        current_edge = torch.gather(edge_inputs, 1, current_index.repeat(1, 1, k_size))
+        current_edge = current_edge.permute(0, 2, 1)
+        embedding_dim = enhanced_node_feature.size()[2]
+
+        neigboring_feature = torch.gather(enhanced_node_feature, 1, current_edge.repeat(1, 1, embedding_dim))
+
+        current_node_feature = torch.gather(enhanced_node_feature, 1, current_index.repeat(1, 1, embedding_dim))
+
+        enhanced_current_node_feature, attention_weights = self.decoder(current_node_feature, enhanced_node_feature, node_padding_mask)
+        action_features = torch.cat((enhanced_current_node_feature.repeat(1, k_size, 1), current_node_feature.repeat(1, k_size, 1), neigboring_feature), dim=-1)
+        action_features = self.action_embedding(action_features)
+        q_values = self.q_values_layer(action_features)
+
+        if edge_padding_mask is not None:
+            current_mask = torch.gather(edge_padding_mask, 1, current_index.repeat(1, 1, k_size)).to(
+                enhanced_node_feature.device)
+        else:
+            current_mask = None
+        current_mask[:, :, 0] = 1  # don't stay at current position
+
+        if not 0 in current_mask:
+            current_mask[:,:,0] = 0
+
+        current_mask = current_mask.permute(0, 2, 1)
+        zero = torch.zeros_like(q_values).to(q_values.device)
+        q_values = torch.where(current_mask == 1, zero, q_values)
+
+        return q_values, attention_weights
+
+    def forward(self, node_inputs, edge_inputs, current_index, node_padding_mask=None, edge_padding_mask=None,
+                edge_mask=None):
+        enhanced_node_feature = self.encode_graph(node_inputs, node_padding_mask, edge_mask)
+        q_values, attention_weights = self.output_q_values(enhanced_node_feature, edge_inputs, current_index, edge_padding_mask, node_padding_mask)
+        return q_values, attention_weights
+

planner/node.py

@@ -0,0 +1,96 @@
+#######################################################################
+# Name: node.py
+#
+# - Contains info per node on graph (edge)
+# - Contains: Position, Utility, Visitation History
+#######################################################################
+
+import sys
+if sys.modules['TRAINING']:
+    from .parameter import *
+else:
+    from .test_parameter import *
+
+import numpy as np
+import shapely.geometry
+
+
+class Node():
+    def __init__(self, coords, frontiers, robot_belief):
+        self.coords = coords
+        self.observable_frontiers = []
+        self.sensor_range = SENSOR_RANGE
+        self.initialize_observable_frontiers(frontiers, robot_belief)
+        self.utility = self.get_node_utility()
+        if self.utility == 0:
+            self.zero_utility_node = True
+        else:
+            self.zero_utility_node = False
+
+    def initialize_observable_frontiers(self, frontiers, robot_belief):
+        dist_list = np.linalg.norm(frontiers - self.coords, axis=-1)
+        frontiers_in_range = frontiers[dist_list < self.sensor_range - 10]
+        for point in frontiers_in_range:
+            collision = self.check_collision(self.coords, point, robot_belief)
+            if not collision:
+                self.observable_frontiers.append(point)
+
+    def get_node_utility(self):
+        return len(self.observable_frontiers)
+
+    def update_observable_frontiers(self, observed_frontiers, new_frontiers, robot_belief):
+        if observed_frontiers != []:
+            observed_index = []
+            for i, point in enumerate(self.observable_frontiers):
+                if point[0] + point[1] * 1j in observed_frontiers[:, 0] + observed_frontiers[:, 1] * 1j:
+                    observed_index.append(i)
+            for index in reversed(observed_index):
+                self.observable_frontiers.pop(index)
+        #
+        if new_frontiers != []:
+            dist_list = np.linalg.norm(new_frontiers - self.coords, axis=-1)
+            new_frontiers_in_range = new_frontiers[dist_list < self.sensor_range - 15]
+            for point in new_frontiers_in_range:
+                collision = self.check_collision(self.coords, point, robot_belief)
+                if not collision:
+                    self.observable_frontiers.append(point)
+
+        self.utility = self.get_node_utility()
+        if self.utility == 0:
+            self.zero_utility_node = True
+        else:
+            self.zero_utility_node = False
+
+    def set_visited(self):
+        self.observable_frontiers = []
+        self.utility = 0
+        self.zero_utility_node = True
+
+    def check_collision(self, start, end, robot_belief):
+        collision = False
+        line = shapely.geometry.LineString([start, end])
+
+        sortx = np.sort([start[0], end[0]])
+        sorty = np.sort([start[1], end[1]])
+
+        robot_belief = robot_belief[sorty[0]:sorty[1] + 1, sortx[0]:sortx[1] + 1]
+
+        occupied_area_index = np.where(robot_belief == 1)
+        occupied_area_coords = np.asarray(
+            [occupied_area_index[1] + sortx[0], occupied_area_index[0] + sorty[0]]).T
+        unexplored_area_index = np.where(robot_belief == 127)
+        unexplored_area_coords = np.asarray(
+            [unexplored_area_index[1] + sortx[0], unexplored_area_index[0] + sorty[0]]).T
+        unfree_area_coords = np.concatenate((occupied_area_coords, unexplored_area_coords))
+
+        for i in range(unfree_area_coords.shape[0]):
+            coords = ([(unfree_area_coords[i][0], unfree_area_coords[i][1]),
+                       (unfree_area_coords[i][0] + 1, unfree_area_coords[i][1]),
+                       (unfree_area_coords[i][0], unfree_area_coords[i][1] + 1),
+                       (unfree_area_coords[i][0] + 1, unfree_area_coords[i][1] + 1)])
+            obstacle = shapely.geometry.Polygon(coords)
+            collision = line.intersects(obstacle)
+            if collision:
+                break
+
+        return collision

planner/robot.py

@@ -0,0 +1,58 @@
+#######################################################################
+# Name: robot.py
+#
+# - Stores S(t), A(t), R(t), S(t+1)
+#######################################################################
+
+import torch
+from copy import deepcopy
+
+class Robot:
+    def __init__(self, robot_id, position, plot=False):
+        self.robot_id = robot_id
+        self.plot = plot
+        self.travel_dist = 0
+        self.robot_position = position
+        self.observations = None
+        self.trajectory_coords = []
+        self.targets_found_on_path = []
+        
+        self.episode_buffer = []
+        for i in range(15):
+            self.episode_buffer.append([])
+
+        if self.plot:
+            self.xPoints = [self.robot_position[0]]
+            self.yPoints = [self.robot_position[1]]
+
+    def save_observations(self, observations):
+        node_inputs, edge_inputs, current_index, node_padding_mask, edge_padding_mask, edge_mask = observations
+        self.episode_buffer[0] += deepcopy(node_inputs).to('cpu')
+        self.episode_buffer[1] += deepcopy(edge_inputs).to('cpu')
+        self.episode_buffer[2] += deepcopy(current_index).to('cpu')
+        self.episode_buffer[3] += deepcopy(node_padding_mask).to('cpu')
+        self.episode_buffer[4] += deepcopy(edge_padding_mask).to('cpu')
+        self.episode_buffer[5] += deepcopy(edge_mask).to('cpu')
+
+    def save_action(self, action_index):
+        self.episode_buffer[6] += action_index.unsqueeze(0).unsqueeze(0)
+
+    def save_reward_done(self, reward, done):
+        self.episode_buffer[7] += deepcopy(torch.FloatTensor([[[reward]]])).to('cpu')
+        self.episode_buffer[8] += deepcopy(torch.tensor([[[(int(done))]]])).to('cpu')
+        if self.plot:
+            self.xPoints.append(self.robot_position[0])
+            self.yPoints.append(self.robot_position[1])
+
+    def save_next_observations(self, observations):
+        node_inputs, edge_inputs, current_index, node_padding_mask, edge_padding_mask, edge_mask = observations
+        self.episode_buffer[9] += deepcopy(node_inputs).to('cpu')
+        self.episode_buffer[10] += deepcopy(edge_inputs).to('cpu')
+        self.episode_buffer[11] += deepcopy(current_index).to('cpu')
+        self.episode_buffer[12] += deepcopy(node_padding_mask).to('cpu')
+        self.episode_buffer[13] += deepcopy(edge_padding_mask).to('cpu')
+        self.episode_buffer[14] += deepcopy(edge_mask).to('cpu')
+
+    def save_trajectory_coords(self, robot_position_coords, num_target_found):
+        self.trajectory_coords.append(robot_position_coords)
+        self.targets_found_on_path.append(num_target_found)

planner/sensor.py

@@ -0,0 +1,128 @@
+#######################################################################
+# Name: sensor.py
+#
+# - Computes sensor related checks (e.g. collision, utility etc)
+#######################################################################
+
+import sys
+if sys.modules['TRAINING']:
+    from .parameter import *
+else:
+    from .test_parameter import *
+
+import math
+import numpy as np
+import copy
+
+def collision_check(x0, y0, x1, y1, ground_truth, robot_belief):
+    x0 = x0.round()
+    y0 = y0.round()
+    x1 = x1.round()
+    y1 = y1.round()
+    dx, dy = abs(x1 - x0), abs(y1 - y0)
+    x, y = x0, y0
+    error = dx - dy
+    x_inc = 1 if x1 > x0 else -1
+    y_inc = 1 if y1 > y0 else -1
+    dx *= 2
+    dy *= 2
+
+    collision_flag = 0
+    max_collision = 10
+
+    while 0 <= x < ground_truth.shape[1] and 0 <= y < ground_truth.shape[0]:
+        k = ground_truth.item(y, x)
+        if k == 1 and collision_flag < max_collision:
+            collision_flag += 1
+            if collision_flag >= max_collision:
+                break
+
+        if k !=1 and collision_flag > 0:
+            break
+
+        if x == x1 and y == y1:
+            break
+
+        robot_belief.itemset((y, x), k)
+
+        if error > 0:
+            x += x_inc
+            error -= dy
+        else:
+            y += y_inc
+            error += dx
+
+    return robot_belief
+
+
+def sensor_work(robot_position, sensor_range, robot_belief, ground_truth, sensor_model=SENSOR_MODEL):
+    x0 = robot_position[0]
+    y0 = robot_position[1]
+    rng_x = 0.5 * (ground_truth.shape[1] / NUM_COORDS_WIDTH)
+    rng_y = 0.5 * (ground_truth.shape[0] / NUM_COORDS_HEIGHT)
+
+    if sensor_model == "rectangular":   # TODO: add collision check 
+        max_x = min(x0 + int(math.ceil(rng_x)), ground_truth.shape[1])
+        min_x = max(x0 - int(math.ceil(rng_x)), 0)
+        max_y = min(y0 + int(math.ceil(rng_y)), ground_truth.shape[0])
+        min_y = max(y0 - int(math.ceil(rng_y)), 0)
+        robot_belief[min_y:max_y, min_x:max_x] = ground_truth[min_y:max_y, min_x:max_x]
+    else:
+        sensor_angle_inc = 0.5 / 180 * np.pi
+        sensor_angle = 0
+        while sensor_angle < 2 * np.pi:
+            x1 = x0 + np.cos(sensor_angle) * sensor_range
+            y1 = y0 + np.sin(sensor_angle) * sensor_range
+            robot_belief = collision_check(x0, y0, x1, y1, ground_truth, robot_belief)
+            sensor_angle += sensor_angle_inc
+    return robot_belief
+
+
+def unexplored_area_check(x0, y0, x1, y1, current_belief):
+    x0 = x0.round()
+    y0 = y0.round()
+    x1 = x1.round()
+    y1 = y1.round()
+    dx, dy = abs(x1 - x0), abs(y1 - y0)
+    x, y = x0, y0
+    error = dx - dy
+    x_inc = 1 if x1 > x0 else -1
+    y_inc = 1 if y1 > y0 else -1
+    dx *= 2
+    dy *= 2
+
+    while 0 <= x < current_belief.shape[1] and 0 <= y < current_belief.shape[0]:
+        k = current_belief.item(y, x)
+        if x == x1 and y == y1:
+            break
+
+        if k == 1:
+            break
+
+        if k == 127:
+            current_belief.itemset((y, x), 0)
+            break
+
+        if error > 0:
+            x += x_inc
+            error -= dy
+        else:
+            y += y_inc
+            error += dx
+
+    return current_belief
+
+
+def calculate_utility(waypoint_position, sensor_range, robot_belief):
+    sensor_angle_inc = 5 / 180 * np.pi
+    sensor_angle = 0
+    x0 = waypoint_position[0]
+    y0 = waypoint_position[1]
+    current_belief = copy.deepcopy(robot_belief)
+    while sensor_angle < 2 * np.pi:
+        x1 = x0 + np.cos(sensor_angle) * sensor_range
+        y1 = y0 + np.sin(sensor_angle) * sensor_range
+        current_belief = unexplored_area_check(x0, y0, x1, y1, current_belief)
+        sensor_angle += sensor_angle_inc
+    utility = np.sum(robot_belief == 127) - np.sum(current_belief == 127)
+    return utility

planner/test_info_surfing.py

@@ -0,0 +1,1071 @@
+#######################################################################
+# Name: test_info_surfing.py
+#
+# - Runs robot in environment using Info Surfing Planner
+#######################################################################
+
+import sys
+sys.modules['TRAINING'] = False           # False = Inference Testing   
+
+import copy
+import os
+import imageio
+import numpy as np
+import matplotlib.pyplot as plt
+from pathlib import Path
+from time import time
+from types import SimpleNamespace
+from skimage.transform import resize
+from taxabind_avs.satbind.kmeans_clustering import CombinedSilhouetteInertiaClusterer
+from .env import Env
+from .test_parameter import *
+
+
+OPPOSITE_ACTIONS = {1: 3, 2: 4, 3: 1, 4: 2, 5: 7, 6: 8, 7: 5, 8: 6}
+# color
+agentColor = (1, 0.2, 0.6)
+agentCommColor = (1, 0.6, 0.2)
+obstacleColor = (0., 0., 0.)
+targetNotFound = (0., 1., 0.)
+targetFound = (0.545, 0.27, 0.075)
+highestProbColor = (1., 0., 0.)
+highestUncertaintyColor = (0., 0., 1.)
+lowestProbColor = (1., 1., 1.)
+
+
+class ISEnv:
+    """Custom Environment that follows gym interface"""
+    metadata = {'render.modes': ['human']}
+
+    def __init__(self, global_step=0, state=None, shape=(24, 24), numAgents=8, observationSize=11, sensorSize=1, diag=False, save_image=False, clip_seg_tta=None):
+        
+        self.global_step = global_step
+        self.infoMap = None
+        self.targetMap = None
+        self.agents = []
+        self.targets = []
+        self.numAgents = numAgents
+        self.found_target = []
+        self.shape = shape
+        self.observationSize = observationSize
+        self.sensorSize = sensorSize
+        self.diag = diag
+        self.communicateCircle = 11
+        self.distribs = []
+        self.mask = None
+        self.finished = False
+        self.action_vects = [[-1., 0.], [0., 1.], [1., 0], [0., -1.]] if not diag else [[-1., 0.], [0., 1.], [1., 0], [0., -1.], [-0.707, -0.707], [-0.707, 0.707], [0.707, 0.707], [0.707, -0.707]]
+        self.actionlist = []
+        self.IS_step = 0
+        self.save_image = save_image
+        self.clip_seg_tta = clip_seg_tta
+        self.perf_metrics = dict()
+        self.steps_to_first_tgt = None
+        self.steps_to_mid_tgt = None
+        self.steps_to_last_tgt = None
+        self.targets_found_on_path = []
+        self.step_since_tta = 0
+        self.IS_frame_files = []
+        self.bad_mask_init = False
+
+        # define env
+        self.env = Env(map_index=self.global_step, n_agent=numAgents, k_size=K_SIZE, plot=save_image, test=True)
+    
+        # Overwrite state
+        if self.clip_seg_tta is not None:
+            self.clip_seg_tta.reset(sample_idx=self.global_step)
+
+            # Override target positions in env
+            self.env.target_positions = [(pose[1], pose[0]) for pose in self.clip_seg_tta.target_positions] 
+
+            # Override segmentation mask
+            if not USE_CLIP_PREDS and OVERRIDE_MASK_DIR != "":
+                score_mask_path = os.path.join(OVERRIDE_MASK_DIR, self.clip_seg_tta.gt_mask_name)
+                print("score_mask_path: ", score_mask_path)
+                if os.path.exists(score_mask_path):
+                    self.env.segmentation_mask = self.env.import_segmentation_mask(score_mask_path)
+                    self.env.begin(self.env.map_start_position)
+                else:
+                    print(f"\n\n{RED}ERROR: Trying to override, but score mask not found at path:{NC} ", score_mask_path)
+                    self.bad_mask_init = True
+
+        # Save clustered embeds from sat encoder
+        if USE_CLIP_PREDS:
+            self.kmeans_clusterer = CombinedSilhouetteInertiaClusterer(
+                k_min=1,
+                k_max=8,
+                k_avg_max=4,
+                silhouette_threshold=0.15,
+                relative_threshold=0.15,
+                random_state=0,
+                min_patch_size=5,   
+                n_smooth_iter=2,    
+                ignore_label=-1,
+                plot=self.save_image,
+                gifs_dir = GIFS_PATH      
+            )    
+            # Generate kmeans clusters
+            self.kmeans_sat_embeds_clusters = self.kmeans_clusterer.fit_predict(
+                patch_embeds=self.clip_seg_tta.patch_embeds,
+                map_shape=(CLIP_GRIDS_DIMS[0], CLIP_GRIDS_DIMS[1]),
+            )
+
+            if EXECUTE_TTA:
+                print("Will execute TTA...")
+
+        IS_info_map = copy.deepcopy(self.env.segmentation_info_mask)
+        IS_agent_loc = copy.deepcopy(self.env.start_positions)
+        IS_target_loc = copy.deepcopy(self.env.target_positions)
+        state=[IS_info_map, IS_agent_loc, IS_target_loc]
+        self.setWorld(state)
+
+
+    def init_render(self):
+        """
+        Call this once (e.g., in __init__ or just before the scenario loop) 
+        to initialize storage for agent paths and turn interactive plotting on.
+        """
+        # Keep track of each agent's trajectory
+        self.trajectories = [[] for _ in range(self.numAgents)]
+        self.trajectories_upscaled = [[] for _ in range(self.numAgents)]
+        
+        # Turn on interactive mode so we can update the same figure repeatedly
+        plt.ion()
+        plt.figure(figsize=(6,6))
+        plt.title("Information Map with Agents, Targets, and Sensor Ranges")
+
+
+    def record_positions(self):
+        """
+        Call this after all agents have moved in a step (or whenever you want to update 
+        the trajectory). It appends the current positions of each agent to `self.trajectories`.
+        """
+        for idx, agent in enumerate(self.agents):
+            self.trajectories[idx].append((agent.row, agent.col))
+            self.trajectories_upscaled[idx].append(self.env.graph_generator.grid_coords[agent.row, agent.col])
+
+
+    def render(self, episode_num, step_num):
+        """
+        Renders the current state in a single matplotlib plot.
+        Ensures consistent image size for GIF generation.
+        """
+        
+        # Completely reset the figure to avoid leftover state
+        plt.close('all')
+        fig = plt.figure(figsize=(6.4, 4.8), dpi=100)
+        ax = fig.add_subplot(111)
+
+        # Plot the information map
+        ax.imshow(self.infoMap, origin='lower', cmap='gray')
+
+        # Show agent positions and their trajectories
+        for idx, agent in enumerate(self.agents):
+            positions = self.trajectories[idx]
+            if len(positions) > 1:
+                rows = [p[0] for p in positions]
+                cols = [p[1] for p in positions]
+                ax.plot(cols, rows, linewidth=1)
+
+            ax.scatter(agent.col, agent.row, marker='o', s=50)
+
+        # Plot target locations
+        for t in self.targets:
+            color = 'green' if np.isnan(t.time_found) else 'red'
+            ax.scatter(t.col, t.row, marker='x', s=100, color=color)
+
+        # Title and axis formatting
+        ax.set_title(f"Step: {self.IS_step}")
+        ax.invert_yaxis()
+
+        # Create output folder if it doesn't exist
+        if not os.path.exists(GIFS_PATH):
+            os.makedirs(GIFS_PATH)
+
+        # Save the frame with consistent canvas
+        frame_path = f'{GIFS_PATH}/IS_{episode_num}_{step_num}.png'
+        plt.savefig(frame_path, bbox_inches='tight', pad_inches=0.1)
+        self.IS_frame_files.append(frame_path)
+
+        # Cleanup
+        plt.close(fig)
+
+
+    def setWorld(self, state=None):
+        """
+        1. empty all the element
+        2. create the new episode
+        """
+        if state is not None:
+            self.infoMap = copy.deepcopy(state[0].reshape(self.shape).T)
+            agents = []
+            self.numAgents = len(state[1])
+            for a in range(1, self.numAgents + 1):
+                abs_pos = state[1].pop(0)
+                abs_pos = np.array(abs_pos)
+                row, col = self.env.graph_generator.find_closest_index_from_grid_coords_2d(np.array(abs_pos))
+                agents.append(Agent(ID=a, row=row, col=col, sensorSize=self.sensorSize, infoMap=np.copy(self.infoMap),
+                                    uncertaintyMap=np.copy(self.infoMap), shape=self.shape, numAgents=self.numAgents))
+            self.agents = agents
+
+            targets, n_targets = [], 1
+            for t in range(len(state[2])):
+                abs_pos = state[2].pop(0)
+                abs_pos = np.array(abs_pos)
+                row, col = self.env.graph_generator.find_closest_index_from_grid_coords_2d(abs_pos)
+                targets.append(Target(ID=n_targets, row=row, col=col, time_found=np.nan))
+                n_targets = n_targets + 1
+            self.targets = targets
+
+    def extractObservation(self, agent):
+        """
+        Extract observations from information map
+        """
+
+        transform_row = self.observationSize // 2 - agent.row
+        transform_col = self.observationSize // 2 - agent.col
+
+        observation_layers = np.zeros((1, self.observationSize, self.observationSize))
+        min_row = max((agent.row - self.observationSize // 2), 0)
+        max_row = min((agent.row + self.observationSize // 2 + 1), self.shape[0])
+        min_col = max((agent.col - self.observationSize // 2), 0)
+        max_col = min((agent.col + self.observationSize // 2 + 1), self.shape[1])
+
+        observation = np.full((self.observationSize, self.observationSize), 0.)
+        infoMap = np.full((self.observationSize, self.observationSize), 0.)
+        densityMap = np.full((self.observationSize, self.observationSize), 0.)
+
+        infoMap[(min_row + transform_row):(max_row + transform_row),
+                    (min_col + transform_col):(max_col + transform_col)] = self.infoMap[
+                                                                           min_row:max_row, min_col:max_col]
+        observation_layers[0] = infoMap
+
+        return observation_layers
+
+
+    def listNextValidActions(self, agent_id, prev_action=0):
+        """
+        No movement: 0
+        North (-1,0): 1
+        East (0,1): 2
+        South (1,0): 3
+        West (0,-1): 4
+        """
+        available_actions = [0]
+        agent = self.agents[agent_id - 1]  
+
+        MOVES = [(-1, 0), (0, 1), (1, 0), (0, -1), (-1, -1), (-1, 1), (1, 1), (1, -1)]  
+        size = 4 + self.diag * 4
+        for action in range(size):
+            out_of_bounds = agent.row + MOVES[action][0] >= self.shape[0] \
+                            or agent.row + MOVES[action][0] < 0\
+                            or agent.col + MOVES[action][1] >= self.shape[1] \
+                            or agent.col + MOVES[action][1] < 0
+
+            if (not out_of_bounds) and not (prev_action == OPPOSITE_ACTIONS[action + 1]):
+                available_actions.append(action + 1)
+
+        return np.array(available_actions)
+
+
+    def executeAction(self, agentID, action, timeStep):
+        """
+        No movement: 0
+        North (-1,0): 1
+        East (0,1): 2
+        South (1,0): 3
+        West (0,-1): 4
+        LeftUp (-1,-1) : 5
+        RightUP (-1,1) :6
+        RightDown (1,1) :7
+        RightLeft (1,-1) :8
+        """
+        agent = self.agents[agentID - 1]
+        origLoc = agent.getLocation()
+
+        if (action >= 1) and (action <= 8):
+            agent.move(action)
+            row, col = agent.getLocation()
+
+            # If the move is not valid, roll it back
+            if (row < 0) or (col < 0) or (row >= self.shape[0]) or (col >= self.shape[1]):
+                self.updateInfoCheckTarget(agentID, timeStep, origLoc)
+                return 0
+
+        elif action == 0:
+            self.updateInfoCheckTarget(agentID, timeStep, origLoc)
+            return 0
+
+        else:
+            print("INVALID ACTION: {}".format(action))
+            sys.exit()
+
+        newLoc = agent.getLocation()
+        self.updateInfoCheckTarget(agentID, timeStep, origLoc)
+        return action
+
+
+    def updateInfoCheckTarget(self, agentID, timeStep, origLoc):
+        """
+        update the self.infoMap and check whether the agent has found a target
+        """
+        agent = self.agents[agentID - 1]
+        transform_row = self.sensorSize // 2 - agent.row
+        transform_col = self.sensorSize // 2 - agent.col
+
+        min_row = max((agent.row - self.sensorSize // 2), 0)
+        max_row = min((agent.row + self.sensorSize // 2 + 1), self.shape[0])
+        min_col = max((agent.col - self.sensorSize // 2), 0)
+        max_col = min((agent.col + self.sensorSize // 2 + 1), self.shape[1])
+        for t in self.targets:
+            if (t.row == agent.row) and (t.col == agent.col):
+                t.updateFound(timeStep)
+                self.found_target.append(t)
+                t.status = True
+
+        self.infoMap[min_row:max_row, min_col:max_col] *= 0.05  
+
+
+    def updateInfoEntireTrajectory(self, agentID):
+        """
+        update the self.infoMap and check whether the agent has found a target
+        """
+        traj = self.trajectories[agentID - 1]
+
+        for (row,col) in traj:
+            min_row = max((row - self.sensorSize // 2), 0)
+            max_row = min((row + self.sensorSize // 2 + 1), self.shape[0])
+            min_col = max((col - self.sensorSize // 2), 0)
+            max_col = min((col + self.sensorSize // 2 + 1), self.shape[1])
+            self.infoMap[min_row:max_row, min_col:max_col] *= 0.05  
+
+
+    # Execute one time step within the environment
+    def step(self, agentID, action, timeStep):
+        """
+        the agents execute the actions
+        No movement: 0
+        North (-1,0): 1
+        East (0,1): 2
+        South (1,0): 3
+        West (0,-1): 4
+        """
+        assert (agentID > 0)
+
+        self.executeAction(agentID, action, timeStep)
+
+
+    def observe(self, agentID):
+        assert (agentID > 0)
+        vectorObs = self.extractObservation(self.agents[agentID - 1])
+        return [vectorObs]
+
+
+    def check_finish(self):
+        if TERMINATE_ON_TGTS_FOUND:
+            found_status = [t.time_found for t in self.targets]
+            d = False
+            if np.isnan(found_status).sum() == 0:
+                d = True
+            return d
+        else:
+            return False    
+
+
+    def gradVec(self, observation, agent):
+        a = observation[0]  
+
+        # Make info & unc cells with low value as 0
+        a[a < 0.0002] = 0.0
+
+        # Center square from 11x11
+        a_11x11 = a[4:7, 4:7]
+        m_11x11 = np.array((a_11x11))
+
+        # Center square from 9x9
+        a_9x9 = self.pooling(a, (3, 3), stride=(1, 1), method='max', pad=False)    
+        a_9x9 = a_9x9[3:6, 3:6]
+        m_9x9 = np.array((a_9x9))
+
+        # Center square from 6x6
+        a_6x6 = self.pooling(a, (6, 6), stride=(1, 1), method='max', pad=False)    
+        a_6x6 = a_6x6[1:4, 1:4]
+        m_6x6 = np.array((a_6x6))
+
+        # Center square from 3x3
+        a_3x3 = self.pooling(a, (5, 5), stride=(3, 3), method='max', pad=False)
+        m_3x3 = np.array((a_3x3))
+
+        # Merging multiScales with weights
+        m = m_3x3 * 0.25 + m_6x6 * 0.25 + m_9x9 * 0.25 + m_11x11 * 0.25
+        a = m
+
+        adx, ady = np.gradient(a)
+        den = np.linalg.norm(np.array([adx[1, 1], ady[1, 1]]))
+        if (den != 0) and (not np.isnan(den)):
+            infovec = np.array([adx[1, 1], ady[1, 1]]) / den
+        else:
+            infovec = 0
+        agentvec = []
+
+        if len(agentvec) == 0:
+            den = np.linalg.norm(infovec)
+            if (den != 0) and (not np.isnan(den)):
+                direction = infovec / den
+            else:
+                direction = self.action_vects[np.random.randint(4 + self.diag * 4)]
+        else:
+            den = np.linalg.norm(np.mean(agentvec, 0))
+            if (den != 0) and (not np.isnan(den)):
+                agentvec = np.mean(agentvec, 0) / den
+            else:
+                agentvec = 0
+
+            den = np.linalg.norm(0.6 * infovec + 0.4 * agentvec)
+            if (den != 0) and (not np.isnan(den)):
+                direction = (0.6 * infovec + 0.4 * agentvec) / den
+            else:
+                direction = self.action_vects[np.random.randint(4 + self.diag * 4)]
+
+        action_vec = [[0, 0], [-1, 0], [0, 1], [1, 0], [0, -1]] if not self.diag else [[0., 0.], [-1., 0.], [0., 1.], [1., 0], [0., -1.], [-0.707, -0.707], [-0.707, 0.707], [0.707, 0.707], [0.707, -0.707]]
+        actionid = np.argmax([np.dot(direction, a) for a in action_vec])
+        actionid = self.best_valid_action(actionid, agent, direction)
+        return actionid
+
+
+    def best_valid_action(self, actionid, agent, direction):
+        if len(self.actionlist) > 1:
+            if self.action_invalid(actionid, agent):
+                action_vec = [[0, 0], [-1, 0], [0, 1], [1, 0], [0, -1]] if not self.diag else [[0., 0.], [-1., 0.], [0., 1.], [1., 0], [0., -1.], [-0.707, -0.707], [-0.707, 0.707], [0.707, 0.707], [0.707, -0.707]]
+                actionid = np.array([np.dot(direction, a) for a in action_vec])
+                actionid = actionid.argsort()
+                pi = 3 + self.diag*4 
+                while self.action_invalid(actionid[pi], agent) and pi >= 0:
+                    pi -= 1
+                if pi == -1:
+                    return OPPOSITE_ACTIONS[self.actionlist[self.IS_step - 1][agent - 1]]
+                elif actionid[pi] == 0:
+                    return OPPOSITE_ACTIONS[self.actionlist[self.IS_step - 1][agent - 1]]
+                else:
+                    return actionid[pi]
+        return actionid
+
+
+    def action_invalid(self, action, agent):
+        # Going back to the previous cell is disabled
+        if action == OPPOSITE_ACTIONS[self.actionlist[self.IS_step - 1][agent - 1]]:
+            return True
+        # Move N,E,S,W
+        if (action >= 1) and (action <= 8):
+            agent = self.agents[agent - 1]
+            agent.move(action)
+            row, col = agent.getLocation()
+
+            # If the move is not valid, roll it back
+            if ((row < 0) or (col < 0) or (row >= self.shape[0]) or (col >= self.shape[1])):
+                agent.reverseMove(action)
+                return True
+
+            agent.reverseMove(action)
+            return False
+        return False
+
+
+    def step_all_parallel(self):
+        actions = []
+        reward = 0
+        # Decide actions for each agent
+        for agent_id in range(1, self.numAgents + 1):
+            o = self.observe(agent_id)
+            actions.append(self.gradVec(o[0], agent_id))
+        self.actionlist.append(actions)
+
+        # Execute those actions
+        for agent_id in range(1, self.numAgents + 1):
+            self.step(agent_id, actions[agent_id - 1], self.IS_step)
+
+        # Record for visualization
+        self.record_positions()
+
+    def is_scenario(self, max_step=512, episode_number=0):
+
+        # Return all metrics as None if faulty mask init
+        if self.bad_mask_init:
+            self.perf_metrics['tax'] = None
+            self.perf_metrics['travel_dist'] = None
+            self.perf_metrics['travel_steps'] = None
+            self.perf_metrics['steps_to_first_tgt'] = None
+            self.perf_metrics['steps_to_mid_tgt'] = None
+            self.perf_metrics['steps_to_last_tgt'] = None
+            self.perf_metrics['explored_rate'] = None
+            self.perf_metrics['targets_found'] = None
+            self.perf_metrics['targets_total'] = None
+            self.perf_metrics['kmeans_k'] = None
+            self.perf_metrics['tgts_gt_score'] = None
+            self.perf_metrics['clip_inference_time'] = None
+            self.perf_metrics['tta_time'] = None
+            self.perf_metrics['success_rate'] = None
+            return
+
+        eps_start = time()
+        self.IS_step = 0
+        self.finished = False
+        reward = 0
+
+        # Initialize the rendering just once before the loop
+        self.init_render()
+        self.record_positions()
+
+        # Initial Setup
+        if LOAD_AVS_BENCH and USE_CLIP_PREDS:
+            if NUM_COORDS_WIDTH != CLIP_GRIDS_DIMS[0] or NUM_COORDS_HEIGHT != CLIP_GRIDS_DIMS[1]:   # If heatmap is resized from clip original dims
+                heatmap = self.convert_heatmap_resolution(self.clip_seg_tta.heatmap, full_dims=(512, 512), new_dims=(NUM_COORDS_WIDTH, NUM_COORDS_HEIGHT))
+                self.env.segmentation_info_mask = np.expand_dims(heatmap.T.flatten(), axis=1)
+                unnormalized_heatmap = self.convert_heatmap_resolution(self.clip_seg_tta.heatmap_unnormalized, full_dims=(512, 512), new_dims=(NUM_COORDS_WIDTH, NUM_COORDS_HEIGHT))
+                self.env.segmentation_info_mask_unnormalized = np.expand_dims(unnormalized_heatmap.T.flatten(), axis=1)
+                self.infoMap = copy.deepcopy(heatmap) 
+                print("Resized heatmap to", NUM_COORDS_WIDTH, "x", NUM_COORDS_HEIGHT)
+            else:   
+                self.env.segmentation_info_mask = np.expand_dims(self.clip_seg_tta.heatmap.T.flatten(), axis=1)
+                self.env.segmentation_info_mask_unnormalized = np.expand_dims(self.clip_seg_tta.heatmap_unnormalized.T.flatten(), axis=1)
+                self.infoMap = copy.deepcopy(self.clip_seg_tta.heatmap) 
+
+        self.targets_found_on_path.append(self.env.num_new_targets_found)
+
+        while self.IS_step < max_step and not self.check_finish():
+            self.step_all_parallel()
+            self.IS_step += 1
+
+            # Render after each step
+            if self.save_image:
+                self.render(episode_num=self.global_step, step_num=self.IS_step)
+
+            # Update in env
+            next_position_list = [self.trajectories_upscaled[i][-1] for i, agent in enumerate(self.agents)]
+            dist_list = [0 for _ in range(self.numAgents)]  
+            travel_dist_list = [self.compute_travel_distance(traj) for traj in self.trajectories]   
+            self.env.multi_robot_step(next_position_list, dist_list, travel_dist_list)
+            self.targets_found_on_path.append(self.env.num_new_targets_found)
+
+            # TTA Update via Poisson Test (with KMeans clustering stats)
+            robot_id = 0     # Assume 1 agent for now
+            robot_traj = self.trajectories[robot_id]
+            if LOAD_AVS_BENCH and USE_CLIP_PREDS and EXECUTE_TTA:  
+                flat_traj_coords = [robot_traj[i][1] * self.shape[0] + robot_traj[i][0] for i in range(len(robot_traj))]
+                robot = SimpleNamespace(
+                    trajectory_coords=flat_traj_coords,
+                    targets_found_on_path=self.targets_found_on_path
+                )
+                self.poisson_tta_update(robot, self.global_step, self.IS_step)
+                self.infoMap = copy.deepcopy(self.env.segmentation_info_mask.reshape((self.shape[1],self.shape[0])).T) 
+                self.updateInfoEntireTrajectory(robot_id)
+
+            # Update metrics
+            self.log_metrics(step=self.IS_step-1)
+
+            ### Save a frame to generate gif of robot trajectories ###
+            if self.save_image:
+                robots_route = [ ([], []) ] # Assume 1 robot
+                for point in self.trajectories_upscaled[robot_id]:
+                    robots_route[robot_id][0].append(point[0])
+                    robots_route[robot_id][1].append(point[1])
+                if not os.path.exists(GIFS_PATH):
+                    os.makedirs(GIFS_PATH)
+                if LOAD_AVS_BENCH:
+                    sound_id_override = None if self.clip_seg_tta.sound_ids == [] else self.clip_seg_tta.sound_ids[0]
+                    self.env.plot_env(
+                        self.global_step, 
+                        GIFS_PATH, 
+                        self.IS_step-1, 
+                        max(travel_dist_list),
+                        robots_route, 
+                        img_path_override=self.clip_seg_tta.img_paths[0],  # Viz 1st 
+                        sat_path_override=self.clip_seg_tta.imo_path,
+                        msk_name_override=self.clip_seg_tta.species_name, 
+                        sound_id_override=sound_id_override,
+                    )
+                else:
+                    self.env.plot_env(
+                        self.global_step, 
+                        GIFS_PATH, 
+                        self.IS_step-1, 
+                        max(travel_dist_list),
+                        robots_route
+                    )
+
+        # Log metrics
+        if LOAD_AVS_BENCH:
+            tax = Path(self.clip_seg_tta.gt_mask_name).stem
+            self.perf_metrics['tax'] = " ".join(tax.split("_")[1:])
+        else:
+            self.perf_metrics['tax'] = None
+        travel_distances = [self.compute_travel_distance(traj) for traj in self.trajectories]
+        self.perf_metrics['travel_dist'] = max(travel_distances)
+        self.perf_metrics['travel_steps'] = self.IS_step
+        self.perf_metrics['steps_to_first_tgt'] = self.steps_to_first_tgt
+        self.perf_metrics['steps_to_mid_tgt'] = self.steps_to_mid_tgt
+        self.perf_metrics['steps_to_last_tgt'] = self.steps_to_last_tgt
+        self.perf_metrics['explored_rate'] = self.env.explored_rate
+        self.perf_metrics['targets_found'] = self.env.targets_found_rate
+        self.perf_metrics['targets_total'] = len(self.env.target_positions)
+        if USE_CLIP_PREDS:
+            self.perf_metrics['kmeans_k'] = self.kmeans_clusterer.final_k
+            self.perf_metrics['tgts_gt_score'] = self.clip_seg_tta.tgts_gt_score
+            self.perf_metrics['clip_inference_time'] = self.clip_seg_tta.clip_inference_time
+            self.perf_metrics['tta_time'] = self.clip_seg_tta.tta_time
+        else:
+            self.perf_metrics['kmeans_k'] = None
+            self.perf_metrics['tgts_gt_score'] = None
+            self.perf_metrics['clip_inference_time'] = None
+            self.perf_metrics['tta_time'] = None
+        if FORCE_LOGGING_DONE_TGTS_FOUND and self.env.targets_found_rate == 1.0:
+            self.perf_metrics['success_rate'] = True
+        else:
+            self.perf_metrics['success_rate'] = self.env.check_done()[0]
+
+        # save gif
+        if self.save_image:
+            path = GIFS_PATH
+            self.make_gif(path, self.global_step)
+
+        print(YELLOW, f"[Eps {episode_number} Completed] Time Taken: {time()-eps_start:.2f}s, Steps: {self.IS_step}", NC)
+
+
+    def asStride(self, arr, sub_shape, stride):
+        """
+        Get a strided sub-matrices view of an ndarray.
+        See also skimage.util.shape.view_as_windows()
+        """
+        s0, s1 = arr.strides[:2]
+        m1, n1 = arr.shape[:2]
+        m2, n2 = sub_shape
+        view_shape = (1+(m1-m2)//stride[0], 1+(n1-n2)//stride[1], m2, n2)+arr.shape[2:]
+        strides = (stride[0]*s0, stride[1]*s1, s0, s1)+arr.strides[2:]
+        subs = np.lib.stride_tricks.as_strided(arr, view_shape, strides=strides)
+        return subs
+
+
+    def pooling(self, mat, ksize, stride=None, method='max', pad=False):
+        """
+        Overlapping pooling on 2D or 3D data.
+
+        <mat>: ndarray, input array to pool.
+        <ksize>: tuple of 2, kernel size in (ky, kx).
+        <stride>: tuple of 2 or None, stride of pooling window.
+                  If None, same as <ksize> (non-overlapping pooling).
+        <method>: str, 'max for max-pooling,
+                       'mean' for mean-pooling.
+        <pad>: bool, pad <mat> or not. If no pad, output has size
+               (n-f)//s+1, n being <mat> size, f being kernel size, s stride.
+               if pad, output has size ceil(n/s).
+
+        Return <result>: pooled matrix.
+        """
+
+        m, n = mat.shape[:2]
+        ky, kx = ksize
+        if stride is None:
+            stride = (ky, kx)
+        sy, sx = stride
+
+        _ceil = lambda x, y: int(np.ceil(x/float(y)))
+
+        if pad:
+            ny = _ceil(m,sy)
+            nx = _ceil(n,sx)
+            size = ((ny-1)*sy+ky, (nx-1)*sx+kx) + mat.shape[2:]
+            mat_pad = np.full(size,np.nan)
+            mat_pad[:m,:n,...] = mat
+        else:
+            mat_pad = mat[:(m-ky)//sy*sy+ky, :(n-kx)//sx*sx+kx, ...]
+
+        view = self.asStride(mat_pad,ksize,stride)
+
+        if method == 'max':
+            result = np.nanmax(view,axis=(2,3))
+        else:
+            result = np.nanmean(view,axis=(2,3))
+
+        return result
+
+
+    def compute_travel_distance(self, trajectory):
+        distance = 0.0
+        for i in range(1, len(trajectory)):
+            # Convert the tuple positions to numpy arrays for easy computation.
+            prev_pos = np.array(trajectory[i-1])
+            curr_pos = np.array(trajectory[i])
+            # Euclidean distance between consecutive positions.
+            distance += np.linalg.norm(curr_pos - prev_pos)
+        return distance
+
+    ################################################################################
+    # SPPP Related Fns
+    ################################################################################
+
+    def log_metrics(self, step):
+        # Update tgt found metrics
+        if self.steps_to_first_tgt is None and self.env.num_targets_found == 1:
+            self.steps_to_first_tgt = step + 1
+        if self.steps_to_mid_tgt is None and self.env.num_targets_found == int(len(self.env.target_positions) / 2):
+            self.steps_to_mid_tgt = step + 1
+        if self.steps_to_last_tgt is None and self.env.num_targets_found == len(self.env.target_positions):
+            self.steps_to_last_tgt = step + 1
+
+
+    def transpose_flat_idx(self, idx, H=NUM_COORDS_HEIGHT, W=NUM_COORDS_WIDTH):
+        """
+        Transpose a flat index from an ``H×W`` grid to the equivalent
+        position in the ``W×H`` transposed grid while **keeping the result
+        in 1-D**.
+        """
+        # --- Safety check to catch out-of-range indices ---
+        assert 0 <= idx < H * W, f"idx {idx} out of bounds for shape ({H}, {W})"
+
+        # Original (row, col)
+        row, col = divmod(idx, W)
+        # After transpose these coordinates swap
+        row_T, col_T = col, row
+
+        # Flatten back into 1-D (row-major) for the W×H grid
+        return row_T * H + col_T
+    
+
+    def poisson_tta_update(self, robot, episode, step):
+
+        # Generate Kmeans Clusters Stats
+        # Scale index back to CLIP_GRIDS_DIMS to be compatible with CLIP patch size
+        if NUM_COORDS_WIDTH != CLIP_GRIDS_DIMS[0] or NUM_COORDS_HEIGHT != CLIP_GRIDS_DIMS[1]:
+            # High-res remap via pixel coordinates preserves exact neighbourhood
+            filt_traj_coords, filt_targets_found_on_path = self.scale_trajectory(
+                robot.trajectory_coords,
+                self.env.target_positions,
+                old_dims=(NUM_COORDS_HEIGHT, NUM_COORDS_WIDTH),
+                full_dims=(512, 512),
+                new_dims=(CLIP_GRIDS_DIMS[0], CLIP_GRIDS_DIMS[1])
+            )
+        else:
+            filt_traj_coords = [self.transpose_flat_idx(idx) for idx in robot.trajectory_coords]
+            filt_targets_found_on_path = robot.targets_found_on_path
+
+        region_stats_dict = self.kmeans_clusterer.compute_region_statistics(
+            self.kmeans_sat_embeds_clusters,
+            self.clip_seg_tta.heatmap_unnormalized,
+            filt_traj_coords,
+            episode_num=episode,
+            step_num=step
+        )
+
+        # Prep & execute TTA
+        self.step_since_tta += 1
+        if robot.targets_found_on_path[-1] or self.step_since_tta % STEPS_PER_TTA == 0:  
+                
+            num_cells = self.clip_seg_tta.heatmap.shape[0] * self.clip_seg_tta.heatmap.shape[1]
+            pos_sample_weight_scale, neg_sample_weight_scale = [], []
+
+            for i, sample_loc in enumerate(filt_traj_coords):
+                label = self.kmeans_clusterer.get_label_id(self.kmeans_sat_embeds_clusters, sample_loc)
+                num_patches = region_stats_dict[label]['num_patches']
+                patches_visited = region_stats_dict[label]['patches_visited']
+                expectation = region_stats_dict[label]['expectation']
+
+                # Exponent like focal loss to wait for more samples before confidently decreasing
+                pos_weight = 4.0 
+                neg_weight = min(1.0, (patches_visited/(3*num_patches))**GAMMA_EXPONENT)  
+                pos_sample_weight_scale.append(pos_weight)
+                neg_sample_weight_scale.append(neg_weight)
+
+            # Adaptative LR (as samples increase, increase LR to fit more datapoints)
+            adaptive_lr = MIN_LR + (MAX_LR - MIN_LR) * (step / num_cells)
+                
+            # TTA Update
+            self.clip_seg_tta.execute_tta(
+                filt_traj_coords, 
+                filt_targets_found_on_path, 
+                tta_steps=NUM_TTA_STEPS, 
+                lr=adaptive_lr,
+                pos_sample_weight=pos_sample_weight_scale,
+                neg_sample_weight=neg_sample_weight_scale,
+                reset_weights=RESET_WEIGHTS
+            )
+            if NUM_COORDS_WIDTH != CLIP_GRIDS_DIMS[0] or NUM_COORDS_HEIGHT != CLIP_GRIDS_DIMS[1]:   # If heatmap is resized from clip original dims
+                heatmap = self.convert_heatmap_resolution(self.clip_seg_tta.heatmap, full_dims=(512, 512), new_dims=(NUM_COORDS_WIDTH, NUM_COORDS_HEIGHT))
+                self.env.segmentation_info_mask = np.expand_dims(heatmap.T.flatten(), axis=1)
+                unnormalized_heatmap = self.convert_heatmap_resolution(self.clip_seg_tta.heatmap_unnormalized, full_dims=(512, 512), new_dims=(NUM_COORDS_WIDTH, NUM_COORDS_HEIGHT))
+                self.env.segmentation_info_mask_unnormalized = np.expand_dims(unnormalized_heatmap.T.flatten(), axis=1)
+                print("~Resized heatmap to", NUM_COORDS_WIDTH, "x", NUM_COORDS_HEIGHT)
+            else:  
+                self.env.segmentation_info_mask = np.expand_dims(self.clip_seg_tta.heatmap.T.flatten(), axis=1)
+                self.env.segmentation_info_mask_unnormalized = np.expand_dims(self.clip_seg_tta.heatmap_unnormalized.T.flatten(), axis=1)
+
+            self.step_since_tta = 0
+
+
+    def convert_heatmap_resolution(self, heatmap, full_dims=(512, 512), new_dims=(24, 24)):
+        heatmap_large = resize(heatmap, full_dims, order=1,  # order=1 → bilinear
+             mode='reflect', anti_aliasing=True)
+        
+        coords = self.env.graph_generator.grid_coords   # (N, N, 2)
+        rows, cols = coords[...,1], coords[...,0]
+        heatmap_resized = heatmap_large[rows, cols]
+        heatmap_resized = heatmap_resized.reshape(new_dims[1], new_dims[0])
+        return heatmap_resized
+
+
+    def convert_labelmap_resolution(self, labelmap, full_dims=(512, 512), new_dims=(24, 24)):
+        """
+        1) Upsample via nearest‐neighbor to full_dims
+        2) Sample back down to your graph grid using grid_coords
+        """
+        # 1) Upsample with nearest‐neighbor, preserving integer labels
+        up = resize(
+            labelmap,
+            full_dims,
+            order=0,                # nearest‐neighbor
+            mode='edge',            # padding mode
+            preserve_range=True,    # don't normalize labels
+            anti_aliasing=False     # must be False for labels
+        ).astype(labelmap.dtype)    # back to original integer dtype
+
+        # 2) Downsample via your precomputed grid coords (N×N×2)
+        coords = self.env.graph_generator.grid_coords   # shape (N, N, 2)
+        rows = coords[...,1].astype(int)
+        cols = coords[...,0].astype(int)
+
+        small = up[rows, cols]      # shape (N, N)
+        small = small.reshape(new_dims[0], new_dims[1])
+        return small
+    
+
+    def scale_trajectory(self,
+                        flat_indices,
+                        targets,
+                        old_dims=(17, 17),
+                        full_dims=(512, 512),
+                        new_dims=(24, 24)):
+        """
+        Args:
+            flat_indices: list of ints in [0..old_H*old_W-1]
+            targets:      list of (y_pix, x_pix) in [0..full_H-1]
+            old_dims:     (old_H, old_W)
+            full_dims:    (full_H, full_W)
+            new_dims:     (new_H, new_W)
+
+        Returns:
+            new_flat_traj: list of unique flattened indices in new_H×new_W
+            counts:        list of ints, same length as new_flat_traj
+        """
+        old_H, old_W = old_dims
+        full_H, full_W = full_dims
+        new_H,  new_W  = new_dims
+
+        # 1) bin targets into new grid
+        cell_h_new = full_H / new_H
+        cell_w_new = full_W / new_W
+        grid_counts = [[0]*new_W for _ in range(new_H)]
+        for x_pix, y_pix in targets:  # note (x, y) order as in original implementation
+            i_t = min(int(y_pix / cell_h_new), new_H - 1)
+            j_t = min(int(x_pix / cell_w_new), new_W - 1)
+            grid_counts[i_t][j_t] += 1
+
+        # 2) Walk the trajectory indices and project each old cell's *entire
+        #    pixel footprint* onto the finer 24×24 grid.
+        cell_h_full = full_H / old_H
+        cell_w_full = full_W / old_W
+
+        seen = set()
+        new_flat_traj = []
+
+        for node_idx in flat_indices:
+            if node_idx < 0 or node_idx >= len(self.env.graph_generator.node_coords):
+                continue
+
+            coord_xy = self.env.graph_generator.node_coords[node_idx]
+            try:
+                row_old, col_old = self.env.graph_generator.find_index_from_grid_coords_2d(coord_xy)
+            except Exception:
+                continue
+
+            # Bounding box of the old cell in full-resolution pixel space
+            y0 = row_old * cell_h_full
+            y1 = (row_old + 1) * cell_h_full
+            x0 = col_old * cell_w_full
+            x1 = (col_old + 1) * cell_w_full
+
+            # Which new-grid rows & cols overlap? (inclusive ranges)
+            i_start = max(0, min(int(y0 / cell_h_new), new_H - 1))
+            i_end   = max(0, min(int((y1 - 1) / cell_h_new), new_H - 1))
+            j_start = max(0, min(int(x0 / cell_w_new), new_W - 1))
+            j_end   = max(0, min(int((x1 - 1) / cell_w_new), new_W - 1))
+
+            for ii in range(i_start, i_end + 1):
+                for jj in range(j_start, j_end + 1):
+                    f_new = ii * new_W + jj
+                    if f_new not in seen:
+                        seen.add(f_new)
+                        new_flat_traj.append(f_new)
+
+        # 3) annotate counts
+        counts = []
+        for f in new_flat_traj:
+            i_new, j_new = divmod(f, new_W)
+            counts.append(grid_counts[i_new][j_new])
+
+        return new_flat_traj, counts
+
+
+    ################################################################################
+
+    def make_gif(self, path, n):
+        """ Generate a gif given list of images """
+        with imageio.get_writer('{}/{}_target_rate_{:.2f}.gif'.format(path, n, self.env.targets_found_rate), mode='I',
+                                fps=5) as writer:
+            for frame in self.env.frame_files:
+                image = imageio.imread(frame)
+                writer.append_data(image)
+        print('gif complete\n')
+
+        # Remove files
+        for filename in self.env.frame_files[:-1]:
+            os.remove(filename)
+
+        # For KMeans gif 
+        if LOAD_AVS_BENCH and USE_CLIP_PREDS:
+            with imageio.get_writer('{}/{}_kmeans_stats.gif'.format(path, n), mode='I',
+                                    fps=5) as writer:
+                for frame in self.kmeans_clusterer.kmeans_frame_files:
+                    image = imageio.imread(frame)
+                    writer.append_data(image)
+            print('Kmeans Clusterer gif complete\n')
+
+            # Remove files
+            for filename in self.kmeans_clusterer.kmeans_frame_files[:-1]:
+                os.remove(filename)
+
+
+        # IS gif
+        with imageio.get_writer('{}/{}_IS.gif'.format(path, n), mode='I',
+                                fps=5) as writer:
+            for frame in self.IS_frame_files:
+                image = imageio.imread(frame)
+                writer.append_data(image)
+        print('Kmeans Clusterer gif complete\n')
+
+        # Remove files
+        for filename in self.IS_frame_files[:-1]:
+            os.remove(filename)
+
+    ################################################################################
+
+
+class Agent:
+    def __init__(self, ID, infoMap=None, uncertaintyMap=None, shape=None, row=0, col=0, sensorSize=9, numAgents=8):
+        self.ID = ID
+        self.row = row
+        self.col = col
+        self.numAgents = numAgents
+        self.sensorSize = sensorSize
+
+    def setLocation(self, row, col):
+        self.row = row
+        self.col = col
+
+    def getLocation(self):
+        return [self.row, self.col]
+
+    def move(self, action):
+        """
+        No movement: 0
+        North (-1,0): 1
+        East (0,1): 2
+        South (1,0): 3
+        West (0,-1): 4
+        LeftUp (-1,-1) : 5
+        RightUP (-1,1) :6
+        RightDown (1,1) :7
+        RightLeft (1,-1) :8
+        check valid action of the agent. be sure not to be out of the boundary
+        """
+        if action == 0:
+            return 0
+        elif action == 1:
+            self.row -= 1
+        elif action == 2:
+            self.col += 1
+        elif action == 3:
+            self.row += 1
+        elif action == 4:
+            self.col -= 1
+        elif action == 5:
+            self.row -= 1
+            self.col -= 1
+        elif action == 6:
+            self.row -= 1
+            self.col += 1
+        elif action == 7:
+            self.row += 1
+            self.col += 1
+        elif action == 8:
+            self.row += 1
+            self.col -= 1
+
+    def reverseMove(self, action):
+        if action == 0:
+            return 0
+        elif action == 1:
+            self.row += 1
+        elif action == 2:
+            self.col -= 1
+        elif action == 3:
+            self.row -= 1
+        elif action == 4:
+            self.col += 1
+        elif action == 5:
+            self.row += 1
+            self.col += 1
+        elif action == 6:
+            self.row += 1
+            self.col -= 1
+        elif action == 7:
+            self.row -= 1
+            self.col -= 1
+        elif action == 8:
+            self.row -= 1
+            self.col += 1
+        else:
+            print("agent can only move NESW/1234")
+            sys.exit()
+
+
+class Target:
+    def __init__(self, row, col, ID, time_found=np.nan):
+        self.row = row
+        self.col = col
+        self.ID = ID
+        self.time_found = time_found
+        self.status = None
+        self.time_visited = time_found
+
+    def getLocation(self):
+        return self.row, self.col
+
+    def updateFound(self, timeStep):
+        if np.isnan(self.time_found):
+            self.time_found = timeStep
+
+    def updateVisited(self, timeStep):
+        if np.isnan(self.time_visited):
+            self.time_visited = timeStep
+
+
+if __name__ == "__main__":
+
+    search_env = Env(map_index=1, k_size=K_SIZE, n_agent=NUM_ROBOTS, plot=SAVE_GIFS)
+
+    IS_info_map = search_env.segmentation_info_mask
+    IS_agent_loc = search_env.start_positions
+    IS_target_loc = [[312, 123], [123, 312], [312, 312], [123, 123]]
+
+    env = ISEnv(state=[IS_info_map, IS_agent_loc, IS_target_loc], shape=(NUM_COORDS_HEIGHT, NUM_COORDS_WIDTH))
+    env.is_scenario(NUM_EPS_STEPS)
+    print()

planner/test_parameter.py

@@ -0,0 +1,118 @@
+############################################################################################
+# Name: test_parameter.py
+#
+# NOTE: Change all your hyper-params here for eval
+# Simple How-To Guide: 
+# 1. CLIP TTA: USE_CLIP_PREDS=True, EXECUTE_TTA=True
+# 2. CLIP (No TTA): USE_CLIP_PREDS=True, EXECUTE_TTA=False 
+# 3. Custom masks (e.g. LISA): USE_CLIP_PREDS=False, EXECUTE_TTA=False 
+############################################################################################
+
+import os
+import sys
+sys.modules['TRAINING'] = False           # False = Inference Testing    
+
+###############################################################
+# Overload Params
+###############################################################
+
+OPT_VARS = {}
+def getenv(var_name, default=None, cast_type=str):
+    try:
+        value = os.environ.get(var_name, None)
+        if value is None:
+            result = default
+        elif cast_type == bool:
+            result = value.lower() in ("true", "1", "yes")
+        else:
+            result = cast_type(value)
+    except (ValueError, TypeError):
+        result = default
+
+    OPT_VARS[var_name] = result  # Log the result
+    return result
+
+###############################################################
+# General
+###############################################################
+
+# --- GENERAL --- #
+USE_GPU = False  
+NUM_GPU = getenv("NUM_GPU", default=1, cast_type=int) # the number of GPUs
+NUM_META_AGENT = getenv("NUM_META_AGENT", default=2, cast_type=int)  # the number of concurrent processes
+NUM_EPS_STEPS = getenv("NUM_EPS_STEPS", default=400, cast_type=int)
+FIX_START_POSITION = getenv("FIX_START_POSITION", default=True, cast_type=bool)  # Whether to fix the starting position of the robots (middle index)
+NUM_ROBOTS = 1                        # Only allow for 1 robot
+NUM_COORDS_WIDTH=24                   # How many node coords across width?
+NUM_COORDS_HEIGHT=24                  # How many node coords across height?
+CLIP_GRIDS_DIMS=[24,24]               # [16,16] if 'openai/clip-vit-large-patch14-336' 
+SENSOR_RANGE=80                       # Only applicable to 'circle' sensor model
+SENSOR_MODEL="rectangular"            # "rectangular", "circle" (NOTE: no colllision check for rectangular)
+TERMINATE_ON_TGTS_FOUND = True        # Whether to terminate episode when all targets found
+FORCE_LOGGING_DONE_TGTS_FOUND = True  # Whether to force csv logging when all targets found
+
+
+# --- Planner Params --- # 
+POLICY = getenv("POLICY", default="RL", cast_type=str)
+NUM_TEST = 800  # Overriden if LOAD_AVS_BENCH 
+NUM_RUN = 1
+MODEL_NAME = "avs_rl_policy.pth" 
+INPUT_DIM = 4
+EMBEDDING_DIM = 128
+K_SIZE = 8   
+
+
+# --- Folders & Visualizations --- #
+GRIDMAP_SET_DIR = "maps/gpt4o/envs_val"
+MASK_SET_DIR = "maps/example/masks_val"      # Overriden if LOAD_AVS_BENCH
+TARGETS_SET_DIR = ""
+# TARGETS_SET_DIR = "maps/example/gt_masks_val_with_tgts"   # Overriden if LOAD_AVS_BENCH
+OVERRIDE_MASK_DIR = getenv("OVERRIDE_MASK_DIR", default="", cast_type=str)  # Override initial score mask from CLIP 
+SAVE_GIFS = getenv("SAVE_GIFS", default=True, cast_type=bool)              # do you want to save GIFs
+FOLDER_NAME = 'avs_search' 
+MODEL_PATH = f'inference/model'
+GIFS_PATH = f'inference/test_results/gifs/{FOLDER_NAME}'
+LOG_PATH = f'inference/test_results/log/{FOLDER_NAME}'
+LOG_TEMPLATE_XLSX = f'inference/template.xlsx'
+CSV_EXPT_NAME = getenv("CSV_EXPT_NAME", default="data", cast_type=str)
+VIZ_GRAPH_EDGES = False  # do you want to visualize the graph edges
+
+
+#######################################################################
+# AVS Params
+#######################################################################
+
+# General PARAMS
+USE_CLIP_PREDS = getenv("USE_CLIP_PREDS", default=True, cast_type=bool)   # If false, use custom masks from OVERRIDE_MASK_DIR
+QUERY_TAX = getenv("QUERY_TAX", default="", cast_type=str)                # "" = Test all tax (can accept taxonomy substrings)
+EXECUTE_TTA = getenv("EXECUTE_TTA", default=True, cast_type=bool)         # Whether to execute TTA mask updates
+QUERY_MODALITY = getenv("QUERY_MODALITY", default="image", cast_type=str) # "image", "text", "sound"
+STEPS_PER_TTA = 20  # no. steps before each TTA series
+NUM_TTA_STEPS = 1   # no. of TTA steps during each series
+RESET_WEIGHTS = True
+MIN_LR = 1e-6
+MAX_LR = 1e-5   
+GAMMA_EXPONENT = 2  
+
+# Paths related to AVS (TRAIN w/ TARGETS)
+LOAD_AVS_BENCH = True     # Whether to init AVS datasets 
+AVS_IMG_DIR = '/mnt/hdd/avs_bench_ds/inat21' 
+AVS_IMO_DIR = '/mnt/hdd/avs_bench_ds/sat_jpg/train_512px' 
+AVS_INAT_JSON_PATH = '/mnt/hdd/avs_bench_ds/inat21/train.json' 
+AVS_SOUND_DIR = '/mnt/hdd/avs_bench_ds/sound_mp3/test'
+AVS_GAUSSIAN_BLUR_KERNEL = (5,5)
+AVS_SAT_TO_IMG_IDS_PATH = getenv("AVS_SAT_TO_IMG_IDS_PATH", default="search_tri_modal|val_in_domain", cast_type=str)
+AVS_LOAD_PRETRAINED_HF_CHECKPOINT = getenv("AVS_LOAD_PRETRAINED_HF_CHECKPOINT", default=True, cast_type=bool)  # If false, load locally using CHECKPOINT_PATHs
+AVS_SAT_CHECKPOINT_PATH = getenv("AVS_SAT_CHECKPOINT_PATH", default="", cast_type=str)                  
+AVS_SOUND_CHECKPOINT_PATH = getenv("AVS_SOUND_CHECKPOINT_PATH", default="", cast_type=str)   
+
+#######################################################################
+# UTILS
+#######################################################################
+
+# COLORS (for printing)
+RED='\033[1;31m'          
+GREEN='\033[1;32m'
+YELLOW='\033[1;93m'       
+NC_BOLD='\033[1m' # Bold, No Color 
+NC='\033[0m'      # No Color 

planner/test_worker.py

@@ -0,0 +1,590 @@
+#######################################################################
+# Name: test_worker.py
+#
+# - Runs robot in environment using RL Planner
+#######################################################################
+
+from .test_parameter import *
+
+import imageio
+import os
+import copy
+import numpy as np
+import torch
+from time import time
+from pathlib import Path
+from skimage.transform import resize
+from taxabind_avs.satbind.kmeans_clustering import CombinedSilhouetteInertiaClusterer
+from .env import Env
+from .robot import Robot
+
+np.seterr(invalid='raise', divide='raise')
+
+
+class TestWorker:
+    def __init__(self, meta_agent_id, n_agent, policy_net, global_step, device='cuda', greedy=False, save_image=False, clip_seg_tta=None):
+        self.device = device
+        self.greedy = greedy
+        self.n_agent = n_agent
+        self.metaAgentID = meta_agent_id
+        self.global_step = global_step
+        self.k_size = K_SIZE
+        self.save_image = save_image
+        self.clip_seg_tta = clip_seg_tta
+        self.execute_tta = EXECUTE_TTA      # Added to interface with app.py
+
+        self.env = Env(map_index=self.global_step, n_agent=n_agent, k_size=self.k_size, plot=save_image, test=True)
+        self.local_policy_net = policy_net
+
+        self.robot_list = []
+        self.all_robot_positions = []
+        for i in range(self.n_agent):
+            robot_position = self.env.start_positions[i]
+            robot = Robot(robot_id=i, position=robot_position, plot=save_image)
+            self.robot_list.append(robot)
+            self.all_robot_positions.append(robot_position)
+
+        self.perf_metrics = dict()
+        self.bad_mask_init = False
+
+        # NOTE: Option to override gifs_path to interface with app.py
+        self.gifs_path = GIFS_PATH
+
+        # NOTE: updated due to app.py (hf does not allow heatmap to persist)
+        if LOAD_AVS_BENCH: 
+            if clip_seg_tta is not None: 
+                heatmap, heatmap_unnormalized, heatmap_unnormalized_initial, patch_embeds = self.clip_seg_tta.reset(sample_idx=self.global_step)
+                self.clip_seg_tta.heatmap = heatmap
+                self.clip_seg_tta.heatmap_unnormalized = heatmap_unnormalized
+                self.clip_seg_tta.heatmap_unnormalized_initial = heatmap_unnormalized_initial
+                self.clip_seg_tta.patch_embeds = patch_embeds   
+
+            # Override target positions in env
+            self.env.target_positions = [(pose[1], pose[0]) for pose in self.clip_seg_tta.target_positions] 
+
+            # Override segmentation mask
+            if not USE_CLIP_PREDS and OVERRIDE_MASK_DIR != "":
+                score_mask_path = os.path.join(OVERRIDE_MASK_DIR, self.clip_seg_tta.gt_mask_name)
+                print("score_mask_path: ", score_mask_path)
+                if os.path.exists(score_mask_path):
+                    self.env.segmentation_mask = self.env.import_segmentation_mask(score_mask_path)
+                    self.env.begin(self.env.map_start_position)
+                else:
+                    print(f"\n\n{RED}ERROR: Trying to override, but score mask not found at path:{NC} ", score_mask_path)
+                    self.bad_mask_init = True
+
+            # Save clustered embeds from sat encoder
+            if USE_CLIP_PREDS:
+                self.kmeans_clusterer = CombinedSilhouetteInertiaClusterer(
+                    k_min=1,
+                    k_max=8,
+                    k_avg_max=4,
+                    silhouette_threshold=0.15,
+                    relative_threshold=0.15,
+                    random_state=0,
+                    min_patch_size=5,   
+                    n_smooth_iter=2,    
+                    ignore_label=-1,
+                    plot=self.save_image,
+                    gifs_dir = GIFS_PATH      
+                )    
+                # Generate kmeans clusters
+                self.kmeans_sat_embeds_clusters = self.kmeans_clusterer.fit_predict(
+                    patch_embeds=self.clip_seg_tta.patch_embeds,
+                    map_shape=(CLIP_GRIDS_DIMS[0], CLIP_GRIDS_DIMS[1]),
+                )
+                print("Chosen k:", self.kmeans_clusterer.final_k)
+
+                # if EXECUTE_TTA:
+                #     print("Will execute TTA...")
+
+        # Define Poisson TTA params
+        self.step_since_tta = 0
+        self.steps_to_first_tgt = None
+        self.steps_to_mid_tgt = None
+        self.steps_to_last_tgt = None
+
+
+    def run_episode(self, curr_episode):
+
+        # Return all metrics as None if faulty mask init
+        if self.bad_mask_init:
+            self.perf_metrics['tax'] = None
+            self.perf_metrics['travel_dist'] = None
+            self.perf_metrics['travel_steps'] = None
+            self.perf_metrics['steps_to_first_tgt'] = None
+            self.perf_metrics['steps_to_mid_tgt'] = None
+            self.perf_metrics['steps_to_last_tgt'] = None
+            self.perf_metrics['explored_rate'] = None
+            self.perf_metrics['targets_found'] = None
+            self.perf_metrics['targets_total'] = None
+            self.perf_metrics['kmeans_k'] = None
+            self.perf_metrics['tgts_gt_score'] = None
+            self.perf_metrics['clip_inference_time'] = None
+            self.perf_metrics['tta_time'] = None
+            self.perf_metrics['success_rate'] = None
+            return
+
+        eps_start = time()
+        done = False
+        for robot_id, deciding_robot in enumerate(self.robot_list):
+            deciding_robot.observations = self.get_observations(deciding_robot.robot_position)
+            if LOAD_AVS_BENCH and USE_CLIP_PREDS:
+                if NUM_COORDS_WIDTH != CLIP_GRIDS_DIMS[0] or NUM_COORDS_HEIGHT != CLIP_GRIDS_DIMS[1]:   # If heatmap is resized from clip original dims
+                    heatmap = self.convert_heatmap_resolution(self.clip_seg_tta.heatmap, full_dims=(512, 512), new_dims=(NUM_COORDS_WIDTH, NUM_COORDS_HEIGHT))
+                    self.env.segmentation_info_mask = np.expand_dims(heatmap.T.flatten(), axis=1)
+                    unnormalized_heatmap = self.convert_heatmap_resolution(self.clip_seg_tta.heatmap_unnormalized, full_dims=(512, 512), new_dims=(NUM_COORDS_WIDTH, NUM_COORDS_HEIGHT))
+                    self.env.segmentation_info_mask_unnormalized = np.expand_dims(unnormalized_heatmap.T.flatten(), axis=1)
+                    print("Resized heatmap to", NUM_COORDS_WIDTH, "x", NUM_COORDS_HEIGHT)
+                else:   
+                    self.env.segmentation_info_mask = np.expand_dims(self.clip_seg_tta.heatmap.T.flatten(), axis=1)
+                    self.env.segmentation_info_mask_unnormalized = np.expand_dims(self.clip_seg_tta.heatmap_unnormalized.T.flatten(), axis=1)
+
+        ### Run episode ###
+        for step in range(NUM_EPS_STEPS):
+
+            next_position_list = []
+            dist_list = []
+            travel_dist_list = []
+            dist_array = np.zeros((self.n_agent, 1))
+            for robot_id, deciding_robot in enumerate(self.robot_list):
+                observations = deciding_robot.observations
+
+                ### Forward pass through policy to get next position ###
+                next_position, action_index = self.select_node(observations)
+                dist = np.linalg.norm(next_position - deciding_robot.robot_position)
+
+                ### Log results of action (e.g. distance travelled) ###
+                dist_array[robot_id] = dist
+                dist_list.append(dist)
+                travel_dist_list.append(deciding_robot.travel_dist)
+                next_position_list.append(next_position)
+                self.all_robot_positions[robot_id] = next_position
+
+            arriving_sequence = np.argsort(dist_list)
+            next_position_list = np.array(next_position_list)
+            dist_list = np.array(dist_list)
+            travel_dist_list = np.array(travel_dist_list)
+            next_position_list = next_position_list[arriving_sequence]
+            dist_list = dist_list[arriving_sequence]
+            travel_dist_list = travel_dist_list[arriving_sequence]
+
+            ### Take Action (Deconflict if 2 agents choose the same target position) ###
+            next_position_list, dist_list = self.solve_conflict(arriving_sequence, next_position_list, dist_list)
+            reward_list, done = self.env.multi_robot_step(next_position_list, dist_list, travel_dist_list)
+
+            ### Update observations + rewards from action ###
+            for reward, robot_id in zip(reward_list, arriving_sequence):
+                robot = self.robot_list[robot_id]
+                robot.save_trajectory_coords(self.env.find_index_from_coords(robot.robot_position), self.env.num_new_targets_found)
+
+                # # TTA Update via Poisson Test (with KMeans clustering stats)
+                if LOAD_AVS_BENCH and USE_CLIP_PREDS and self.execute_tta:  
+                    self.poisson_tta_update(robot, self.global_step, step)
+
+                robot.observations = self.get_observations(robot.robot_position)
+                robot.save_reward_done(reward, done)
+
+                # Update metrics
+                self.log_metrics(step=step) 
+
+            ### Save a frame to generate gif of robot trajectories ###
+            if self.save_image:
+                robots_route = []
+                for robot in self.robot_list:
+                    robots_route.append([robot.xPoints, robot.yPoints])
+                if not os.path.exists(self.gifs_path):
+                    os.makedirs(self.gifs_path)
+                if LOAD_AVS_BENCH:
+                    # NOTE: Replaced since using app.py
+                    self.env.plot_heatmap(self.gifs_path, step, max(travel_dist_list), robots_route)
+
+            if done:
+                break
+
+        if LOAD_AVS_BENCH:
+            tax = Path(self.clip_seg_tta.gt_mask_name).stem
+            self.perf_metrics['tax'] = " ".join(tax.split("_")[1:])
+        else:
+            self.perf_metrics['tax'] = None
+        self.perf_metrics['travel_dist'] = max(travel_dist_list)
+        self.perf_metrics['travel_steps'] = step + 1
+        self.perf_metrics['steps_to_first_tgt'] = self.steps_to_first_tgt
+        self.perf_metrics['steps_to_mid_tgt'] = self.steps_to_mid_tgt
+        self.perf_metrics['steps_to_last_tgt'] = self.steps_to_last_tgt
+        self.perf_metrics['explored_rate'] = self.env.explored_rate
+        self.perf_metrics['targets_found'] = self.env.targets_found_rate
+        self.perf_metrics['targets_total'] = len(self.env.target_positions)
+        if USE_CLIP_PREDS:
+            self.perf_metrics['kmeans_k'] = self.kmeans_clusterer.final_k
+            self.perf_metrics['tgts_gt_score'] = self.clip_seg_tta.tgts_gt_score
+            self.perf_metrics['clip_inference_time'] = self.clip_seg_tta.clip_inference_time
+            self.perf_metrics['tta_time'] = self.clip_seg_tta.tta_time
+        else:
+            self.perf_metrics['kmeans_k'] = None
+            self.perf_metrics['tgts_gt_score'] = None
+            self.perf_metrics['clip_inference_time'] = None
+            self.perf_metrics['tta_time'] = None
+        if FORCE_LOGGING_DONE_TGTS_FOUND and self.env.targets_found_rate == 1.0:
+            self.perf_metrics['success_rate'] = True
+        else:
+            self.perf_metrics['success_rate'] = done
+
+        # save gif
+        if self.save_image:
+            path = self.gifs_path   # NOTE: Set to self.gifs_path since using app.py
+            self.make_gif(path, curr_episode)
+
+        print(YELLOW, f"[Eps {curr_episode} Completed] Time Taken: {time()-eps_start:.2f}s, Steps: {step+1}", NC)
+
+    def get_observations(self, robot_position):
+        """ Get robot's sensor observation of environment given position """
+        current_node_index = self.env.find_index_from_coords(robot_position)
+        current_index = torch.tensor([current_node_index]).unsqueeze(0).unsqueeze(0).to(self.device)  # (1,1,1)
+    
+        node_coords = copy.deepcopy(self.env.node_coords)
+        graph = copy.deepcopy(self.env.graph)
+        node_utility = copy.deepcopy(self.env.node_utility)
+        guidepost = copy.deepcopy(self.env.guidepost)
+        segmentation_info_mask = copy.deepcopy(self.env.filtered_seg_info_mask)
+
+        n_nodes = node_coords.shape[0]
+        node_coords = node_coords / 640
+        node_utility = node_utility / 50
+        node_utility_inputs = node_utility.reshape((n_nodes, 1))
+
+        occupied_node = np.zeros((n_nodes, 1))
+        for position in self.all_robot_positions:
+            index = self.env.find_index_from_coords(position)
+            if index == current_index.item():
+                occupied_node[index] = -1
+            else:
+                occupied_node[index] = 1
+
+        node_inputs = np.concatenate((node_coords, segmentation_info_mask, guidepost), axis=1)
+        node_inputs = torch.FloatTensor(node_inputs).unsqueeze(0).to(self.device)  
+        node_padding_mask = None
+
+        graph = list(graph.values())
+        edge_inputs = []
+        for node in graph:
+            node_edges = list(map(int, node))
+            edge_inputs.append(node_edges)
+
+        bias_matrix = self.calculate_edge_mask(edge_inputs)
+        edge_mask = torch.from_numpy(bias_matrix).float().unsqueeze(0).to(self.device)
+
+        for edges in edge_inputs:
+            while len(edges) < self.k_size:
+                edges.append(0)        
+
+        edge_inputs = torch.tensor(edge_inputs).unsqueeze(0).to(self.device)  
+        edge_padding_mask = torch.zeros((1, len(edge_inputs), K_SIZE), dtype=torch.int64).to(self.device)
+        one = torch.ones_like(edge_padding_mask, dtype=torch.int64).to(self.device)
+        edge_padding_mask = torch.where(edge_inputs == 0, one, edge_padding_mask)
+
+        observations = node_inputs, edge_inputs, current_index, node_padding_mask, edge_padding_mask, edge_mask
+        return observations
+    
+
+    def select_node(self, observations):
+        """ Forward pass through policy to get next position to go to on map """
+        node_inputs, edge_inputs, current_index, node_padding_mask, edge_padding_mask, edge_mask = observations
+        with torch.no_grad():
+            logp_list = self.local_policy_net(node_inputs, edge_inputs, current_index, node_padding_mask, edge_padding_mask, edge_mask)
+
+        if self.greedy:
+            action_index = torch.argmax(logp_list, dim=1).long()
+        else:
+            action_index = torch.multinomial(logp_list.exp(), 1).long().squeeze(1)
+
+        next_node_index = edge_inputs[:, current_index.item(), action_index.item()]
+
+        next_position = self.env.node_coords[next_node_index]
+
+        return next_position, action_index
+
+    def solve_conflict(self, arriving_sequence, next_position_list, dist_list):
+        """ Deconflict if 2 agents choose the same target position """
+        for j, [robot_id, next_position] in enumerate(zip(arriving_sequence, next_position_list)):
+            moving_robot = self.robot_list[robot_id]
+            # if next_position[0] + next_position[1] * 1j in (next_position_list[:, 0] + next_position_list[:, 1] * 1j)[:j]:
+            #     dist_to_next_position = np.argsort(np.linalg.norm(self.env.node_coords - next_position, axis=1))
+            #     k = 0
+            #     while next_position[0] + next_position[1] * 1j in (next_position_list[:, 0] + next_position_list[:, 1] * 1j)[:j]:
+            #         k += 1
+            #         next_position = self.env.node_coords[dist_to_next_position[k]]
+
+            dist = np.linalg.norm(next_position - moving_robot.robot_position)
+            next_position_list[j] = next_position
+            dist_list[j] = dist
+            moving_robot.travel_dist += dist
+            moving_robot.robot_position = next_position
+
+        return next_position_list, dist_list
+
+    def work(self, currEpisode):
+        '''
+        Interacts with the environment. The agent gets either gradients or experience buffer
+        '''
+        self.run_episode(currEpisode)
+
+    def calculate_edge_mask(self, edge_inputs):
+        size = len(edge_inputs)
+        bias_matrix = np.ones((size, size))
+        for i in range(size):
+            for j in range(size):
+                if j in edge_inputs[i]:
+                    bias_matrix[i][j] = 0
+        return bias_matrix
+
+    def make_gif(self, path, n):
+        """ Generate a gif given list of images """
+        with imageio.get_writer('{}/{}_target_rate_{:.2f}.gif'.format(path, n, self.env.targets_found_rate), mode='I',
+                                fps=5) as writer:
+            for frame in self.env.frame_files:
+                image = imageio.imread(frame)
+                writer.append_data(image)
+        print('gif complete\n')
+
+        # Remove files
+        for filename in self.env.frame_files[:-1]:
+            os.remove(filename)
+
+        # For gif during TTA
+        if LOAD_AVS_BENCH:
+            with imageio.get_writer('{}/{}_kmeans_stats.gif'.format(path, n), mode='I',
+                                    fps=5) as writer:
+                for frame in self.kmeans_clusterer.kmeans_frame_files:
+                    image = imageio.imread(frame)
+                    writer.append_data(image)
+            print('Kmeans Clusterer gif complete\n')
+
+            # Remove files
+            for filename in self.kmeans_clusterer.kmeans_frame_files[:-1]:
+                os.remove(filename)
+
+    ################################################################################
+    # SPPP Related Fns
+    ################################################################################
+
+    def log_metrics(self, step):
+        # Update tgt found metrics
+        if self.steps_to_first_tgt is None and self.env.num_targets_found == 1:
+            self.steps_to_first_tgt = step + 1
+        if self.steps_to_mid_tgt is None and self.env.num_targets_found == int(len(self.env.target_positions) / 2):
+            self.steps_to_mid_tgt = step + 1
+        if self.steps_to_last_tgt is None and self.env.num_targets_found == len(self.env.target_positions):
+            self.steps_to_last_tgt = step + 1
+
+
+    def transpose_flat_idx(self, idx, H=NUM_COORDS_HEIGHT, W=NUM_COORDS_WIDTH):
+        """
+        Transpose a flat index from an ``H×W`` grid to the equivalent
+        position in the ``W×H`` transposed grid while **keeping the result
+        in 1-D**.
+        """
+        # --- Safety check to catch out-of-range indices ---
+        assert 0 <= idx < H * W, f"idx {idx} out of bounds for shape ({H}, {W})"
+
+        # Original (row, col)
+        row, col = divmod(idx, W)
+        # After transpose these coordinates swap
+        row_T, col_T = col, row
+
+        # Flatten back into 1-D (row-major) for the W×H grid
+        return row_T * H + col_T
+    
+
+    def poisson_tta_update(self, robot, episode, step):
+
+        # Generate Kmeans Clusters Stats
+        # Scale index back to CLIP_GRIDS_DIMS to be compatible with CLIP patch size
+        if NUM_COORDS_WIDTH != CLIP_GRIDS_DIMS[0] or NUM_COORDS_HEIGHT != CLIP_GRIDS_DIMS[1]:
+            # High-res remap via pixel coordinates preserves exact neighbourhood
+            filt_traj_coords, filt_targets_found_on_path = self.scale_trajectory(
+                robot.trajectory_coords,
+                self.env.target_positions,
+                old_dims=(NUM_COORDS_HEIGHT, NUM_COORDS_WIDTH),
+                full_dims=(512, 512),
+                new_dims=(CLIP_GRIDS_DIMS[0], CLIP_GRIDS_DIMS[1])
+            )
+        else:
+            filt_traj_coords = [self.transpose_flat_idx(idx) for idx in robot.trajectory_coords]
+            filt_targets_found_on_path = robot.targets_found_on_path
+
+        region_stats_dict = self.kmeans_clusterer.compute_region_statistics(
+            self.kmeans_sat_embeds_clusters,
+            self.clip_seg_tta.heatmap_unnormalized,
+            filt_traj_coords,
+            episode_num=episode,
+            step_num=step
+        )
+
+        # Prep & execute TTA
+        self.step_since_tta += 1
+        if robot.targets_found_on_path[-1] or self.step_since_tta % STEPS_PER_TTA == 0:    
+
+            # NOTE: integration with app.py on hf
+            self.clip_seg_tta.executing_tta = True
+
+            num_cells = self.clip_seg_tta.heatmap.shape[0] * self.clip_seg_tta.heatmap.shape[1]
+            pos_sample_weight_scale, neg_sample_weight_scale = [], []
+
+            for i, sample_loc in enumerate(filt_traj_coords):
+                label = self.kmeans_clusterer.get_label_id(self.kmeans_sat_embeds_clusters, sample_loc)
+                num_patches = region_stats_dict[label]['num_patches']
+                patches_visited = region_stats_dict[label]['patches_visited']
+                expectation = region_stats_dict[label]['expectation']
+
+                # Exponent like focal loss to wait for more samples before confidently decreasing
+                pos_weight = 4.0 
+                neg_weight = min(1.0, (patches_visited/(3*num_patches))**GAMMA_EXPONENT)  
+                pos_sample_weight_scale.append(pos_weight)
+                neg_sample_weight_scale.append(neg_weight)
+
+            # # # Adaptative LR (as samples increase, increase LR to fit more datapoints)
+            adaptive_lr = MIN_LR + (MAX_LR - MIN_LR) * (step / num_cells)
+                
+            # TTA Update
+            # NOTE: updated due to app.py (hf does not allow heatmap to persist)
+            heatmap = self.clip_seg_tta.execute_tta(
+                filt_traj_coords, 
+                filt_targets_found_on_path, 
+                tta_steps=NUM_TTA_STEPS, 
+                lr=adaptive_lr,
+                pos_sample_weight=pos_sample_weight_scale,
+                neg_sample_weight=neg_sample_weight_scale,
+                reset_weights=RESET_WEIGHTS
+            )
+            self.clip_seg_tta.heatmap = heatmap
+
+            if NUM_COORDS_WIDTH != CLIP_GRIDS_DIMS[0] or NUM_COORDS_HEIGHT != CLIP_GRIDS_DIMS[1]:   # If heatmap is resized from clip original dims
+                heatmap = self.convert_heatmap_resolution(self.clip_seg_tta.heatmap, full_dims=(512, 512), new_dims=(NUM_COORDS_WIDTH, NUM_COORDS_HEIGHT))
+                self.env.segmentation_info_mask = np.expand_dims(heatmap.T.flatten(), axis=1)
+                unnormalized_heatmap = self.convert_heatmap_resolution(self.clip_seg_tta.heatmap_unnormalized, full_dims=(512, 512), new_dims=(NUM_COORDS_WIDTH, NUM_COORDS_HEIGHT))
+                self.env.segmentation_info_mask_unnormalized = np.expand_dims(unnormalized_heatmap.T.flatten(), axis=1)
+                print("~Resized heatmap to", NUM_COORDS_WIDTH, "x", NUM_COORDS_HEIGHT)
+            else:  
+                self.env.segmentation_info_mask = np.expand_dims(self.clip_seg_tta.heatmap.T.flatten(), axis=1)
+                self.env.segmentation_info_mask_unnormalized = np.expand_dims(self.clip_seg_tta.heatmap_unnormalized.T.flatten(), axis=1)
+
+            self.step_since_tta = 0
+
+            # NOTE: integration with app.py on hf
+            self.clip_seg_tta.executing_tta = False
+    
+
+    def convert_heatmap_resolution(self, heatmap, full_dims=(512, 512), new_dims=(24, 24)):
+        heatmap_large = resize(heatmap, full_dims, order=1,  # order=1 → bilinear
+             mode='reflect', anti_aliasing=True)
+        
+        coords = self.env.graph_generator.grid_coords   # (N, N, 2)
+        rows, cols = coords[...,1], coords[...,0]
+        heatmap_resized = heatmap_large[rows, cols]
+        heatmap_resized = heatmap_resized.reshape(new_dims[1], new_dims[0])
+        return heatmap_resized
+
+
+    def convert_labelmap_resolution(self, labelmap, full_dims=(512, 512), new_dims=(24, 24)):
+        """
+        1) Upsample via nearest‐neighbor to full_dims
+        2) Sample back down to your graph grid using grid_coords
+        """
+        # 1) Upsample with nearest‐neighbor, preserving integer labels
+        up = resize(
+            labelmap,
+            full_dims,
+            order=0,                # nearest‐neighbor
+            mode='edge',            # padding mode
+            preserve_range=True,    # don't normalize labels
+            anti_aliasing=False     # must be False for labels
+        ).astype(labelmap.dtype)    # back to original integer dtype
+
+        # 2) Downsample via your precomputed grid coords
+        coords = self.env.graph_generator.grid_coords   # shape (N, N, 2)
+        rows = coords[...,1].astype(int)
+        cols = coords[...,0].astype(int)
+
+        small = up[rows, cols]      # shape (N, N)
+        small = small.reshape(new_dims[0], new_dims[1])
+        return small
+    
+
+    def scale_trajectory(self,
+                        flat_indices,
+                        targets,
+                        old_dims=(17, 17),
+                        full_dims=(512, 512),
+                        new_dims=(24, 24)):
+        """
+        Args:
+            flat_indices: list of ints in [0..old_H*old_W-1]
+            targets:      list of (y_pix, x_pix) in [0..full_H-1]
+            old_dims:     (old_H, old_W)
+            full_dims:    (full_H, full_W)
+            new_dims:     (new_H, new_W)
+
+        Returns:
+            new_flat_traj: list of unique flattened indices in new_H×new_W
+            counts:        list of ints, same length as new_flat_traj
+        """
+        old_H, old_W = old_dims
+        full_H, full_W = full_dims
+        new_H,  new_W  = new_dims
+
+        # 1) bin targets into new grid
+        cell_h_new = full_H / new_H
+        cell_w_new = full_W / new_W
+        grid_counts = [[0]*new_W for _ in range(new_H)]
+        for x_pix, y_pix in targets:  # note (x, y) order as in original implementation
+            i_t = min(int(y_pix / cell_h_new), new_H - 1)
+            j_t = min(int(x_pix / cell_w_new), new_W - 1)
+            grid_counts[i_t][j_t] += 1
+
+        # 2) Walk the trajectory indices and project each old cell's *entire
+        #    pixel footprint* onto the finer 24×24 grid. 
+        cell_h_full = full_H / old_H
+        cell_w_full = full_W / old_W
+
+        seen = set()
+        new_flat_traj = []
+
+        for node_idx in flat_indices:
+            if node_idx < 0 or node_idx >= len(self.env.graph_generator.node_coords):
+                continue
+
+            coord_xy = self.env.graph_generator.node_coords[node_idx]
+            try:
+                row_old, col_old = self.env.graph_generator.find_index_from_grid_coords_2d(coord_xy)
+            except Exception:
+                continue
+
+            # Bounding box of the old cell in full-resolution pixel space
+            y0 = row_old * cell_h_full
+            y1 = (row_old + 1) * cell_h_full
+            x0 = col_old * cell_w_full
+            x1 = (col_old + 1) * cell_w_full
+
+            # Which new-grid rows & cols overlap? (inclusive ranges)
+            i_start = max(0, min(int(y0 / cell_h_new), new_H - 1))
+            i_end   = max(0, min(int((y1 - 1) / cell_h_new), new_H - 1))
+            j_start = max(0, min(int(x0 / cell_w_new), new_W - 1))
+            j_end   = max(0, min(int((x1 - 1) / cell_w_new), new_W - 1))
+
+            for ii in range(i_start, i_end + 1):
+                for jj in range(j_start, j_end + 1):
+                    f_new = ii * new_W + jj
+                    if f_new not in seen:
+                        seen.add(f_new)
+                        new_flat_traj.append(f_new)
+
+        # 3) annotate counts
+        counts = []
+        for f in new_flat_traj:
+            i_new, j_new = divmod(f, new_W)
+            counts.append(grid_counts[i_new][j_new])
+
+        return new_flat_traj, counts
+
+    ################################################################################

planner/worker.py

@@ -0,0 +1,272 @@
+#######################################################################
+# Name: worker.py
+#
+# - Runs robot in environment for N steps
+# - Collects & Returns S(t), A(t), R(t), S(t+1)
+#######################################################################
+
+from .parameter import *
+
+import os
+import json
+import copy
+import imageio
+import numpy as np
+import torch
+from time import time
+from .env import Env
+from .robot import Robot
+
+class Worker:
+    def __init__(self, meta_agent_id, n_agent, policy_net, q_net, global_step, device='cuda', greedy=False, save_image=False, clip_seg_tta=None):
+        self.device = device
+        self.greedy = greedy
+        self.n_agent = n_agent
+        self.metaAgentID = meta_agent_id
+        self.global_step = global_step
+        self.node_padding_size = NODE_PADDING_SIZE
+        self.k_size = K_SIZE
+        self.save_image = save_image
+        self.clip_seg_tta = clip_seg_tta
+
+        # Randomize map_index
+        mask_index = None
+        if MASKS_RAND_INDICES_PATH != "":
+            with open(MASKS_RAND_INDICES_PATH, 'r') as f:
+                mask_index_rand_json = json.load(f)
+            mask_index = mask_index_rand_json[self.global_step % len(mask_index_rand_json)]
+            print("mask_index: ", mask_index)
+
+        self.env = Env(map_index=self.global_step, n_agent=n_agent, k_size=self.k_size, plot=save_image, mask_index=mask_index)
+        self.local_policy_net = policy_net
+        self.local_q_net = q_net
+
+        self.robot_list = []
+        self.all_robot_positions = []
+
+        for i in range(self.n_agent):
+            robot_position = self.env.start_positions[i]
+            robot = Robot(robot_id=i, position=robot_position, plot=save_image)
+            self.robot_list.append(robot)
+            self.all_robot_positions.append(robot_position)
+
+        self.perf_metrics = dict()
+        self.episode_buffer = []
+        for i in range(15):
+            self.episode_buffer.append([])
+
+
+    def run_episode(self, curr_episode):
+
+        eps_start = time()
+        done = False
+        for robot_id, deciding_robot in enumerate(self.robot_list):
+            deciding_robot.observations = self.get_observations(deciding_robot.robot_position)
+
+        ### Run episode ###
+        for step in range(NUM_EPS_STEPS):
+
+            next_position_list = []
+            dist_list = []
+            travel_dist_list = []
+            dist_array = np.zeros((self.n_agent, 1))
+            for robot_id, deciding_robot in enumerate(self.robot_list):
+                observations = deciding_robot.observations
+                deciding_robot.save_observations(observations)
+
+                ### Forward pass through policy to get next position ###
+                next_position, action_index = self.select_node(observations)
+                deciding_robot.save_action(action_index)
+
+                dist = np.linalg.norm(next_position - deciding_robot.robot_position)
+
+                ### Log results of action (e.g. distance travelled) ###
+                dist_array[robot_id] = dist
+                dist_list.append(dist)
+                travel_dist_list.append(deciding_robot.travel_dist)
+                next_position_list.append(next_position)
+                self.all_robot_positions[robot_id] = next_position
+
+            arriving_sequence = np.argsort(dist_list)
+            next_position_list = np.array(next_position_list)
+            dist_list = np.array(dist_list)
+            travel_dist_list = np.array(travel_dist_list)
+            next_position_list = next_position_list[arriving_sequence]
+            dist_list = dist_list[arriving_sequence]
+            travel_dist_list = travel_dist_list[arriving_sequence]
+
+            ### Take Action (Deconflict if 2 agents choose the same target position) ###
+            next_position_list, dist_list = self.solve_conflict(arriving_sequence, next_position_list, dist_list)
+            reward_list, done = self.env.multi_robot_step(next_position_list, dist_list, travel_dist_list)
+
+            ### Update observations + rewards from action ###
+            for reward, robot_id in zip(reward_list, arriving_sequence):
+                robot = self.robot_list[robot_id]
+                robot.observations = self.get_observations(robot.robot_position)
+                robot.save_trajectory_coords(self.env.find_index_from_coords(robot.robot_position), self.env.num_new_targets_found)
+                robot.save_reward_done(reward, done)
+                robot.save_next_observations(robot.observations)
+
+            ### Save a frame to generate gif of robot trajectories ###
+            if self.save_image:
+                robots_route = []
+                for robot in self.robot_list:
+                    robots_route.append([robot.xPoints, robot.yPoints])
+                if not os.path.exists(GIFS_PATH):
+                    os.makedirs(GIFS_PATH)
+                self.env.plot_env(self.global_step, GIFS_PATH, step, max(travel_dist_list), robots_route)
+    
+            if done:
+                break
+
+        for robot in self.robot_list:
+            for i in range(15):
+                self.episode_buffer[i] += robot.episode_buffer[i]
+
+        self.perf_metrics['travel_dist'] = max(travel_dist_list)
+        self.perf_metrics['explored_rate'] = self.env.explored_rate
+        self.perf_metrics['targets_found'] = self.env.targets_found_rate
+        self.perf_metrics['success_rate'] = done
+
+        # save gif
+        if self.save_image:
+            path = GIFS_PATH
+            self.make_gif(path, curr_episode)
+
+        print(YELLOW, f"[Eps {curr_episode} Completed] Time Taken: {time()-eps_start:.2f}s, Steps: {step+1}", NC)
+
+
+    def get_observations(self, robot_position):
+        """ Get robot's sensor observation of environment given position """
+        current_node_index = self.env.find_index_from_coords(robot_position)
+        current_index = torch.tensor([current_node_index]).unsqueeze(0).unsqueeze(0).to(self.device)  # (1,1,1)
+    
+        node_coords = copy.deepcopy(self.env.node_coords)
+        graph = copy.deepcopy(self.env.graph)
+        node_utility = copy.deepcopy(self.env.node_utility)
+        guidepost = copy.deepcopy(self.env.guidepost)
+        segmentation_info_mask = copy.deepcopy(self.env.filtered_seg_info_mask)
+
+        n_nodes = node_coords.shape[0]
+        node_coords = node_coords / 640
+        node_utility = node_utility / 50
+
+        node_utility_inputs = node_utility.reshape((n_nodes, 1))
+
+        occupied_node = np.zeros((n_nodes, 1))
+        for position in self.all_robot_positions:
+            index = self.env.find_index_from_coords(position)
+            if index == current_index.item():
+                occupied_node[index] = -1
+            else:
+                occupied_node[index] = 1
+
+        node_inputs = np.concatenate((node_coords, segmentation_info_mask, guidepost), axis=1)
+        node_inputs = torch.FloatTensor(node_inputs).unsqueeze(0).to(self.device)  # (1, node_padding_size+1, 3)
+
+        assert node_coords.shape[0] < self.node_padding_size
+        padding = torch.nn.ZeroPad2d((0, 0, 0, self.node_padding_size - node_coords.shape[0]))
+        node_inputs = padding(node_inputs)
+
+        node_padding_mask = torch.zeros((1, 1, node_coords.shape[0]), dtype=torch.int64).to(self.device)
+        node_padding = torch.ones((1, 1, self.node_padding_size - node_coords.shape[0]), dtype=torch.int64).to(
+            self.device)
+        node_padding_mask = torch.cat((node_padding_mask, node_padding), dim=-1)
+
+        graph = list(graph.values())
+        edge_inputs = []
+        for node in graph:
+            node_edges = list(map(int, node))
+            edge_inputs.append(node_edges)
+
+        bias_matrix = self.calculate_edge_mask(edge_inputs)
+        edge_mask = torch.from_numpy(bias_matrix).float().unsqueeze(0).to(self.device)
+
+        assert len(edge_inputs) < self.node_padding_size
+        padding = torch.nn.ConstantPad2d(
+            (0, self.node_padding_size - len(edge_inputs), 0, self.node_padding_size - len(edge_inputs)), 1)
+        edge_mask = padding(edge_mask)
+        padding2 = torch.nn.ZeroPad2d((0, 0, 0, self.node_padding_size - len(edge_inputs)))
+
+        for edges in edge_inputs:
+            while len(edges) < self.k_size:
+                edges.append(0)        
+
+        edge_inputs = torch.tensor(edge_inputs).unsqueeze(0).to(self.device)  # (1, node_padding_size+1, k_size)
+        edge_inputs = padding2(edge_inputs)
+
+        edge_padding_mask = torch.zeros((1, len(edge_inputs), K_SIZE), dtype=torch.int64).to(self.device)
+        one = torch.ones_like(edge_padding_mask, dtype=torch.int64).to(self.device)
+        edge_padding_mask = torch.where(edge_inputs == 0, one, edge_padding_mask)
+
+        observations = node_inputs, edge_inputs, current_index, node_padding_mask, edge_padding_mask, edge_mask
+        return observations
+
+
+    def select_node(self, observations):
+        """ Forward pass through policy to get next position to go to on map """
+        node_inputs, edge_inputs, current_index, node_padding_mask, edge_padding_mask, edge_mask = observations
+        with torch.no_grad():
+            logp_list = self.local_policy_net(node_inputs, edge_inputs, current_index, node_padding_mask,
+                                              edge_padding_mask, edge_mask)
+
+        if self.greedy:
+            action_index = torch.argmax(logp_list, dim=1).long()
+        else:
+            action_index = torch.multinomial(logp_list.exp(), 1).long().squeeze(1)
+
+        next_node_index = edge_inputs[:, current_index.item(), action_index.item()]
+
+        next_position = self.env.node_coords[next_node_index]
+
+        return next_position, action_index
+
+
+    def solve_conflict(self, arriving_sequence, next_position_list, dist_list):
+        """ Deconflict if 2 agents choose the same target position """
+        for j, [robot_id, next_position] in enumerate(zip(arriving_sequence, next_position_list)):
+            moving_robot = self.robot_list[robot_id]
+            # if next_position[0] + next_position[1] * 1j in (next_position_list[:, 0] + next_position_list[:, 1] * 1j)[:j]:
+            #     dist_to_next_position = np.argsort(np.linalg.norm(self.env.node_coords - next_position, axis=1))
+            #     k = 0
+            #     while next_position[0] + next_position[1] * 1j in (next_position_list[:, 0] + next_position_list[:, 1] * 1j)[:j]:
+            #         k += 1
+            #         next_position = self.env.node_coords[dist_to_next_position[k]]
+
+            dist = np.linalg.norm(next_position - moving_robot.robot_position)
+            next_position_list[j] = next_position
+            dist_list[j] = dist
+            moving_robot.travel_dist += dist
+            moving_robot.robot_position = next_position
+
+        return next_position_list, dist_list
+
+
+    def work(self, currEpisode):
+        '''
+        Interacts with the environment. The agent gets either gradients or experience buffer
+        '''
+        self.run_episode(currEpisode)
+
+    def calculate_edge_mask(self, edge_inputs):
+        size = len(edge_inputs)
+        bias_matrix = np.ones((size, size))
+        for i in range(size):
+            for j in range(size):
+                if j in edge_inputs[i]:
+                    bias_matrix[i][j] = 0
+        return bias_matrix
+
+
+    def make_gif(self, path, n):
+        """ Generate a gif given list of images """
+        with imageio.get_writer('{}/{}_target_rate_{:.2f}.gif'.format(path, n, self.env.targets_found_rate), mode='I',
+                                fps=5) as writer:
+            for frame in self.env.frame_files:
+                image = imageio.imread(frame)
+                writer.append_data(image)
+        print('gif complete\n')
+
+        # Remove files
+        for filename in self.env.frame_files[:-1]:
+            os.remove(filename)