Спасибо за ответ, использовать другой input_shape модели, к сожаленью не могу, так как датасет на 256x256. Однако, думаю что 2. и 3. советы помогут. Возможно, нужно просто правильно проаугментировать данные, для более корректного сшивания.
Спасибо, попробую, однако, я думаю, что проблема всё-таки глобальнее, потому что когда я смотрел другие примеры сделанные не с помощью tf-agents, то обучение двигалось быстрее
