fix(utils): add MPS branch to torch RNG state serialization (#4014)

serialize_torch_rng_state/deserialize_torch_rng_state only handled CPU
and CUDA generators. On MPS, resumed training was not bit-exact for any
stochastic op (dropout, ACT's CVAE noise) since the MPS generator's state
was never saved or restored. Mirrors the existing CUDA branch using
torch.mps.get_rng_state/set_rng_state (available since torch 2.11).

Note: get_rng_state()/set_rng_state() (used by seeded_context()) have the
same gap but are out of scope here — happy to follow up separately if
useful.

Co-authored-by: Sunny Dave <sunnydave@Sunnys-Mac-Studio.local>
Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>

src/lerobot/utils/random_utils.py

@@ -85,6 +85,8 @@ def serialize_torch_rng_state() -> dict[str, torch.Tensor]:
     torch_rng_state_dict = {"torch_rng_state": torch.get_rng_state()}
     if torch.cuda.is_available():
         torch_rng_state_dict["torch_cuda_rng_state"] = torch.cuda.get_rng_state()
+    if torch.backends.mps.is_available():
+        torch_rng_state_dict["torch_mps_rng_state"] = torch.mps.get_rng_state()
     return torch_rng_state_dict
 
 
@@ -95,6 +97,8 @@ def deserialize_torch_rng_state(rng_state_dict: dict[str, torch.Tensor]) -> None
     torch.set_rng_state(rng_state_dict["torch_rng_state"])
     if torch.cuda.is_available() and "torch_cuda_rng_state" in rng_state_dict:
         torch.cuda.set_rng_state(rng_state_dict["torch_cuda_rng_state"])
+    if torch.backends.mps.is_available() and "torch_mps_rng_state" in rng_state_dict:
+        torch.mps.set_rng_state(rng_state_dict["torch_mps_rng_state"])
 
 
 def serialize_rng_state() -> dict[str, torch.Tensor]:

tests/utils/test_random_utils.py

@@ -73,6 +73,17 @@ def test_serialize_deserialize_torch_rng(fixed_seed):
     assert val2 == val3
 
 
+@pytest.mark.skipif(not torch.backends.mps.is_available(), reason="MPS not available")
+def test_serialize_deserialize_torch_rng_mps(fixed_seed):
+    _ = torch.rand(1, device="mps").item()
+    st = serialize_torch_rng_state()
+    assert "torch_mps_rng_state" in st
+    val2 = torch.rand(1, device="mps").item()
+    deserialize_torch_rng_state(st)
+    val3 = torch.rand(1, device="mps").item()
+    assert val2 == val3
+
+
 def test_serialize_deserialize_rng(fixed_seed):
     # Generate one from each library
     _ = random.random()