[BE] Add type annotation for tensor_keys to facilitate auto-complete #3987

Job	Run time
unittests-meltingpot / linux-job	7m 18s
unittests-pettingzoo / linux-job	10m 25s
unittests-atari-dqn (3.9, 12.1) / linux-job	13m 4s
unittests-brax (3.11, 12.1) / linux-job	11m 33s
unittests-d4rl (3.9, 12.1) / linux-job	9m 21s
unittests-chess (3.9, 12.1) / linux-job	8m 46s
unittests-envpool (3.9, 12.1) / linux-job	9m 22s
unittests-gendgrl (3.9, 12.1) / linux-job	7m 54s
unittests-gym (3.9, 12.1) / linux-job	19m 46s
unittests-jumanji (3.9, 12.1) / linux-job	45m 27s
unittests-open_spiel (3.9, 12.1) / linux-job	9m 2s
unittests-minari (3.9, 12.1) / linux-job	6m 48s
unittests-robohive (3.9, 12.1) / linux-job	50m 5s
unittests-openx (3.9, 12.1) / linux-job	25m 11s
unittests-roboset (3.9, 12.1) / linux-job	11m 48s
unittests-sklearn (3.9, 12.1) / linux-job	2h 0m 47s
unittests-unity_mlagents (3.10.12, 12.1) / linux-job	9m 37s
unittests-smacv2 (3.9, 12.1) / linux-job	12m 50s
unittests-vd4rl (3.9, 12.1) / linux-job	6m 54s
unittests-vmas (3.9, 12.1) / linux-job	13m 28s
	6h 49m 26s

Provide feedback