Merge pull request #68 from meta-pytorch/mortimer/resetcode

[BUG] reset coding env properly

Author: pankit-eng

pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta"
 [project]
 name = "openenv"
 version = "0.1.0"
-requires-python = ">=3.8"
+requires-python = ">=3.10"
 dependencies = [
     "torch>=1.9.0",
     "numpy>=1.19.0",

src/envs/coding_env/server/python_codeact_env.py

@@ -13,12 +13,13 @@
 
 import uuid
 
-from core.env_server import Action, Environment, Observation, Transform
+from core.env_server import Action, Environment, Observation
 from core.tools import PyExecutor
 
 from ..models import CodeAction, CodeObservation, CodeState
 from .transforms import create_safe_coding_transform
 
+
 class PythonCodeActEnv(Environment):
     """
     Python Code Action Environment for executing code and tracking state.
@@ -61,6 +62,12 @@ def reset(self) -> Observation:
         # Add last_exit_code to state
         self._state.last_exit_code = 0
 
+        # Reset executor to clear any previously defined variables/functions
+        self._executor = PyExecutor()
+
+        # Reset transform to clear any accumulated state
+        self.transform = create_safe_coding_transform()
+
         # Return initial observation
         observation = CodeObservation(
             stdout="",

tests/envs/test_python_codeact_reset.py

@@ -0,0 +1,160 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Test that PythonCodeActEnv.reset() properly resets executor state."""
+
+import sys
+from pathlib import Path
+
+from envs.coding_env.models import CodeAction
+from envs.coding_env.server.python_codeact_env import PythonCodeActEnv
+
+# Add src to path
+sys.path.insert(0, str(Path(__file__).parent.parent.parent / "src"))
+
+
+def test_reset_clears_executor_state():
+    """Test that reset() clears functions and variables defined in
+    previous execution."""
+    env = PythonCodeActEnv()
+
+    # Initial reset
+    obs = env.reset()
+    assert obs.exit_code == 0
+    assert env.state.step_count == 0
+
+    # Define a function in the executor
+    action1 = CodeAction(code="def my_function():\n    return 'Hello from function'\n")
+    obs1 = env.step(action1)
+    assert obs1.exit_code == 0
+
+    # Call the function to verify it exists
+    action2 = CodeAction(code="result = my_function()\nprint(result)")
+    obs2 = env.step(action2)
+    assert obs2.exit_code == 0
+    assert "Hello from function" in obs2.stdout
+
+    # Reset the environment
+    obs_reset = env.reset()
+    assert obs_reset.exit_code == 0
+    assert env.state.step_count == 0
+
+    # Try to call the function again - should fail because executor was reset
+    action3 = CodeAction(code="result = my_function()\nprint(result)")
+    obs3 = env.step(action3)
+
+    # Should get an error because my_function is no longer defined
+    assert obs3.exit_code == 1  # Error exit code
+    assert "my_function" in obs3.stderr or "NameError" in obs3.stderr
+
+
+def test_reset_clears_variables():
+    """Test that reset() clears variables defined in previous execution."""
+    env = PythonCodeActEnv()
+
+    # Initial reset
+    env.reset()
+
+    # Define a variable
+    action1 = CodeAction(code="my_variable = 42\n")
+    obs1 = env.step(action1)
+    assert obs1.exit_code == 0
+
+    # Use the variable to verify it exists
+    action2 = CodeAction(code="print(my_variable)")
+    obs2 = env.step(action2)
+    assert obs2.exit_code == 0
+    assert "42" in obs2.stdout
+
+    # Reset the environment
+    env.reset()
+
+    # Try to use the variable again - should fail
+    action3 = CodeAction(code="print(my_variable)")
+    obs3 = env.step(action3)
+
+    # Should get an error because my_variable is no longer defined
+    assert obs3.exit_code == 1
+    assert "my_variable" in obs3.stderr or "NameError" in obs3.stderr
+
+
+def test_reset_clears_imports():
+    """Test that reset() clears imported modules from previous execution."""
+    env = PythonCodeActEnv()
+
+    # Initial reset
+    env.reset()
+
+    # Import a module and define an alias
+    action1 = CodeAction(code="import math as m\n")
+    obs1 = env.step(action1)
+    assert obs1.exit_code == 0
+
+    # Use the alias to verify it exists
+    action2 = CodeAction(code="print(m.pi)")
+    obs2 = env.step(action2)
+    assert obs2.exit_code == 0
+    assert "3.14" in obs2.stdout
+
+    # Reset the environment
+    env.reset()
+
+    # Try to use the alias again - should fail
+    action3 = CodeAction(code="print(m.pi)")
+    obs3 = env.step(action3)
+
+    # Should get an error because 'm' is no longer defined
+    assert obs3.exit_code == 1
+    assert (
+        "NameError" in obs3.stderr
+        or "'m'" in obs3.stderr
+        or "variable `m` is not defined" in obs3.stderr
+    )
+
+
+def test_reset_preserves_step_count_reset():
+    """Test that reset() properly resets step count."""
+    env = PythonCodeActEnv()
+
+    # Initial reset
+    env.reset()
+    assert env.state.step_count == 0
+
+    # Execute some steps
+    for i in range(5):
+        action = CodeAction(code=f"print({i})")
+        env.step(action)
+
+    assert env.state.step_count == 5
+
+    # Reset should reset step count
+    env.reset()
+    assert env.state.step_count == 0
+
+    # Execute another step
+    action = CodeAction(code="print('test')")
+    env.step(action)
+    assert env.state.step_count == 1
+
+
+def test_reset_changes_episode_id():
+    """Test that reset() generates a new episode ID."""
+    env = PythonCodeActEnv()
+
+    # Initial reset
+    env.reset()
+    episode_id_1 = env.state.episode_id
+
+    # Execute some steps
+    action = CodeAction(code="print('test')")
+    env.step(action)
+
+    # Reset and get new episode ID
+    env.reset()
+    episode_id_2 = env.state.episode_id
+
+    # Episode IDs should be different
+    assert episode_id_1 != episode_id_2