Record max_steps in base_agent so that an agent may have different be…

…havior when reaching the last step.

PiperOrigin-RevId: 680600677

android_world/agents/base_agent.py

@@ -72,6 +72,8 @@ def __init__(
       )
     self._transition_pause = transition_pause
 
+    self._max_steps = None
+
   @property
   def transition_pause(self) -> float | None:
     return self._transition_pause
@@ -88,6 +90,9 @@ def env(self) -> interface.AsyncEnv:
   def env(self, env: interface.AsyncEnv) -> None:
     self._env = env
 
+  def set_max_steps(self, max_steps: int) -> None:
+    self._max_steps = max_steps
+
   def reset(self, go_home: bool = False) -> None:
     """Resets the agent."""
     self.env.reset(go_home=go_home)