Upload folder using huggingface_hub

server/meta_optimizer_environment.py

@@ -29,6 +29,8 @@ from .tasks import TRAIN_TASK_IDS, get_task, task_spec_from_dict, TaskSpec
 LOSS_THRESHOLD = 0.1
 MAX_STEPS = 100
 BATCH_SIZE = 32
+# Dense reward scale: reward += DENSE_REWARD_SCALE * (prev_loss - current_loss) each step (potential-based, helps credit assignment)
+DENSE_REWARD_SCALE = 0.2
 
 
 def _build_model(spec: TaskSpec) -> nn.Module:
@@ -200,7 +202,9 @@ def run_meta_optimizer_trajectory(
 class MetaOptimizerEnvironment(Environment[MetaOptimizerAction, MetaOptimizerObservation, State]):
     """
     Meta-learning optimizer environment: agent chooses LR scale, momentum, grad clip, weight decay per step.
-    Reward = -steps_to_reach_threshold (convergence speed). Supports 50 train tasks and held-out eval.
+    Reward: dense term = scale * (prev_loss - current_loss) each step (loss decrease); terminal = -steps_to_threshold
+    when episode ends. Episode ends at max_steps or as soon as loss < threshold (early termination). Supports 50 train
+    tasks and held-out eval.
     """
 
     SUPPORTS_CONCURRENT_SESSIONS: bool = True
@@ -222,6 +226,7 @@ class MetaOptimizerEnvironment(Environment[MetaOptimizerAction, MetaOptimizerObs
         self._velocities: Optional[List[torch.Tensor]] = None
         self._step_count: int = 0
         self._current_loss: float = 0.0
+        self._prev_loss: float = 0.0  # for dense reward (loss decrease)
         self._steps_to_threshold: Optional[int] = None
         self._action_log: List[Dict[str, Any]] = []
         self._episode_id: Optional[str] = None
@@ -254,6 +259,7 @@ class MetaOptimizerEnvironment(Environment[MetaOptimizerAction, MetaOptimizerObs
         with torch.no_grad():
             out = self._model(X)
             self._current_loss = nn.functional.mse_loss(out, y).item()
+        self._prev_loss = self._current_loss
 
         return self._observation(reward=None, grad_norm=None)
 
@@ -264,6 +270,7 @@ class MetaOptimizerEnvironment(Environment[MetaOptimizerAction, MetaOptimizerObs
         **kwargs: Any,
     ) -> MetaOptimizerObservation:
         assert self._model is not None and self._task_spec is not None
+        prev_loss = self._prev_loss
         lr = action.lr_scale
         momentum = action.momentum_coef
         clip = action.grad_clip_threshold
@@ -310,11 +317,17 @@ class MetaOptimizerEnvironment(Environment[MetaOptimizerAction, MetaOptimizerObs
         if self._steps_to_threshold is None and self._current_loss < self.loss_threshold:
             self._steps_to_threshold = self._step_count
 
-        done = self._step_count >= self.max_steps
+        # Dense reward: reward loss decrease (potential-based shaping, does not change optimal policy)
+        dense_reward = DENSE_REWARD_SCALE * (prev_loss - self._current_loss)
+        self._prev_loss = self._current_loss
+
+        # End episode when we hit max_steps or when loss first crosses threshold (early termination)
+        done = self._step_count >= self.max_steps or self._steps_to_threshold is not None
         if done:
-            reward = -(self._steps_to_threshold if self._steps_to_threshold is not None else self.max_steps)
+            terminal = -(self._steps_to_threshold if self._steps_to_threshold is not None else self.max_steps)
+            reward = dense_reward + terminal
         else:
-            reward = 0.0
+            reward = dense_reward
 
         return self._observation(reward=reward, grad_norm=grad_norm, done=done)