feat(train): route reward model training through rewards/factory instead of policies/factory

examples/tutorial/rl/reward_classifier_example.py

@@ -1,9 +1,7 @@
 import torch
 
 from lerobot.datasets import LeRobotDataset
-from lerobot.policies import make_policy
-from lerobot.rewards.classifier.configuration_classifier import RewardClassifierConfig
-from lerobot.rewards.factory import make_reward_pre_post_processors
+from lerobot.rewards import RewardClassifierConfig, make_reward_model, make_reward_pre_post_processors
 
 
 def main():
@@ -24,9 +22,9 @@ def main():
         model_name="microsoft/resnet-18",
     )
 
-    # Make policy, preprocessor, and optimizer
-    policy = make_policy(config, ds_meta=dataset.meta)
-    optimizer = config.get_optimizer_preset().build(policy.parameters())
+    # Make reward model, preprocessor, and optimizer
+    reward_model = make_reward_model(config, dataset_stats=dataset.meta.stats)
+    optimizer = config.get_optimizer_preset().build(reward_model.parameters())
     preprocessor, _ = make_reward_pre_post_processors(config, dataset_stats=dataset.meta.stats)
 
     classifier_id = "<user>/reward_classifier_hil_serl_example"
@@ -44,7 +42,7 @@ def main():
             batch = preprocessor(batch)
 
             # Forward pass
-            loss, output_dict = policy.forward(batch)
+            loss, output_dict = reward_model.forward(batch)
 
             # Backward pass and optimization
             optimizer.zero_grad()
@@ -60,8 +58,8 @@ def main():
 
     print("Training finished!")
 
-    # You can now save the trained policy.
-    policy.push_to_hub(classifier_id)
+    # You can now save the trained reward model.
+    reward_model.push_to_hub(classifier_id)
 
 
 if __name__ == "__main__":