for my thesis I’m studying the algorithm Neural A Star (https://github.com/omron-sinicx/neural-astar) and to take familiarity with the code and after created the SDD datasets with

sh 3_SDD.sh

in https://github.com/omron-sinicx/planning-datasets/tree/7f8953c4b0f511c2cc03410abfdb3687d54deafb, i wanted to train the encoder with these, so i built a dataloader to extract SDD dataset:

    from __future__ import annotations, print_function

import numpy as np
import torch
import torch.utils.data as data
from neural_astar.planner.differentiable_astar import AstarOutput
from PIL import Image
from torchvision.utils import make_grid
import os


def create_sdd_dataloader(
        dirname: str,
        batch_size: int,
) -> data.DataLoader:
    
    dataset = SDD_Dataset(dirname)
    return data.DataLoader(
        dataset, batch_size, num_workers=0
    )


class SDD_Dataset(data.Dataset):
    def __init__(
            self,
            dirname: str,        
    ):
        self.dirname = dirname
        dir = os.fsdecode(dirname)
        images = []
        start_images = []
        goal_images = []
        traj_images = []
        for file in os.listdir(dir):
            filename = os.fsdecode(file)
            if filename.endswith("npz"):
                (
                    image,
                    start_image,
                    goal_image,
                    traj_image,
                ) = self._process(filename)
    
                images.append(image)
                start_images.append(start_image)
                goal_images.append(goal_image)
                traj_images.append(traj_image)
        self.images = self.toTensor(tlist=images)
        self.start_images = self.toTensor(tlist=start_images).unsqueeze(-1)
        self.goal_images = self.toTensor(tlist=goal_images).unsqueeze(-1)
        self.traj_images = self.toTensor(tlist=traj_images)

                


    def _process(self, filename: str):
        with np.load(self.dirname + "/" +filename) as f:
            image = f["image"]
            start_image = f["start_image"]
            goal_image = f["goal_image"]
            traj_image = f["traj_image"]
        image = image.astype(np.float32)
        start_image = start_image.astype(np.float32)
        goal_image = goal_image.astype(np.float32)
        traj_image = traj_image.astype(np.float32)

        return torch.tensor(image), torch.tensor(start_image), torch.tensor(goal_image), torch.tensor(traj_image)


    def __getitem__(self, index: int):
        image = self.images[index]
        start_map = self.start_images[index]
        goal_map = self.goal_images[index]
        traj_image = self.traj_images[index]

        return image, start_map, goal_map, traj_image
    
    def __len__(self):
        return self.images.shape[0]


    def toTensor(self, tlist: list):
        dims = list(tlist[0].shape)
        result = torch.zeros(len(tlist), *dims)
        i = 0
        for image in tlist:
            result[i] = image
            i+=1
        return result

                

And a training script:

from __future__ import annotations

import os

import hydra
import pytorch_lightning as pl
import torch
from neural_astar.planner import NeuralAstar
from neural_astar.utils.data_sdd import create_sdd_dataloader
from neural_astar.utils.training import PlannerModule, set_global_seeds
from pytorch_lightning.callbacks import ModelCheckpoint 

@hydra.main(config_path="config", config_name="train_sdd")
def main(config):

    set_global_seeds(config.seed)
    train_loader1 = create_sdd_dataloader(config.dataset + "video1", config.params.batch_size)
    val_loader1 = create_sdd_dataloader(config.dataset + "video2", config.params.batch_size)

    neural_astar = NeuralAstar(   
        encoder_arch=config.encoder.arch,
        encoder_depth=config.encoder.depth,
        encoder_input=config.encoder.input,
        #const=config.encoder.const,
        learn_obstacles=True,
        Tmax=config.Tmax,
    )
     
    

    checkpoint_callback = ModelCheckpoint(
        monitor="metrics/val_loss", save_weights_only=True, mode="max"
    )

    module = PlannerModule(neural_astar, config)
    logdir= f"{config.logdir}/{os.path.basename(config.dataset)}"
    trainer = pl.Trainer(
        accelerator= "gpu" if torch.cuda.is_available() else "cpu",
        log_every_n_steps=1,
        default_root_dir=logdir,
        max_epochs=config.params.num_epochs,
        callbacks=[checkpoint_callback],
    )

    trainer.fit(module, train_loader1, val_loader1)


if __name__ == "__main__":
    main()

With this yaml file:

seed: 1234
dataset: "planning-datasets/data/sdd/s064_0.5_128_300/bookstore/"
logdir: "model"
Tmax: 0.25

encoder:
  input: "rgb+"
  arch: "CNNDownSize"
  depth: 3
  const: 10.0s

params:
  batch_size: 100
  num_epochs: 5
  lr: 0.001


hydra:
  job:
    chdir: False

    

But when i try to run it, this is the output:

scripts/train_sdd.py:13: UserWarning: 
The version_base parameter is not specified.
Please specify a compatability version level, or None.
Will assume defaults for version 1.1
  @hydra.main(config_path="config", config_name="train_sdd")
WARNING: learn_obstacles has been set to True
GPU available: False, used: False
TPU available: False, using: 0 TPU cores
IPU available: False, using: 0 IPUs
HPU available: False, using: 0 HPUs
Missing logger folder: model/lightning_logs

  | Name          | Type         | Params
-----------------------------------------------
0 | planner       | NeuralAstar  | 95.1 K
1 | vanilla_astar | VanillaAstar | 9     
-----------------------------------------------
95.1 K    Trainable params
18        Non-trainable params
95.2 K    Total params
0.381     Total estimated model params size (MB)
Sanity Checking: 0it [00:00, ?it/s]/home/sam/Desktop/Tesi/neural-astar/.venv/lib/python3.8/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:224: PossibleUserWarning: The dataloader, val_dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 4 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.
  rank_zero_warn(
Sanity Checking DataLoader 0:   0%|                       | 0/2 [00:00<?, ?it/s]Error executing job with overrides: []
Traceback (most recent call last):
  File "scripts/train_sdd.py", line 45, in main
    trainer.fit(module, train_loader1, val_loader1)
  File "/home/sam/Desktop/Tesi/neural-astar/.venv/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 603, in fit
    call._call_and_handle_interrupt(
  File "/home/sam/Desktop/Tesi/neural-astar/.venv/lib/python3.8/site-packages/pytorch_lightning/trainer/call.py", line 38, in _call_and_handle_interrupt
    return trainer_fn(*args, **kwargs)
  File "/home/sam/Desktop/Tesi/neural-astar/.venv/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 645, in _fit_impl
    self._run(model, ckpt_path=self.ckpt_path)
  File "/home/sam/Desktop/Tesi/neural-astar/.venv/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1098, in _run
    results = self._run_stage()
  File "/home/sam/Desktop/Tesi/neural-astar/.venv/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1177, in _run_stage
    self._run_train()
  File "/home/sam/Desktop/Tesi/neural-astar/.venv/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1190, in _run_train
    self._run_sanity_check()
  File "/home/sam/Desktop/Tesi/neural-astar/.venv/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1262, in _run_sanity_check
    val_loop.run()
  File "/home/sam/Desktop/Tesi/neural-astar/.venv/lib/python3.8/site-packages/pytorch_lightning/loops/loop.py", line 199, in run
    self.advance(*args, **kwargs)
  File "/home/sam/Desktop/Tesi/neural-astar/.venv/lib/python3.8/site-packages/pytorch_lightning/loops/dataloader/evaluation_loop.py", line 152, in advance
    dl_outputs = self.epoch_loop.run(self._data_fetcher, dl_max_batches, kwargs)
  File "/home/sam/Desktop/Tesi/neural-astar/.venv/lib/python3.8/site-packages/pytorch_lightning/loops/loop.py", line 199, in run
    self.advance(*args, **kwargs)
  File "/home/sam/Desktop/Tesi/neural-astar/.venv/lib/python3.8/site-packages/pytorch_lightning/loops/epoch/evaluation_epoch_loop.py", line 137, in advance
    output = self._evaluation_step(**kwargs)
  File "/home/sam/Desktop/Tesi/neural-astar/.venv/lib/python3.8/site-packages/pytorch_lightning/loops/epoch/evaluation_epoch_loop.py", line 234, in _evaluation_step
    output = self.trainer._call_strategy_hook(hook_name, *kwargs.values())
  File "/home/sam/Desktop/Tesi/neural-astar/.venv/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1480, in _call_strategy_hook
    output = fn(*args, **kwargs)
  File "/home/sam/Desktop/Tesi/neural-astar/.venv/lib/python3.8/site-packages/pytorch_lightning/strategies/strategy.py", line 390, in validation_step
    return self.model.validation_step(*args, **kwargs)
  File "/home/sam/Desktop/Tesi/neural-astar/.venv/lib/python3.8/site-packages/neural_astar/utils/training.py", line 68, in validation_step
    outputs = self.forward(map_designs, start_maps, goal_maps)
  File "/home/sam/Desktop/Tesi/neural-astar/.venv/lib/python3.8/site-packages/neural_astar/utils/training.py", line 53, in forward
    return self.planner(map_designs, start_maps, goal_maps)
  File "/home/sam/Desktop/Tesi/neural-astar/.venv/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/sam/Desktop/Tesi/neural-astar/.venv/lib/python3.8/site-packages/neural_astar/planner/astar.py", line 202, in forward
    cost_maps = self.encode(map_designs, start_maps, goal_maps)
  File "/home/sam/Desktop/Tesi/neural-astar/.venv/lib/python3.8/site-packages/neural_astar/planner/astar.py", line 178, in encode
    cost_maps = self.encoder(inputs)
  File "/home/sam/Desktop/Tesi/neural-astar/.venv/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/sam/Desktop/Tesi/neural-astar/.venv/lib/python3.8/site-packages/neural_astar/planner/encoder.py", line 33, in forward
    y = torch.sigmoid(self.model(x))
  File "/home/sam/Desktop/Tesi/neural-astar/.venv/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/sam/Desktop/Tesi/neural-astar/.venv/lib/python3.8/site-packages/torch/nn/modules/container.py", line 139, in forward
    input = module(input)
  File "/home/sam/Desktop/Tesi/neural-astar/.venv/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/sam/Desktop/Tesi/neural-astar/.venv/lib/python3.8/site-packages/torch/nn/modules/conv.py", line 457, in forward
    return self._conv_forward(input, self.weight, self.bias)
  File "/home/sam/Desktop/Tesi/neural-astar/.venv/lib/python3.8/site-packages/torch/nn/modules/conv.py", line 453, in _conv_forward
    return F.conv2d(input, weight, bias, self.stride,
RuntimeError: Given groups=1, weight of size [32, 4, 3, 3], expected input[100, 128, 64, 3] to have 4 channels, but got 128 channels instead

Set the environment variable HYDRA_FULL_ERROR=1 for a complete stack trace.

And if I uncomment the cost in model initialization, it gives me another error:

scripts/train_sdd.py:13: UserWarning: 
The version_base parameter is not specified.
Please specify a compatability version level, or None.
Will assume defaults for version 1.1
  @hydra.main(config_path="config", config_name="train_sdd")
Error executing job with overrides: []
Traceback (most recent call last):
  File "scripts/train_sdd.py", line 20, in main
    neural_astar = NeuralAstar(
  File "/home/sam/Desktop/Tesi/neural-astar/.venv/lib/python3.8/site-packages/neural_astar/planner/astar.py", line 147, in __init__
    self.encoder = encoder_arch(len(self.encoder_input), encoder_depth, const)
  File "/home/sam/Desktop/Tesi/neural-astar/.venv/lib/python3.8/site-packages/neural_astar/planner/encoder.py", line 25, in __init__
    self.const = nn.Parameter(torch.ones(1) * const)  #nn.Parameter fa si che la costante compaia in model.parameters()
TypeError: only integer tensors of a single element can be converted to an index

Set the environment variable HYDRA_FULL_ERROR=1 for a complete stack trace.