BC и GAIL падали на любой машине с GPU

[DenisDrobyshev]

agent = BC(env)                       # device="auto" -> cuda
data  = collect_expert_dataset(...)   # device="cpu", всегда
agent.train(data)                     # RuntimeError на первом же батче

RuntimeError: Expected all tensors to be on the same device, but got mat1
is on cpu, different from other tensors on cuda:0

Это ровно то, как BC используется в README. И это никогда не работало на машине с GPU — а для библиотеки обучения с подкреплением это большинство машин.

Ничего придумывать не пришлось

ReplayBatch.to() уже существует, и его докстринг — это описание данного бага:

Lets an agent train from a dataset stored on a different device (for example a CPU-collected offline dataset feeding a CUDA agent) without a device mismatch.

CQL, IQL, TD3BC и DiffusionPolicy — все зовут его на той же строке, где сэмплят:

batch = dataset.sample(self.batch_size).to(self.device)

У decisionrl.imitation было два места, где сэмплят, и ни одно этого не делало. Теперь делают оба, ровно как четыре соседа.

У GAIL случай тоньше: свой policy-датасет он строит уже на self.device, так что «извне» приходит только экспертная сторона — и приходит всегда, потому что она каллерская, собранная collect_expert_dataset, у которого "cpu" захардкожен дефолтом и сказать иначе нечем.

Почему этого никто не видел

test_bc_predicts_and_round_trips, test_bc_imitates_expert и test_gail_imitates_expert падают на любой машине с GPU и проходят на каждом раннере CI, потому что у раннеров GPU нет.

Библиотека была зелёной ровно столько, сколько её проверяли только там, где баг невозможен.

Поэтому регрессионным тестам GPU тоже не нужен

Два из трёх проверяют шаг, а не падение — что каждый батч проходит через .to(self.device), — подглядывая за собственным to у батча. На раннере перенос — no-op, но его отсутствие всё равно роняет тест. Это и есть нужное свойство:

$ CUDA_VISIBLE_DEVICES=-1 pytest tests/test_imitation.py -q      # на НЕисправленном исходнике
FAILED test_bc_moves_every_batch_to_the_actors_device
FAILED test_gail_moves_the_expert_batch_to_the_discriminators_device
  Right contains 2 more items, first extra item: device(type='cpu')

То есть этот баг был бы пойман в CI, без всякого GPU.

Третий тест утверждает падение целиком и пропускается без CUDA — и говорит в докстринге, что CI его запустить не может, потому что именно необъявленность этой разницы и есть вся история.

Измерено

	до	после
test_imitation.py на CUDA-машине	3 failed, 2 passed	5 passed
test_imitation.py на CPU	—	7 passed, 1 skipped
test_offline.py, test_diffusion.py	passed	passed (их не трогал — они были правы)

ruff check чист.