chore: update

Author: Lawhy

examples/math_env.py

@@ -1,3 +1,17 @@
+# Copyright 2025 Horizon RL Contributors
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 """Math environment example — run math problems through a Strands agent with a calculator tool.
 
 Usage:
@@ -15,11 +29,10 @@
 import logging
 
 import httpx
-from strands_tools import calculator
 
-from strands_env.core.environment import Environment
 from strands_env.core.models import ModelFactory, bedrock_model_factory, sglang_model_factory
 from strands_env.core.types import Action, TaskContext
+from strands_env.environments.simple_math_env import SimpleMathEnv
 from strands_env.rewards.math_reward import MathRewardFunction
 
 logger = logging.getLogger(__name__)
@@ -31,18 +44,6 @@
 ]
 
 
-# ---------------------------------------------------------------------------
-# Environment
-# ---------------------------------------------------------------------------
-
-
-class MathEnvironment(Environment):
-    """Environment that gives the agent a calculator tool to solve math problems."""
-
-    def get_tools(self) -> list:
-        return [calculator]
-
-
 # ---------------------------------------------------------------------------
 # Model factory helpers
 # ---------------------------------------------------------------------------
@@ -96,7 +97,7 @@ async def main() -> None:
     args = parser.parse_args()
 
     model_factory = create_model_factory(args)
-    env = MathEnvironment(model_factory=model_factory, reward_fn=MathRewardFunction(), verbose=args.verbose)
+    env = SimpleMathEnv(model_factory=model_factory, reward_fn=MathRewardFunction(), verbose=args.verbose)
 
     for question, ground_truth in MATH_PROBLEMS:
         print(f"\n{'=' * 60}")