refactor(sarm): update reset method and remove unused action prediction methods

This commit is contained in:
Khalil Meftah
2026-04-22 13:38:23 +02:00
parent 6b518a3de7
commit c6b8d36d48
+1 -9
View File
@@ -645,17 +645,9 @@ class SARMRewardModel(PreTrainedRewardModel):
return self.parameters()
def reset(self):
"""Required by PreTrainedPolicy but not used for reward models."""
"""SARM has no episode-level state to reset."""
pass
def predict_action_chunk(self, batch: dict[str, Tensor]) -> Tensor:
"""Required by PreTrainedPolicy but not used for reward models."""
raise NotImplementedError("SARM model does not predict action chunks")
def select_action(self, batch: dict[str, Tensor]) -> Tensor:
"""Required by PreTrainedPolicy but not used for SARM."""
raise NotImplementedError("SARM model does not select actions")
def _train_step(
self,
img_emb: torch.Tensor, # (B, N, T, D)