chore: Add support for more experimental config with the ServerlessBackend

Author: bradhilton

src/art/dev/train.py

@@ -18,6 +18,7 @@ class TrainConfig(TypedDict, total=False):
     max_negative_advantage_importance_sampling_weight: float
     num_trajectories_learning_rate_multiplier_power: float
     plot_tensors: bool
+    ppo: bool
     precalculate_logprobs: bool
     scale_learning_rate_by_reward_std_dev: bool
     scale_rewards: bool

src/art/serverless/backend.py

@@ -129,7 +129,12 @@ async def _train_model(
             trajectory_groups=trajectory_groups,
             experimental_config=ExperimentalTrainingConfig(
                 advantage_balance=dev_config.get("advantage_balance"),
+                epsilon=dev_config.get("epsilon"),
+                epsilon_high=dev_config.get("epsilon_high"),
+                importance_sampling_level=dev_config.get("importance_sampling_level"),
                 learning_rate=config.learning_rate,
+                max_negative_advantage_importance_sampling_weight=dev_config.get("max_negative_advantage_importance_sampling_weight"),
+                ppo=dev_config.get("ppo"),
                 precalculate_logprobs=dev_config.get("precalculate_logprobs"),
                 scale_rewards=dev_config.get("scale_rewards"),
             ),

src/art/serverless/client.py

@@ -51,7 +51,12 @@ class DeleteCheckpointsResponse(BaseModel):
 
 class ExperimentalTrainingConfig(TypedDict, total=False):
     advantage_balance: float | None
+    epsilon: float | None
+    epsilon_high: float | None
+    importance_sampling_level: Literal["token", "sequence"] | None
     learning_rate: float | None
+    max_negative_advantage_importance_sampling_weight: float | None
+    ppo: bool | None
     precalculate_logprobs: bool | None
     scale_rewards: bool | None