├── overview.jpg ├── requirements.txt ├── .gitignore ├── CONTRIBUTING.md ├── src ├── agent_debugger │ ├── pycoworld │ │ ├── types.py │ │ ├── debugger.py │ │ ├── extractors.py │ │ └── interventions.py │ ├── types.py │ ├── agent.py │ ├── node.py │ ├── default_interventions.py │ └── debugger.py ├── pycoworld │ ├── serializable_environment.py │ ├── levels │ │ ├── color_memory.py │ │ ├── level_names.py │ │ ├── apples.py │ │ ├── key_door.py │ │ ├── utils.py │ │ ├── red_green_apples.py │ │ ├── grass_sand.py │ │ └── base_level.py │ ├── default_constants.py │ ├── default_sprites_and_drapes.py │ └── environment.py ├── impala_agent.py └── impala_net.py ├── README.md ├── LICENSE └── colabs └── experiments.ipynb /overview.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/agent_debugger/HEAD/overview.jpg -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | chex 2 | dill 3 | dm-env 4 | dm-haiku 5 | dm-tree 6 | jax 7 | numpy 8 | pycolab 9 | scipy 10 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # Distribution / packaging 7 | .Python 8 | build/ 9 | develop-eggs/ 10 | dist/ 11 | downloads/ 12 | eggs/ 13 | .eggs/ 14 | lib/ 15 | lib64/ 16 | parts/ 17 | sdist/ 18 | var/ 19 | wheels/ 20 | share/python-wheels/ 21 | *.egg-info/ 22 | .installed.cfg 23 | *.egg 24 | MANIFEST 25 | -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- 1 | # How to Contribute 2 | 3 | ## Contributor License Agreement 4 | 5 | Contributions to this project must be accompanied by a Contributor License 6 | Agreement. You (or your employer) retain the copyright to your contribution, 7 | this simply gives us permission to use and redistribute your contributions as 8 | part of the project. Head over to to see 9 | your current agreements on file or to sign a new one. 10 | 11 | You generally only need to submit a CLA once, so if you've already submitted one 12 | (even if it was for a different project), you probably don't need to do it 13 | again. 14 | 15 | ## Code reviews 16 | 17 | All submissions, including submissions by project members, require review. We 18 | use GitHub pull requests for this purpose. Consult 19 | [GitHub Help](https://help.github.com/articles/about-pull-requests/) for more 20 | information on using pull requests. 21 | 22 | ## Community Guidelines 23 | 24 | This project follows [Google's Open Source Community 25 | Guidelines](https://opensource.google/conduct/). 26 | -------------------------------------------------------------------------------- /src/agent_debugger/pycoworld/types.py: -------------------------------------------------------------------------------- 1 | # Copyright 2023 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | 16 | """Specific types for the Pycoworld debugger.""" 17 | 18 | from typing import Union 19 | 20 | from pycolab import engine as engine_lib 21 | from pycolab import things 22 | 23 | PycolabPosition = things.Sprite.Position 24 | Position = Union[PycolabPosition, tuple[float, float]] 25 | PycolabEngine = engine_lib.Engine 26 | 27 | 28 | def to_pycolab_position(pos: Position) -> PycolabPosition: 29 | """Returns a PycolabPosition from a tuple or a PycolabPosition.""" 30 | if isinstance(pos, tuple): 31 | return PycolabPosition(*pos) 32 | return pos 33 | -------------------------------------------------------------------------------- /src/agent_debugger/types.py: -------------------------------------------------------------------------------- 1 | # Copyright 2023 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | 16 | """Custom types used throughout the codebase. 17 | 18 | A seed is required for the agent. For the environment, it can be included 19 | but it's not mandatory in the dm_env standard. 20 | """ 21 | 22 | from typing import Any, Callable, NamedTuple 23 | 24 | import dm_env 25 | 26 | from agent_debugger.src.pycoworld import serializable_environment 27 | 28 | Action = Any 29 | Observation = Any 30 | TimeStep = dm_env.TimeStep 31 | FirstTimeStep = dm_env.StepType.FIRST 32 | MidTimeStep = dm_env.StepType.MID 33 | LastTimeStep = dm_env.StepType.LAST 34 | Environment = dm_env.Environment 35 | SerializableEnvironment = serializable_environment.SerializableEnvironment 36 | Agent = Any 37 | Rng = Any 38 | EnvState = Any 39 | EnvBuilder = Callable[[], Environment] 40 | SerializableEnvBuilder = Callable[[], SerializableEnvironment] 41 | 42 | 43 | # We don't use dataclasses as they are not supported by jax. 44 | class AgentState(NamedTuple): 45 | internal_state: Any 46 | seed: Rng 47 | -------------------------------------------------------------------------------- /src/agent_debugger/agent.py: -------------------------------------------------------------------------------- 1 | # Copyright 2023 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | 16 | """Base class of agent recognized by the Agent Debugger.""" 17 | 18 | import abc 19 | from typing import Any 20 | 21 | import chex 22 | import dm_env 23 | 24 | from agent_debugger.src.agent_debugger import types 25 | 26 | 27 | class DebuggerAgent(abc.ABC): 28 | """The standard agent interface to be used in the debugger. 29 | 30 | The internal state is wrapped with the seed. The step method takes an 31 | observation and not a timestep. The agent parameters (mostly neural network 32 | parameters) are fixed and are NOT part of the state. 33 | """ 34 | 35 | @abc.abstractmethod 36 | def initial_state(self, rng: chex.PRNGKey) -> types.AgentState: 37 | """Returns the initial state of the agent.""" 38 | 39 | @abc.abstractmethod 40 | def step( 41 | self, 42 | timestep: dm_env.TimeStep, 43 | state: types.AgentState, 44 | ) -> tuple[types.Action, Any, types.AgentState]: 45 | """Steps the agent, taking a timestep (including observation) and a state.""" 46 | -------------------------------------------------------------------------------- /src/pycoworld/serializable_environment.py: -------------------------------------------------------------------------------- 1 | # Copyright 2023 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | 16 | """Serializable environment class. 17 | 18 | The main feature is to ability to retrieve and set the state of an environment. 19 | """ 20 | 21 | import abc 22 | from typing import Any 23 | 24 | import dm_env 25 | 26 | 27 | class SerializableEnvironment(dm_env.Environment, abc.ABC): 28 | """Abstract class with methods to set and get states. 29 | 30 | The state can be anything even though we prefer mappings like dicts for 31 | readability. It must contain all information, no compression is allowed. The 32 | state must be a parse of the environment containing only the stateful 33 | variables. 34 | There is no assumption on how the agent would use this object: please provide 35 | copies of internal attributes to avoid side effects. 36 | """ 37 | 38 | @abc.abstractmethod 39 | def get_state(self) -> Any: 40 | """Returns the state of the environment.""" 41 | 42 | @abc.abstractmethod 43 | def set_state(self, state: Any) -> Any: 44 | """Sets the state of the environment. 45 | 46 | Args: 47 | state: The state to set. 48 | """ 49 | -------------------------------------------------------------------------------- /src/agent_debugger/pycoworld/debugger.py: -------------------------------------------------------------------------------- 1 | # Copyright 2023 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | 16 | """Specific implementation of the debugger for pycoworld.""" 17 | 18 | from agent_debugger.src.agent_debugger import agent as debugger_agent 19 | from agent_debugger.src.agent_debugger import debugger as dbg 20 | from agent_debugger.src.agent_debugger.pycoworld import extractors as pycolab_extractors 21 | from agent_debugger.src.agent_debugger.pycoworld import interventions as pycolab_interventions 22 | from agent_debugger.src.pycoworld import serializable_environment 23 | 24 | 25 | class PycoworldDebugger(dbg.Debugger): 26 | """Overriding the Debugger with additional interventions and extractors. 27 | 28 | Attributes: 29 | interventions: An object to intervene on pycoworld nodes. 30 | extractors: An object to extract information from pycoworld nodes. 31 | """ 32 | 33 | def __init__( 34 | self, 35 | agent: debugger_agent.DebuggerAgent, 36 | env: serializable_environment.SerializableEnvironment, 37 | ) -> None: 38 | """Initializes the object.""" 39 | super().__init__(agent, env) 40 | 41 | self.interventions = pycolab_interventions.PycoworldInterventions(self._env) 42 | self.extractors = pycolab_extractors.PycoworldExtractors() 43 | -------------------------------------------------------------------------------- /src/agent_debugger/node.py: -------------------------------------------------------------------------------- 1 | # Copyright 2023 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | 16 | """The base node interface used in the debugger.""" 17 | 18 | import dataclasses 19 | from typing import Any, Optional, Sequence 20 | 21 | from agent_debugger.src.agent_debugger import types 22 | 23 | 24 | @dataclasses.dataclass 25 | class Node: 26 | """A node is the concatenation of the agent_state and the env_state. 27 | 28 | Besides the agent and env state, it also carries the action taken by the agent 29 | and the timestep returned by the environment in the last transition. This 30 | transition can be seen as the arrow leading to the state, therefore we call it 31 | 'last_action' and 'last_timestep'. The latter will be used to step to the next 32 | state. 33 | Finally, the node also carries some future agent actions that the user wants 34 | to enforce. The first element of the list will be used to force the next 35 | transition action, and the list minus this first element will be passed to the 36 | next state. 37 | """ 38 | 39 | agent_state: types.AgentState 40 | last_action: types.Action 41 | env_state: types.EnvState 42 | last_timestep: types.TimeStep 43 | forced_next_actions: Sequence[types.Action] = () 44 | episode_step: int = 0 45 | last_agent_output: Optional[Any] = None 46 | 47 | @property 48 | def is_terminal(self) -> bool: 49 | """Returns whether the last timestep is of step_type LAST.""" 50 | return self.last_timestep.step_type == types.LastTimeStep 51 | -------------------------------------------------------------------------------- /src/pycoworld/levels/color_memory.py: -------------------------------------------------------------------------------- 1 | # Copyright 2023 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | 16 | """Color memory level.""" 17 | 18 | import numpy as np 19 | 20 | from agent_debugger.src.pycoworld.levels import grass_sand 21 | 22 | 23 | class ColorMemoryLevel(grass_sand.GrassSandLevel): 24 | """A modification of the grass sand environment to test memory. 25 | 26 | The position of the reward is correlated with the floor type, as 27 | in grass_sand, but the agent has an egocentric view. Therefore, it must 28 | remember the color at the beginning to get the reward as fast as possible. 29 | """ 30 | 31 | def __init__(self, corr: float = 1.0, large: bool = False) -> None: 32 | """Initializes the level. 33 | 34 | Args: 35 | corr: see grass_sand 36 | large: whether to enlarge the initial corridor. 37 | """ 38 | if not large: 39 | above = np.array([[0, 0, 0, 0, 4, 4, 4], [0, 0, 0, 0, 4, 0, 4], 40 | [4, 4, 4, 4, 4, 0, 4], [4, 4, 4, 4, 4, 0, 4], 41 | [4, 99, 8, 0, 0, 0, 4], [4, 4, 4, 4, 4, 0, 4], 42 | [4, 4, 4, 4, 4, 0, 4], [0, 0, 0, 0, 4, 0, 4], 43 | [0, 0, 0, 0, 4, 4, 4]]) 44 | else: 45 | above = np.array([[0, 0, 0, 0, 4, 4, 4], [0, 0, 0, 0, 4, 0, 4], 46 | [4, 4, 4, 4, 4, 0, 4], [4, 4, 0, 0, 0, 0, 4], 47 | [4, 99, 8, 0, 0, 0, 4], [4, 4, 0, 0, 0, 0, 4], 48 | [4, 4, 4, 4, 4, 0, 4], [0, 0, 0, 0, 4, 0, 4], 49 | [0, 0, 0, 0, 4, 4, 4]]) 50 | below = np.zeros(above.shape, dtype=np.uint8) 51 | reward_pos = (np.array([1, 5]), np.array([7, 5])) 52 | super().__init__(corr, above, below, reward_pos, double_terminal_event=True) 53 | -------------------------------------------------------------------------------- /src/pycoworld/levels/level_names.py: -------------------------------------------------------------------------------- 1 | # Copyright 2023 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | 16 | """Level names for backward compatibility. 17 | 18 | This module lists the names of levels for backward compatibililty with old code. 19 | We plan to deprecate this module in the future so please do not rely on it for 20 | new code. 21 | """ 22 | 23 | from agent_debugger.src.pycoworld.levels import apples 24 | from agent_debugger.src.pycoworld.levels import base_level 25 | from agent_debugger.src.pycoworld.levels import color_memory 26 | from agent_debugger.src.pycoworld.levels import grass_sand 27 | from agent_debugger.src.pycoworld.levels import key_door 28 | from agent_debugger.src.pycoworld.levels import red_green_apples 29 | 30 | 31 | # Provided for backward compatibility only. Please do not use in new code. 32 | def level_by_name(level_name: str) -> base_level.PycoworldLevel: 33 | """Returns the level corresponding to a given level name. 34 | 35 | Provided for backward compatibility only. Please do not use in new code. 36 | 37 | Args: 38 | level_name: Name of the level to construct. See NAMED_LEVELS for the 39 | supported level names. 40 | """ 41 | if level_name == 'apples_corner': 42 | return apples.ApplesLevel(start_type='corner') 43 | elif level_name == 'apples_full': 44 | return apples.ApplesLevel() 45 | elif level_name == 'grass_sand': 46 | return grass_sand.GrassSandLevel() 47 | elif level_name == 'grass_sand_uncorrelated': 48 | return grass_sand.GrassSandLevel(corr=0.5, double_terminal_event=True) 49 | elif level_name == 'key_door': 50 | return key_door.KeyDoorLevel() 51 | elif level_name == 'key_door_closed': 52 | return key_door.KeyDoorLevel(door_type='closed') 53 | elif level_name == 'large_color_memory': 54 | return color_memory.ColorMemoryLevel(large=True) 55 | elif level_name == 'red_green_apples': 56 | return red_green_apples.RedGreenApplesLevel() 57 | 58 | raise ValueError(f'Unknown level {level_name}.') 59 | -------------------------------------------------------------------------------- /src/agent_debugger/pycoworld/extractors.py: -------------------------------------------------------------------------------- 1 | # Copyright 2023 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | 16 | """Pycoworld extractors. 17 | 18 | The extractors act on the node directly, and extract information from it. 19 | """ 20 | 21 | import numpy as np 22 | from pycolab import things 23 | 24 | from agent_debugger.src.agent_debugger import node as node_lib 25 | from agent_debugger.src.agent_debugger.pycoworld import types 26 | 27 | 28 | class PycoworldExtractors: 29 | """Object containing the pycoworld extractors.""" 30 | 31 | # pylint: disable=protected-access 32 | def get_element_positions( 33 | self, 34 | node: node_lib.Node, 35 | element_id: int, 36 | ) -> list[types.Position]: 37 | """Returns the position of a sprite/drape/background element. 38 | 39 | Args: 40 | node: The node to intervene on. 41 | element_id: The identifier of the element. A full list can be found in 42 | the pycoworld/default_constants.py file. 43 | """ 44 | engine = node.env_state.current_game 45 | element_key = chr(element_id) 46 | if element_key in engine._sprites_and_drapes: 47 | # If it is a sprite. 48 | if isinstance(engine._sprites_and_drapes[element_key], things.Sprite): 49 | return [engine._sprites_and_drapes[element_key].position] 50 | # Otherwise, it's a drape. 51 | else: 52 | curtain = engine._sprites_and_drapes[element_key].curtain 53 | list_tuples = list(zip(*np.where(curtain))) 54 | return [types.PycolabPosition(*x) for x in list_tuples] 55 | 56 | # Last resort: we look at the board and the backdrop. 57 | board = engine._board.board 58 | list_tuples = list(zip(*np.where(board == ord(element_key)))) 59 | backdrop = engine._backdrop.curtain 60 | list_tuples += list(zip(*np.where(backdrop == ord(element_key)))) 61 | return [types.PycolabPosition(*x) for x in list_tuples] 62 | 63 | def get_backdrop_curtain(self, node: node_lib.Node) -> np.ndarray: 64 | """Returns the backdrop of a pycolab engine from a node.""" 65 | return node.env_state.current_game._backdrop.curtain 66 | 67 | # pylint: enable=protected-access 68 | -------------------------------------------------------------------------------- /src/impala_agent.py: -------------------------------------------------------------------------------- 1 | # Copyright 2023 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | 16 | """Impala agent class, which follows the DebuggerAgent interface.""" 17 | 18 | from typing import Any, Callable 19 | 20 | import dm_env 21 | import haiku as hk 22 | import jax 23 | import jax.numpy as jnp 24 | 25 | from agent_debugger.src.agent_debugger import agent 26 | from agent_debugger.src.agent_debugger import types 27 | 28 | 29 | class ImpalaAgent(agent.DebuggerAgent): 30 | """Impala agent class.""" 31 | 32 | def __init__( 33 | self, 34 | net_factory: Callable[[], hk.RNNCore], 35 | params: hk.Params, 36 | ) -> None: 37 | """Initializes the agent. 38 | 39 | Args: 40 | net_factory: Function to create the model. 41 | params: The parameters of the agent. 42 | """ 43 | _, self._initial_state = hk.transform( 44 | lambda batch_size: net_factory().initial_state(batch_size)) 45 | 46 | self._init_fn, apply_fn = hk.without_apply_rng( 47 | hk.transform(lambda obs, state: net_factory().__call__(obs, state))) 48 | self._apply_fn = jax.jit(apply_fn) 49 | 50 | self._params = params 51 | 52 | def initial_state(self, rng: jnp.ndarray) -> types.AgentState: 53 | """Returns the agent initial state.""" 54 | # Wrapper method to avoid pytype attribute-error. 55 | return types.AgentState( 56 | internal_state=self._initial_state(self._params, rng, batch_size=1), 57 | seed=rng) 58 | 59 | def step( 60 | self, 61 | timestep: dm_env.TimeStep, 62 | state: types.AgentState, 63 | ) -> tuple[types.Action, Any, types.AgentState]: 64 | """Steps the agent in the environment.""" 65 | net_out, next_state = self._apply_fn(self._params, timestep, 66 | state.internal_state) 67 | 68 | # Sample an action and return. 69 | action = hk.multinomial(state.seed, net_out.policy_logits, num_samples=1) 70 | action = jnp.squeeze(action, axis=-1) 71 | action = int(action) 72 | 73 | new_rng, _ = jax.random.split(state.seed) 74 | new_agent_state = types.AgentState(internal_state=next_state, seed=new_rng) 75 | return action, net_out, new_agent_state 76 | -------------------------------------------------------------------------------- /src/pycoworld/default_constants.py: -------------------------------------------------------------------------------- 1 | # Copyright 2023 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | 16 | """Default constants for Pycoworld.""" 17 | 18 | import enum 19 | 20 | 21 | class Tile(enum.IntEnum): 22 | """Available pycoworld tiles with their IDs. 23 | 24 | This IntEnum is only provided for better readability. There is no guarantee 25 | that the raw IDs are not used anywhere in the code and no requirement to use 26 | the enum in cases where the raw IDs are more readable. 27 | """ 28 | FLOOR = 0 29 | FLOOR_R = 1 30 | FLOOR_B = 2 31 | FLOOR_G = 3 32 | 33 | WALL = 4 34 | WALL_R = 5 35 | WALL_B = 6 36 | WALL_G = 7 37 | 38 | PLAYER = 8 39 | PLAYER_R = 9 40 | PLAYER_B = 10 41 | PLAYER_G = 11 42 | 43 | TERMINAL = 12 44 | TERMINAL_R = 13 45 | TERMINAL_B = 14 46 | TERMINAL_G = 15 47 | 48 | BLOCK = 16 49 | BLOCK_R = 17 50 | BLOCK_B = 18 51 | BLOCK_G = 19 52 | 53 | REWARD = 20 54 | REWARD_R = 21 55 | REWARD_B = 22 56 | REWARD_G = 23 57 | 58 | BIG_REWARD = 24 59 | BIG_REWARD_R = 25 60 | BIG_REWARD_B = 26 61 | BIG_REWARD_G = 27 62 | 63 | HOLE = 28 64 | HOLE_R = 29 65 | HOLE_B = 30 66 | HOLE_G = 31 67 | 68 | SAND = 32 69 | GRASS = 33 70 | LAVA = 34 71 | WATER = 35 72 | 73 | KEY = 36 74 | KEY_R = 37 75 | KEY_B = 38 76 | KEY_G = 39 77 | 78 | DOOR = 40 79 | DOOR_R = 41 80 | DOOR_B = 42 81 | DOOR_G = 43 82 | 83 | SENSOR = 44 84 | SENSOR_R = 45 85 | SENSOR_B = 46 86 | SENSOR_G = 47 87 | 88 | OBJECT = 48 89 | OBJECT_R = 49 90 | OBJECT_B = 50 91 | OBJECT_G = 51 92 | 93 | 94 | # This dict defines which sensor matches which object. 95 | MATCHING_OBJECT_DICT = { 96 | Tile.SENSOR: Tile.OBJECT, 97 | Tile.SENSOR_R: Tile.OBJECT_R, 98 | Tile.SENSOR_G: Tile.OBJECT_G, 99 | Tile.SENSOR_B: Tile.OBJECT_B, 100 | } 101 | 102 | # Impassable tiles: these tiles cannot be traversed by the agent. 103 | IMPASSABLE = [ 104 | Tile.WALL, 105 | Tile.WALL_R, 106 | Tile.WALL_B, 107 | Tile.WALL_G, 108 | Tile.BLOCK, 109 | Tile.BLOCK_R, 110 | Tile.BLOCK_B, 111 | Tile.BLOCK_G, 112 | Tile.DOOR, 113 | Tile.DOOR_R, 114 | Tile.DOOR_B, 115 | Tile.DOOR_G, 116 | ] 117 | 118 | REWARD_LAVA = -10. 119 | REWARD_WATER = -1. 120 | REWARD_SMALL = 1. 121 | REWARD_BIG = 5. 122 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Causal Analysis of Agent Behavior for AI Safety 2 | 3 |

4 | Overview figure 5 |

6 | 7 | This repository provides an implementation of our paper [Causal Analysis of Agent Behavior for AI Safety](https://arxiv.org/abs/2103.03938). 8 | 9 | >As machine learning systems become more powerful they also become increasingly unpredictable and opaque. 10 | Yet, finding human-understandable explanations of how they work is essential for their safe deployment. 11 | This technical report illustrates a methodology for investigating the causal mechanisms that drive the behaviour of artificial agents. 12 | Six use cases are covered, each addressing a typical question an analyst might ask about an agent. 13 | In particular, we show that each question cannot be addressed by pure observation alone, but instead requires conducting experiments with systematically chosen manipulations so as to generate the correct causal evidence. 14 | 15 | The main tool is the "Agent Debugger", which can be used to perform causal interventions on the environment to infer the causal model of an agent. 16 | We currently only support the environment Pycoworld, a 2D gridworld based on the open source game engine [Pycolab](https://github.com/deepmind/pycolab). 17 | 18 | 19 | ## Usage 20 | 21 | To reproduce the experiments of the paper, run the [experiments notebook](https://colab.research.google.com/github/deepmind/agent_debugger/blob/master/colabs/experiments.ipynb). 22 | 23 | 24 | ## Citing this work 25 | 26 | ```bibtex 27 | @article{deletang2021causal, 28 | author = {Gr{\'{e}}goire Del{\'{e}}tang and 29 | Jordi Grau{-}Moya and 30 | Miljan Martic and 31 | Tim Genewein and 32 | Tom McGrath and 33 | Vladimir Mikulik and 34 | Markus Kunesch and 35 | Shane Legg and 36 | Pedro A. Ortega}, 37 | title = {Causal Analysis of Agent Behavior for {AI} Safety}, 38 | journal = {arXiv:2103.03938}, 39 | year = {2021}, 40 | } 41 | ``` 42 | 43 | 44 | ## License and disclaimer 45 | 46 | Copyright 2022 DeepMind Technologies Limited 47 | 48 | All software is licensed under the Apache License, Version 2.0 (Apache 2.0); 49 | you may not use this file except in compliance with the Apache 2.0 license. 50 | You may obtain a copy of the Apache 2.0 license at: 51 | https://www.apache.org/licenses/LICENSE-2.0 52 | 53 | All other materials are licensed under the Creative Commons Attribution 4.0 54 | International License (CC-BY). You may obtain a copy of the CC-BY license at: 55 | https://creativecommons.org/licenses/by/4.0/legalcode 56 | 57 | Unless required by applicable law or agreed to in writing, all software and 58 | materials distributed here under the Apache 2.0 or CC-BY licenses are 59 | distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, 60 | either express or implied. See the licenses for the specific language governing 61 | permissions and limitations under those licenses. 62 | 63 | This is not an official Google product. 64 | -------------------------------------------------------------------------------- /src/pycoworld/levels/apples.py: -------------------------------------------------------------------------------- 1 | # Copyright 2023 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | 16 | """Apples level.""" 17 | 18 | import numpy as np 19 | 20 | from agent_debugger.src.pycoworld import default_constants 21 | from agent_debugger.src.pycoworld.levels import base_level 22 | from agent_debugger.src.pycoworld.levels import utils 23 | 24 | Tile = default_constants.Tile 25 | 26 | 27 | class ApplesLevel(base_level.PycoworldLevel): 28 | """Level where the goal is to pick up the reward.""" 29 | 30 | def __init__( 31 | self, 32 | start_type: str = 'full_room', 33 | height: int = 8, 34 | width: int = 8, 35 | ) -> None: 36 | """Initializes the level. 37 | 38 | Args: 39 | start_type: The type of random initialization: 'full_room' forces random 40 | initialization using all cells in the room, 'corner' forces random 41 | initialization of the player around one corner of the room and the 42 | reward is initialized close to the other corner. 43 | height: Height of the grid. 44 | width: Width of the grid. 45 | """ 46 | if start_type not in ['full_room', 'corner']: 47 | raise ValueError('Unrecognised start type.') 48 | self._start_type = start_type 49 | self._height = height 50 | self._width = width 51 | 52 | def foreground_and_background( 53 | self, 54 | rng: np.ndarray, 55 | ) -> tuple[np.ndarray, np.ndarray]: 56 | """See base class.""" 57 | foreground = utils.room(self._height, self._width) 58 | 59 | # Sample player's and reward's initial position in the interior of the room. 60 | if self._start_type == 'full_room': 61 | player_pos, reward_pos = utils.sample_positions( 62 | rng, 63 | height_range=(1, self._height - 1), 64 | width_range=(1, self._width - 1), 65 | number_samples=2, 66 | replace=False) 67 | elif self._start_type == 'corner': 68 | # Sample positions in the top left quadrant for the player 69 | player_pos = utils.sample_positions( 70 | rng, 71 | height_range=(1, self._height // 2), 72 | width_range=(1, self._width // 2), 73 | number_samples=1)[0] 74 | # Sample positions in the bottom right quadrant for the reward 75 | reward_pos = utils.sample_positions( 76 | rng, 77 | height_range=((self._height + 1) // 2, self._height - 1), 78 | width_range=((self._width + 1) // 2, self._width - 1), 79 | number_samples=1)[0] 80 | 81 | foreground[player_pos] = Tile.PLAYER 82 | foreground[reward_pos] = Tile.REWARD 83 | 84 | background = np.full_like(foreground, Tile.FLOOR) 85 | background[reward_pos] = Tile.TERMINAL_R 86 | return foreground, background 87 | -------------------------------------------------------------------------------- /src/pycoworld/levels/key_door.py: -------------------------------------------------------------------------------- 1 | # Copyright 2023 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | 16 | """Key Door level. 17 | 18 | In this level there is a key and a door. Behind the door there is a reward. The 19 | agent must learn to collect the key, open the door and obtain the reward. 20 | """ 21 | 22 | import jax 23 | import jax.numpy as jnp 24 | import numpy as np 25 | 26 | from agent_debugger.src.pycoworld import default_constants 27 | from agent_debugger.src.pycoworld.levels import base_level 28 | from agent_debugger.src.pycoworld.levels import utils 29 | 30 | _VALID_DOORS = ['random', 'closed', 'open'] 31 | 32 | 33 | class KeyDoorLevel(base_level.PycoworldLevel): 34 | """Level in which the agent has to unlock a door to collect the reward.""" 35 | 36 | def __init__(self, door_type: str = 'random') -> None: 37 | """Initializes the level. 38 | 39 | Args: 40 | door_type: The way we sample the door state (open or closed). Must be in 41 | {'random', 'closed', 'open'}, otherwise a ValueError is raised. 42 | """ 43 | if door_type not in _VALID_DOORS: 44 | raise ValueError( 45 | f'Argument door_type has an incorrect value. Expected one ' 46 | f'of {_VALID_DOORS}, but got {door_type} instead.') 47 | self._door_type = door_type 48 | 49 | def foreground_and_background( 50 | self, 51 | rng: jnp.ndarray, 52 | ) -> tuple[np.ndarray, np.ndarray]: 53 | """Returns a tuple with the level foreground and background. 54 | 55 | We first determine the state of the door and use the tile associated with 56 | it (closed and open doors don't have the same id). 57 | Then, we sample the random player and key positions and add their tiles to 58 | the board. The background is full of floor tiles and a terminal event 59 | beneath the reward tile, so that the episode is terminated when the agent 60 | gets the reward. 61 | 62 | Args: 63 | rng: The jax random seed. Standard name for jax random seeds. 64 | """ 65 | if self._door_type == 'closed': 66 | door_state = 'closed' 67 | elif self._door_type == 'open': 68 | door_state = 'open' 69 | elif self._door_type == 'random': 70 | rng, rng1 = jax.random.split(rng) 71 | if jax.random.uniform(rng1) < 0.5: 72 | door_state = 'open' 73 | else: 74 | door_state = 'closed' 75 | 76 | foreground = np.array([[4, 4, 4, 4, 4, 0, 0], [4, 0, 0, 0, 4, 0, 0], 77 | [4, 0, 0, 0, 4, 4, 4], [4, 0, 0, 0, 0, 20, 4], 78 | [4, 4, 4, 4, 4, 4, 4]]) 79 | 80 | if door_state == 'closed': 81 | foreground[3, 4] = default_constants.Tile.DOOR_R 82 | 83 | player_pos, key_pos = utils.sample_positions( 84 | rng, 85 | height_range=(1, 4), 86 | width_range=(1, 4), 87 | number_samples=2, 88 | replace=False) 89 | 90 | foreground[player_pos] = default_constants.Tile.PLAYER 91 | foreground[key_pos] = default_constants.Tile.KEY_R 92 | 93 | background = np.zeros(foreground.shape, dtype=int) 94 | background[3, 5] = default_constants.Tile.TERMINAL_R 95 | return foreground, background 96 | -------------------------------------------------------------------------------- /src/pycoworld/levels/utils.py: -------------------------------------------------------------------------------- 1 | # Copyright 2023 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | 16 | """Utilities for constructing pycoworld levels.""" 17 | 18 | from typing import Optional, Sequence 19 | 20 | import chex 21 | import jax 22 | import numpy as np 23 | 24 | from agent_debugger.src.pycoworld import default_constants 25 | 26 | Tile = default_constants.Tile 27 | 28 | # A 2d slice for masking out the walls of a room. 29 | _INTERIOR = (slice(1, -1),) * 2 30 | 31 | 32 | def room( 33 | height: int, 34 | width: int, 35 | floor_tile: Tile = Tile.FLOOR, 36 | wall_tile: Tile = Tile.WALL, 37 | ) -> np.ndarray: 38 | """Returns a pycoworld representation of a room (floor surrounded by walls). 39 | 40 | Args: 41 | height: Height of the environment. 42 | width: Width of the environment. 43 | floor_tile: Tile to use for the floor. 44 | wall_tile: Tile to use for the wall. 45 | 46 | Returns: 47 | The array representation of the room. 48 | """ 49 | room_array = np.full([height, width], wall_tile, dtype=np.uint8) 50 | room_array[_INTERIOR] = floor_tile 51 | return room_array 52 | 53 | 54 | def sample_positions( 55 | rng: np.ndarray, 56 | height_range: tuple[int, int], 57 | width_range: tuple[int, int], 58 | number_samples: int, 59 | replace: bool = True, 60 | exclude: Optional[np.ndarray] = None, 61 | ) -> Sequence[tuple[int, int]]: 62 | """Uniformly samples random positions in a 2d grid. 63 | 64 | Args: 65 | rng: Jax random key to use for sampling. 66 | height_range: Range (min, 1+max) in the height of the grid to sample from. 67 | width_range: Range (min, 1+max) in the width of the grid to sample from. 68 | number_samples: Number of positions to sample. 69 | replace: Whether to sample with replacement. 70 | exclude: Array of positions to exclude from the sampling. Each row of the 71 | array represents the x,y coordinates of one position. 72 | 73 | Returns: 74 | A sequence of x,y index tuples which can be used to index a 2d numpy array. 75 | 76 | Raises: 77 | ValueError: if more positions are requested than are available when sampling 78 | without replacement. 79 | """ 80 | height = height_range[1] - height_range[0] 81 | width = width_range[1] - width_range[0] 82 | offset = np.array([height_range[0], width_range[0]]) 83 | 84 | choices = np.arange(height * width) 85 | if exclude is not None: 86 | exclude = np.asarray(exclude) # Allow the user to pass any array-like type. 87 | chex.assert_shape(exclude, (None, 2)) 88 | 89 | exclude_offset = exclude - offset 90 | exclude_indices = np.ravel_multi_index(exclude_offset.T, (height, width)) 91 | mask = np.ones_like(choices, dtype=bool) 92 | mask[exclude_indices] = False 93 | choices = choices[mask] 94 | 95 | flat_indices = jax.random.choice( 96 | rng, choices, (number_samples,), replace=replace) 97 | positions_offset = np.unravel_index(flat_indices, (height, width)) 98 | positions = positions_offset + offset[:, np.newaxis] 99 | return list(zip(*positions)) 100 | -------------------------------------------------------------------------------- /src/impala_net.py: -------------------------------------------------------------------------------- 1 | # Copyright 2023 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | 16 | """Neural network used in Impala trained agents.""" 17 | 18 | from typing import NamedTuple, Optional, Sequence, Union 19 | 20 | import dm_env 21 | import haiku as hk 22 | import jax.nn 23 | import jax.numpy as jnp 24 | 25 | NetState = Union[jnp.ndarray, hk.LSTMState] 26 | 27 | 28 | # Kept as a NamedTuple as jax does not support dataclasses. 29 | class NetOutput(NamedTuple): 30 | """Dataclass to define network outputs.""" 31 | policy_logits: jnp.ndarray 32 | value: jnp.ndarray 33 | 34 | 35 | class RecurrentConvNet(hk.RNNCore): 36 | """A class for Impala nets. 37 | 38 | Architecture: 39 | MLP torso -> LSTM -> MLP head -> (linear policy logits, linear value) 40 | 41 | If initialised with lstm_width=0, skips the LSTM layer. 42 | """ 43 | 44 | def __init__( 45 | self, 46 | num_actions: int, 47 | conv_widths: Sequence[int], 48 | conv_kernels: Sequence[int], 49 | padding: str, 50 | torso_widths: Sequence[int], 51 | lstm_width: Optional[int], 52 | head_widths: Sequence[int], 53 | name: str = None, 54 | ) -> None: 55 | """Initializes the impala net.""" 56 | super().__init__(name=name) 57 | self._num_actions = num_actions 58 | self._torso_widths = torso_widths 59 | self._head_widths = head_widths 60 | self._core = hk.LSTM(lstm_width) if lstm_width else None 61 | 62 | conv_layers = [] 63 | for width, kernel_size in zip(conv_widths, conv_kernels): 64 | layer = hk.Conv2D( 65 | width, kernel_shape=[kernel_size, kernel_size], padding=padding) 66 | conv_layers += [layer, jax.nn.relu] 67 | self._conv_net = hk.Sequential(conv_layers + [hk.Flatten()]) 68 | 69 | def initial_state(self, batch_size: int) -> NetState: 70 | """Returns a fresh hidden state for the LSTM core.""" 71 | return self._core.initial_state(batch_size) if self._core else jnp.zeros(()) 72 | 73 | def __call__( 74 | self, 75 | x: dm_env.TimeStep, 76 | state: NetState, 77 | ) -> tuple[NetOutput, NetState]: 78 | """Steps the net, applying a forward pass of the neural network.""" 79 | # Apply torso. 80 | observation = x.observation['board'].astype(dtype=jnp.float32) / 255 81 | observation = jnp.expand_dims(observation, axis=0) 82 | output = self._torso(observation) 83 | 84 | if self._core is not None: 85 | output, state = self._core(output, state) 86 | 87 | policy_logits, value = self._head(output) 88 | return NetOutput(policy_logits=policy_logits[0], value=value[0]), state 89 | 90 | def _head(self, activations: jnp.ndarray) -> jnp.ndarray: 91 | """Returns new activations after applying the head network.""" 92 | pre_outputs = hk.nets.MLP(self._head_widths)(activations) 93 | policy_logits = hk.Linear(self._num_actions)(pre_outputs) 94 | value = hk.Linear(1)(pre_outputs) 95 | value = jnp.squeeze(value, axis=-1) 96 | return policy_logits, value 97 | 98 | def _torso(self, inputs: jnp.ndarray) -> jnp.ndarray: 99 | """Returns activations after applying the torso to the inputs.""" 100 | return hk.Sequential([self._conv_net, 101 | hk.nets.MLP(self._torso_widths)])( 102 | inputs) 103 | -------------------------------------------------------------------------------- /src/agent_debugger/default_interventions.py: -------------------------------------------------------------------------------- 1 | # Copyright 2023 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | 16 | """Base interventions for the debugger.""" 17 | 18 | import copy 19 | import dataclasses 20 | from typing import Union 21 | 22 | import jax 23 | 24 | from agent_debugger.src.agent_debugger import node as node_lib 25 | from agent_debugger.src.agent_debugger import types 26 | from agent_debugger.src.pycoworld import serializable_environment 27 | 28 | 29 | class DefaultInterventions: 30 | """Object containing the default debugger interventions.""" 31 | 32 | def __init__( 33 | self, 34 | env: serializable_environment.SerializableEnvironment, 35 | ) -> None: 36 | """Initializes the default interventions object. 37 | 38 | Args: 39 | env: The environment to be able to reset after a seed intervention. 40 | """ 41 | self._env = env 42 | 43 | def change_agent_seed( 44 | self, 45 | node: node_lib.Node, 46 | seed: Union[int, types.Rng], 47 | ) -> node_lib.Node: 48 | """Intervenes to replace the agent seed, used for acting. 49 | 50 | Args: 51 | node: The node to intervene on. 52 | seed: The new seed to put to the agent state of the node. Can be int or 53 | jax.random.PRNGKey directly. 54 | 55 | Returns: 56 | A new node, whose agent seed has been updated. 57 | """ 58 | if isinstance(seed, int): 59 | seed = jax.random.PRNGKey(seed) 60 | new_agent_state = node.agent_state._replace(seed=seed) 61 | return dataclasses.replace(node, agent_state=new_agent_state) 62 | 63 | def change_env_seed( 64 | self, 65 | node: node_lib.Node, 66 | seed: int, 67 | ) -> node_lib.Node: 68 | """Intervenes to replace the seed of the environment. 69 | 70 | The state of the environment must contain an attribute 'seed' to be able to 71 | use this intervention. 72 | 73 | Args: 74 | node: The node to intervene on. 75 | seed: The new seed to put to the environment state of the node. 76 | 77 | Returns: 78 | A new node, which environment state has been modified. 79 | """ 80 | # Create the new state. 81 | state = copy.deepcopy(node.env_state) 82 | state.seed = seed 83 | self._env.set_state(state) 84 | 85 | new_timestep = self._env.reset() 86 | new_state = self._env.get_state() 87 | return dataclasses.replace( 88 | node, env_state=new_state, last_timestep=new_timestep) 89 | 90 | def change_agent_next_actions( 91 | self, 92 | node: node_lib.Node, 93 | forced_next_actions: list[types.Action], 94 | ) -> node_lib.Node: 95 | """Changes the next actions of the agent at a given node. 96 | 97 | This intervention allows the user to change the N next actions of the agent, 98 | not only the next one. When stepping the debugger from this node, the list 99 | for the next node is actualised by removing the first action (since it has 100 | just been executed). 101 | 102 | Args: 103 | node: The node to intervene on. 104 | forced_next_actions: The next actions to be taken by the agent. 105 | 106 | Returns: 107 | A new node, which, when stepped from, will force the N next actions taken 108 | by the agent. 109 | """ 110 | return dataclasses.replace(node, forced_next_actions=forced_next_actions) 111 | -------------------------------------------------------------------------------- /src/pycoworld/levels/red_green_apples.py: -------------------------------------------------------------------------------- 1 | # Copyright 2023 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | 16 | """Red-Green Apples level. 17 | 18 | In this level there are two rooms, each containing two apples, one is red and 19 | the other is green. On each episode only one of the rooms is opened (and the 20 | other is closed by a door) and the agent's goal is to enter the opened room 21 | and reach the green apple. 22 | """ 23 | 24 | from typing import Any, MutableMapping 25 | 26 | import jax 27 | import numpy as np 28 | from pycolab import engine 29 | from pycolab import things as plab_things 30 | 31 | from agent_debugger.src.pycoworld import default_sprites_and_drapes 32 | from agent_debugger.src.pycoworld.levels import base_level 33 | 34 | 35 | class RedGreenApplesLevel(base_level.PycoworldLevel): 36 | """Red-green apples level.""" 37 | 38 | def __init__(self, red_lover: bool = False): 39 | """Initializes the level. 40 | 41 | Args: 42 | red_lover: Whether the red marble gives a positive or a negative reward. 43 | True means a positive reward. 44 | """ 45 | super().__init__() 46 | self._red_lover = red_lover 47 | 48 | def foreground_and_background( 49 | self, 50 | rng: np.ndarray, 51 | ) -> tuple[np.ndarray, np.ndarray]: 52 | """See base class.""" 53 | if jax.random.uniform(rng) < 0.5: 54 | # Room on the left is open. 55 | foreground = np.array([[4, 4, 4, 4, 4, 4, 4], [4, 0, 21, 4, 21, 0, 4], 56 | [4, 0, 23, 4, 23, 0, 4], [4, 0, 4, 4, 4, 42, 4], 57 | [4, 0, 0, 8, 0, 0, 4], [4, 4, 4, 4, 4, 4, 4]]) 58 | else: 59 | # Room on the right is open. 60 | foreground = np.array([[4, 4, 4, 4, 4, 4, 4], [4, 0, 21, 4, 21, 0, 4], 61 | [4, 0, 23, 4, 23, 0, 4], [4, 42, 4, 4, 4, 0, 4], 62 | [4, 0, 0, 8, 0, 0, 4], [4, 4, 4, 4, 4, 4, 4]]) 63 | 64 | background = np.array([ 65 | [0, 0, 0, 0, 0, 0, 0], 66 | [0, 0, 13, 0, 13, 0, 0], 67 | [0, 0, 13, 0, 13, 0, 0], 68 | [0, 0, 0, 0, 0, 0, 0], 69 | [0, 0, 0, 0, 0, 0, 0], 70 | [0, 0, 0, 0, 0, 0, 0], 71 | ]) 72 | return foreground, background 73 | 74 | def sample_game(self, rng: np.ndarray) -> engine.Engine: 75 | """See base class.""" 76 | foreground, background = self.foreground_and_background(rng) 77 | drapes = base_level.default_drapes() 78 | if self._red_lover: 79 | drapes[chr(23)] = BadAppleDrape 80 | drapes[chr(21)] = default_sprites_and_drapes.RewardDrape 81 | else: 82 | drapes[chr(21)] = BadAppleDrape 83 | drapes[chr(23)] = default_sprites_and_drapes.RewardDrape 84 | return base_level.make_pycolab_engine(foreground, background, drapes=drapes) 85 | 86 | 87 | class BadAppleDrape(plab_things.Drape): 88 | """A drape for a bad apple. 89 | 90 | Collecting the bad apple punishes the agent with a negative reward of -1. 91 | """ 92 | 93 | def update( 94 | self, 95 | actions: int, 96 | board: np.ndarray, 97 | layers: MutableMapping[str, np.ndarray], 98 | backdrop: np.ndarray, 99 | things: MutableMapping[str, Any], 100 | the_plot: engine.plot.Plot 101 | ) -> None: 102 | ypos, xpos = things[chr(8)].position # Get agent's position. 103 | 104 | # If the agent is in the same position as the bad apple give -1 reward and 105 | # consume the bad apple. 106 | if self.curtain[ypos, xpos]: 107 | the_plot.add_reward(-1.) 108 | 109 | # Remove bad apple. 110 | self.curtain[ypos, xpos] = False 111 | -------------------------------------------------------------------------------- /src/pycoworld/levels/grass_sand.py: -------------------------------------------------------------------------------- 1 | # Copyright 2023 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | 16 | """Grass sand level.""" 17 | from typing import Optional 18 | 19 | import jax 20 | import numpy as np 21 | 22 | from agent_debugger.src.pycoworld import default_constants 23 | from agent_debugger.src.pycoworld.levels import base_level 24 | 25 | Tile = default_constants.Tile 26 | 27 | 28 | class GrassSandLevel(base_level.PycoworldLevel): 29 | """The grass sand level. 30 | 31 | The goal position causally depends on the type of world. Similarly, the floor 32 | tiles (grass or sand) also causally depend on the type of world. Therefore, 33 | type of floor and reward position are correlated, but there is no causal link 34 | between them. 35 | """ 36 | 37 | def __init__( 38 | self, 39 | corr: float = 1.0, 40 | above: Optional[np.ndarray] = None, 41 | below: Optional[np.ndarray] = None, 42 | reward_pos: Optional[tuple[np.ndarray, np.ndarray]] = None, 43 | double_terminal_event: Optional[bool] = False, 44 | ) -> None: 45 | """Initializes the level. 46 | 47 | Args: 48 | corr: "correlation" between reward position and world. 1.0 means that the 49 | reward position will always be on "north" for sand world and on "south" 50 | for grass world. A value of 0.5 corresponds to fully randomized position 51 | independent of the world. 52 | above: a map to use. 53 | below: a background to use. 54 | reward_pos: position of the rewards. 55 | double_terminal_event: whether to have a terminal event on both sides of 56 | the maze or only beneath the reward. 57 | """ 58 | if not 0.0 <= corr <= 1.0: 59 | raise ValueError('corr variable must be a float between 0.0 and 1.0') 60 | 61 | self._corr = corr 62 | if above is None: 63 | above = np.array([[0, 0, 0, 4, 4, 4], [4, 4, 4, 4, 99, 4], 64 | [4, 4, 4, 4, 99, 4], [4, 8, 99, 99, 99, 4], 65 | [4, 4, 4, 4, 99, 4], [4, 4, 4, 4, 99, 4], 66 | [0, 0, 0, 4, 4, 4]]) 67 | self._above = above 68 | if below is None: 69 | below = np.full_like(above, Tile.FLOOR) 70 | self._below = below 71 | 72 | if reward_pos is None: 73 | self._reward_pos_top = np.array([1, 4]) 74 | self._reward_pos_bottom = np.array([5, 4]) 75 | else: 76 | self._reward_pos_top, self._reward_pos_bottom = reward_pos 77 | 78 | self._double_terminal_event = double_terminal_event 79 | 80 | def foreground_and_background( 81 | self, 82 | rng: np.ndarray, 83 | ) -> tuple[np.ndarray, np.ndarray]: 84 | """See base class.""" 85 | rng1, rng2 = jax.random.split(rng, 2) 86 | # Do not modify the base maps during sampling 87 | sampled_above = self._above.copy() 88 | sampled_below = self._below.copy() 89 | 90 | # Select world. 91 | if jax.random.uniform(rng1) < 0.5: 92 | world = 'sand' 93 | floor_type = Tile.SAND 94 | else: 95 | world = 'grass' 96 | floor_type = Tile.GRASS 97 | 98 | # Sample reward location depending on corr variable. 99 | if jax.random.uniform(rng2) <= self._corr: 100 | # Standard reward position for both worlds. 101 | if world == 'sand': 102 | used_reward_pos = self._reward_pos_top 103 | elif world == 'grass': 104 | used_reward_pos = self._reward_pos_bottom 105 | else: 106 | # Alternative reward position for both worlds. 107 | if world == 'sand': 108 | used_reward_pos = self._reward_pos_bottom 109 | elif world == 'grass': 110 | used_reward_pos = self._reward_pos_top 111 | 112 | if self._double_terminal_event: 113 | # Put terminal events everywhere first. 114 | sampled_above[self._reward_pos_top[0], 115 | self._reward_pos_top[1]] = Tile.TERMINAL_R 116 | sampled_above[self._reward_pos_bottom[0], 117 | self._reward_pos_bottom[1]] = Tile.TERMINAL_R 118 | # Draw the reward ball. 119 | sampled_above[used_reward_pos[0], used_reward_pos[1]] = Tile.REWARD 120 | 121 | # Substitute tiles with code 99 with the corresponding floor_type code. 122 | floor_tiles = np.argwhere(sampled_above == 99) 123 | for tile in floor_tiles: 124 | sampled_above[tile[0], tile[1]] = floor_type 125 | 126 | if self._double_terminal_event: 127 | # Add terminal events on both sides to the below. 128 | sampled_below[self._reward_pos_top[0], 129 | self._reward_pos_top[1]] = Tile.TERMINAL_R 130 | sampled_below[self._reward_pos_bottom[0], 131 | self._reward_pos_bottom[1]] = Tile.TERMINAL_R 132 | else: 133 | sampled_below[used_reward_pos[0], used_reward_pos[1]] = Tile.TERMINAL_R 134 | 135 | return sampled_above, sampled_below 136 | -------------------------------------------------------------------------------- /src/pycoworld/levels/base_level.py: -------------------------------------------------------------------------------- 1 | # Copyright 2023 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | 16 | """Level config.""" 17 | 18 | import abc 19 | from typing import Any, Mapping, Optional, Sequence, Union 20 | 21 | import numpy as np 22 | from pycolab import ascii_art 23 | from pycolab import engine 24 | from pycolab import things as pycolab_things 25 | 26 | from agent_debugger.src.pycoworld import default_constants 27 | from agent_debugger.src.pycoworld import default_sprites_and_drapes 28 | 29 | _ACTIONS = range(4) 30 | Tile = default_constants.Tile 31 | 32 | 33 | def default_sprites() -> dict[str, Any]: 34 | """Returns the default mapping of characters to sprites in the levels.""" 35 | return {chr(Tile.PLAYER): default_sprites_and_drapes.PlayerSprite} 36 | 37 | 38 | def default_drapes() -> dict[str, Any]: 39 | """Returns the default mapping of characters to drapes in the levels.""" 40 | return { 41 | chr(Tile.WALL): default_sprites_and_drapes.WallDrape, 42 | chr(Tile.WALL_R): default_sprites_and_drapes.WallDrape, 43 | chr(Tile.WALL_B): default_sprites_and_drapes.WallDrape, 44 | chr(Tile.WALL_G): default_sprites_and_drapes.WallDrape, 45 | chr(Tile.BLOCK): default_sprites_and_drapes.BoxDrape, 46 | chr(Tile.BLOCK_R): default_sprites_and_drapes.BoxDrape, 47 | chr(Tile.BLOCK_B): default_sprites_and_drapes.BoxDrape, 48 | chr(Tile.BLOCK_G): default_sprites_and_drapes.BoxDrape, 49 | chr(Tile.REWARD): default_sprites_and_drapes.RewardDrape, 50 | chr(Tile.REWARD_R): default_sprites_and_drapes.RewardDrape, 51 | chr(Tile.REWARD_B): default_sprites_and_drapes.RewardDrape, 52 | chr(Tile.REWARD_G): default_sprites_and_drapes.RewardDrape, 53 | chr(Tile.BIG_REWARD): default_sprites_and_drapes.BigRewardDrape, 54 | chr(Tile.BIG_REWARD_R): default_sprites_and_drapes.BigRewardDrape, 55 | chr(Tile.BIG_REWARD_B): default_sprites_and_drapes.BigRewardDrape, 56 | chr(Tile.BIG_REWARD_G): default_sprites_and_drapes.BigRewardDrape, 57 | chr(Tile.KEY): default_sprites_and_drapes.ObjectDrape, 58 | chr(Tile.KEY_R): default_sprites_and_drapes.ObjectDrape, 59 | chr(Tile.KEY_B): default_sprites_and_drapes.ObjectDrape, 60 | chr(Tile.KEY_G): default_sprites_and_drapes.ObjectDrape, 61 | chr(Tile.DOOR): default_sprites_and_drapes.DoorDrape, 62 | chr(Tile.DOOR_R): default_sprites_and_drapes.DoorDrape, 63 | chr(Tile.DOOR_B): default_sprites_and_drapes.DoorDrape, 64 | chr(Tile.DOOR_G): default_sprites_and_drapes.DoorDrape, 65 | chr(Tile.OBJECT): default_sprites_and_drapes.ObjectDrape, 66 | chr(Tile.OBJECT_R): default_sprites_and_drapes.ObjectDrape, 67 | chr(Tile.OBJECT_B): default_sprites_and_drapes.ObjectDrape, 68 | chr(Tile.OBJECT_G): default_sprites_and_drapes.ObjectDrape, 69 | } 70 | 71 | 72 | def default_schedule() -> list[str]: 73 | """Returns the default update schedule of sprites and drapes in the levels.""" 74 | return [ 75 | chr(Tile.PLAYER), # PlayerSprite 76 | chr(Tile.WALL), 77 | chr(Tile.WALL_R), 78 | chr(Tile.WALL_B), 79 | chr(Tile.WALL_G), 80 | chr(Tile.REWARD), 81 | chr(Tile.REWARD_R), 82 | chr(Tile.REWARD_B), 83 | chr(Tile.REWARD_G), 84 | chr(Tile.BIG_REWARD), 85 | chr(Tile.BIG_REWARD_R), 86 | chr(Tile.BIG_REWARD_B), 87 | chr(Tile.BIG_REWARD_G), 88 | chr(Tile.KEY), 89 | chr(Tile.KEY_R), 90 | chr(Tile.KEY_B), 91 | chr(Tile.KEY_G), 92 | chr(Tile.OBJECT), 93 | chr(Tile.OBJECT_R), 94 | chr(Tile.OBJECT_B), 95 | chr(Tile.OBJECT_G), 96 | chr(Tile.DOOR), 97 | chr(Tile.DOOR_R), 98 | chr(Tile.DOOR_B), 99 | chr(Tile.DOOR_G), 100 | chr(Tile.BLOCK), 101 | chr(Tile.BLOCK_R), 102 | chr(Tile.BLOCK_B), 103 | chr(Tile.BLOCK_G), 104 | ] 105 | 106 | 107 | def _numpy_to_str(array: np.ndarray) -> list[str]: 108 | """Converts numpy array into a list of strings. 109 | 110 | Args: 111 | array: a 2-D np.darray of np.uint8. 112 | 113 | Returns: 114 | A list of strings of equal length, corresponding to the entries in A. 115 | """ 116 | return [''.join(map(chr, row.tolist())) for row in array] 117 | 118 | 119 | def make_pycolab_engine( 120 | foreground: np.ndarray, 121 | background: Union[np.ndarray, int], 122 | sprites: Optional[Mapping[str, pycolab_things.Sprite]] = None, 123 | drapes: Optional[Mapping[str, pycolab_things.Drape]] = None, 124 | update_schedule: Optional[Sequence[str]] = None, 125 | rng: Optional[Any] = None, 126 | ) -> engine.Engine: 127 | """Builds and returns a pycoworld game engine. 128 | 129 | Args: 130 | foreground: Array of foreground tiles. 131 | background: Array of background tiles or a single tile to use as the 132 | background everywhere. 133 | sprites: Pycolab sprites. See pycolab.ascii_art.ascii_art_to_game for more 134 | information. 135 | drapes: Pycolab drapes. 136 | update_schedule: Update schedule for sprites and drapes. 137 | rng: Random key to use for pycolab. 138 | 139 | Returns: 140 | A pycolab engine with the pycoworld game. 141 | """ 142 | sprites = sprites if sprites is not None else default_sprites() 143 | drapes = drapes if drapes is not None else default_drapes() 144 | if update_schedule is None: 145 | update_schedule = default_schedule() 146 | 147 | # The pycolab engine constructor requires arrays of strings 148 | above_str = _numpy_to_str(foreground) 149 | below_str = _numpy_to_str(background) 150 | 151 | pycolab_engine = ascii_art.ascii_art_to_game( 152 | above_str, below_str, sprites, drapes, update_schedule=update_schedule) 153 | # Pycolab does not allow to add a global seed in the engine constructor. 154 | # Therefore, we have to set it manually. 155 | pycolab_engine._the_plot['rng'] = rng # pylint: disable=protected-access 156 | return pycolab_engine 157 | 158 | 159 | class PycoworldLevel(abc.ABC): 160 | """Abstract class representing a pycoworld level. 161 | 162 | A pycoworld level captures all the data that is required to define the level 163 | and implements a function that returns a pycolab game engine. 164 | """ 165 | 166 | @abc.abstractmethod 167 | def foreground_and_background( 168 | self, 169 | rng: np.ndarray, 170 | ) -> tuple[np.ndarray, np.ndarray]: 171 | """Generates the foreground and background arrays of the level.""" 172 | 173 | def sample_game(self, rng: np.ndarray) -> engine.Engine: 174 | """Samples and returns a game from this level. 175 | 176 | A level may contain random elements (e.g. a random goal location). This 177 | function samples and returns a particular game from this level. The base 178 | version calls foreground_and_background and constructs a pycolab engine 179 | using the arrays. In order to use custom sprites, drapes, and update 180 | schedule override this method and use make_pycolab_engine to create an 181 | engine. 182 | 183 | Args: 184 | rng: Random key to use for sampling. 185 | 186 | Returns: 187 | A pycolab game engine for the sampled game. 188 | """ 189 | foreground, background = self.foreground_and_background(rng) 190 | return make_pycolab_engine(foreground, background) 191 | 192 | @property 193 | def actions(self) -> Sequence[int]: 194 | return _ACTIONS 195 | -------------------------------------------------------------------------------- /src/agent_debugger/debugger.py: -------------------------------------------------------------------------------- 1 | # Copyright 2023 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | 16 | """Debugger object, main interface of the framework.""" 17 | 18 | import contextlib 19 | import operator 20 | from typing import Any, Callable, Optional, Sequence 21 | 22 | import jax.random as jrandom 23 | 24 | from agent_debugger.src.agent_debugger import agent as debugger_agent 25 | from agent_debugger.src.agent_debugger import default_interventions 26 | from agent_debugger.src.agent_debugger import node as node_lib 27 | from agent_debugger.src.pycoworld import serializable_environment 28 | 29 | 30 | class NotReachedError(RuntimeError): 31 | """Raised when a rollout did not reach the required point.""" 32 | 33 | 34 | class Debugger: 35 | """Debugger to analyze agents in an environment step by step. 36 | 37 | The debugger can be used to step an agent in an environment, intervene on the 38 | agent and environment, and extract information. The debugger object is the 39 | main API of the Agent Debugger. 40 | 41 | Attributes: 42 | interventions: The interventions module to intervene on nodes. 43 | """ 44 | 45 | def __init__( 46 | self, 47 | agent: debugger_agent.DebuggerAgent, 48 | env: serializable_environment.SerializableEnvironment, 49 | ) -> None: 50 | """Initializes the debugger with an agent and an environment.""" 51 | self._agent = agent 52 | self._env = env 53 | self._rollout_mode = False 54 | self._is_first_rollout_step = False 55 | 56 | self.interventions = default_interventions.DefaultInterventions(self._env) 57 | 58 | def get_root_node(self, seed: int = 1) -> node_lib.Node: 59 | """Returns a root node, containing an initial world state. 60 | 61 | Args: 62 | seed: The seed to use to sample the world state. 63 | """ 64 | rng = jrandom.PRNGKey(seed=seed) 65 | timestep = self._env.reset() 66 | env_state = self._env.get_state() 67 | agent_state = self._agent.initial_state(rng) 68 | return node_lib.Node( 69 | agent_state=agent_state, 70 | env_state=env_state, 71 | last_action=None, 72 | last_timestep=timestep) 73 | 74 | @contextlib.contextmanager 75 | def rollout_mode(self): 76 | """Context manager for the rollout mode. 77 | 78 | Being in rollout mode means some performance optimisation will be done when 79 | we step the agent and the environment together. Specifically, to avoid 80 | getting and setting their states at each step of the rollout rather than 81 | just once at the beginning, we pass an option to deactivate it if the 82 | rollout mode is on. 83 | In practice, this mode reduces the compute time of debugger.get_rollout 84 | while returning exactly the same result. 85 | 86 | Yields: 87 | None. The method is used inside the same object, so no need to yield any 88 | reference. 89 | """ 90 | self._rollout_mode = True 91 | self._is_first_rollout_step = True 92 | try: 93 | yield None 94 | finally: 95 | self._rollout_mode = False 96 | self._is_first_rollout_step = False 97 | 98 | def step_forward(self, node: node_lib.Node) -> node_lib.Node: 99 | """Steps from a world state, using the environment and agent dynamics. 100 | 101 | Args: 102 | node: The node to step from. 103 | 104 | Returns: 105 | The new node obtained after stepping. 106 | """ 107 | action, agent_output, new_agent_state = self._agent.step( 108 | node.last_timestep, node.agent_state) 109 | 110 | if node.forced_next_actions: 111 | action = node.forced_next_actions[0] 112 | 113 | if not self._rollout_mode or self._is_first_rollout_step: 114 | self._env.set_state(node.env_state) 115 | timestep = self._env.step(action) 116 | new_env_state = self._env.get_state() 117 | self._is_first_rollout_step = False 118 | 119 | return node_lib.Node( 120 | agent_state=new_agent_state, 121 | env_state=new_env_state, 122 | last_action=action, 123 | last_timestep=timestep, 124 | forced_next_actions=node.forced_next_actions[1:], 125 | last_agent_output=agent_output, 126 | episode_step=node.episode_step + 1) 127 | 128 | def get_rollout( 129 | self, 130 | initial_node: node_lib.Node, 131 | maximum_length: int, 132 | endpoint: Optional[Callable[[node_lib.Node], bool]] = operator.attrgetter( 133 | 'is_terminal' 134 | ), 135 | raise_endpoint_not_reached: bool = False, 136 | ) -> list[node_lib.Node]: 137 | """Returns the list of nodes obtained by stepping from an initial node. 138 | 139 | Args: 140 | initial_node: The initial node to step from. It is the first element of 141 | the returned list. 142 | maximum_length: The maximum length of the returned list. 143 | endpoint: Function that specifies when to stop the rollout. Defaulted to 144 | is_terminal. If None, the rollout stops after maximum_length iterations, 145 | and some intermediate nodes in the list can therefore be terminal. If we 146 | fell upon a terminal node in this case, we just keep on stepping the 147 | environment, which must return a FIRST node, following the dm_env API. 148 | raise_endpoint_not_reached: Whether to raise an error if the endpoint is 149 | not reached within the maximum rollout length. 150 | 151 | Returns: 152 | A list of nodes. 153 | 154 | Raises: 155 | NotReachedError: If raise_endpoint_not_reached = True and the 156 | endpoint was not reached within the maximum rollout length. 157 | ValueError: If raise_endpoint_not_reached=True while endpoint=None. 158 | """ 159 | if endpoint is None and raise_endpoint_not_reached: 160 | raise ValueError( 161 | 'Cannot raise endpoint not reached error when endpoint is None.') 162 | endpoint_reached = False 163 | with self.rollout_mode(): 164 | visited_nodes = [initial_node] 165 | for _ in range(maximum_length - 1): 166 | current_node = self.step_forward(visited_nodes[-1]) 167 | visited_nodes.append(current_node) 168 | if endpoint is not None and endpoint(current_node): 169 | endpoint_reached = True 170 | break 171 | 172 | if raise_endpoint_not_reached and not endpoint_reached: 173 | raise NotReachedError( 174 | 'Rollout did not reach the end point within the maximum length.') 175 | 176 | return visited_nodes 177 | 178 | def get_intervened_rollout( 179 | self, 180 | initial_node: node_lib.Node, 181 | maximum_length: int, 182 | rollout_breakpoint: Callable[[node_lib.Node], bool], 183 | intervention_at_breakpoint: Callable[[node_lib.Node], node_lib.Node], 184 | endpoint: Optional[Callable[[node_lib.Node], bool]] = operator.attrgetter( 185 | 'is_terminal' 186 | ), 187 | raise_breakpoint_not_reached: bool = True, 188 | ) -> list[node_lib.Node]: 189 | """Returns a modified version of get_rollout, with an intervention. 190 | 191 | The user can provide a breakpoint function to branch at a given point in 192 | the rollout. For now it is possible to branch only once in the rollout. 193 | This is a separate function from get_rollout because we must set and unset 194 | the rollout_mode of the PTS during the intervention. 195 | Args: 196 | initial_node: The initial node to step from. It is the first element of 197 | the returned list. 198 | maximum_length: The maximum length of the returned list. 199 | rollout_breakpoint: Function that specifies when to branch in the rollout, 200 | to intervene. If None, we never branch. 201 | intervention_at_breakpoint: The intervention to perform when the 202 | breakpoint is reached. 203 | endpoint: See get_rollout. 204 | raise_breakpoint_not_reached: Whether to raise an error if the breakpoint 205 | is not reached within the maximum rollout length. 206 | 207 | Returns: 208 | A list of nodes, where an intervention has been performed on one of them. 209 | 210 | Raises: 211 | NotReachedError: If raise_breakpoint_not_reached = True and the 212 | breakpoint was not reached within the maximum rollout length. 213 | ValueError: If raise_breakpoint_not_reached=True while endpoint=None. 214 | """ 215 | visited_nodes = self.get_rollout( 216 | initial_node, 217 | maximum_length, 218 | endpoint=rollout_breakpoint, 219 | raise_endpoint_not_reached=raise_breakpoint_not_reached) 220 | visited_nodes[-1] = intervention_at_breakpoint(visited_nodes[-1]) 221 | visited_nodes.extend( 222 | self.get_rollout( 223 | initial_node=visited_nodes[-1], 224 | maximum_length=maximum_length, 225 | endpoint=endpoint)[1:]) 226 | return visited_nodes 227 | 228 | def get_actions_rewards( 229 | self, 230 | rollout: Sequence[node_lib.Node], 231 | cast_action: Callable[[Any], Any] = lambda x: x, 232 | ) -> tuple[Sequence[Optional[Any]], Sequence[float]]: 233 | """Returns the actions and rewards of a rollout in a nice format. 234 | 235 | Args: 236 | rollout: The rollout of nodes to be parsed. 237 | cast_action: A function to cast the actions to the right format. 238 | """ 239 | curate = lambda x: cast_action(x) if x is not None else x 240 | actions = [curate(node.last_action) for node in rollout] 241 | rewards = [node.last_timestep.reward for node in rollout] 242 | return actions, rewards 243 | -------------------------------------------------------------------------------- /src/agent_debugger/pycoworld/interventions.py: -------------------------------------------------------------------------------- 1 | # Copyright 2023 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | 16 | """Pycoworld interventions. 17 | 18 | These are methods of the form (node, **params) -> new_node. 19 | """ 20 | 21 | import copy 22 | import dataclasses 23 | from typing import Any 24 | 25 | import dm_env 26 | 27 | from agent_debugger.src.agent_debugger import default_interventions 28 | from agent_debugger.src.agent_debugger import node as node_lib 29 | from agent_debugger.src.agent_debugger.pycoworld import types 30 | 31 | 32 | def _check_positive_position(position: types.Position) -> None: 33 | """Checks if the position is positive, required by pycolab.""" 34 | if position.row < 0 or position.col < 0: # pytype: disable=attribute-error # enable-nested-classes 35 | raise ValueError('The new position must have positive coordinates. Got ' 36 | f'{position}.') 37 | 38 | 39 | class InterventionContext: 40 | """Context used for interventions, exposing the engine and postprocessing. 41 | 42 | Attributes: 43 | engine: The internal Pycolab game engine, useful to perform interventions. 44 | new_node: The node after intervention, which will be returned. 45 | """ 46 | 47 | def __init__(self, node: node_lib.Node): 48 | """Initializes.""" 49 | self._node = node 50 | self._env_state = copy.deepcopy(self._node.env_state) 51 | self.engine = self._env_state.current_game 52 | self.new_node = None 53 | 54 | def __enter__(self) -> type['InterventionContext']: 55 | """Returns the object itself when we enter the context.""" 56 | return self # pytype: disable=bad-return-type 57 | 58 | def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: 59 | """Applies some postprocessing and computes the new node to return.""" 60 | self._env_state.current_game._render() 61 | 62 | # Compute the new timestep that the agent will see. 63 | new_observation = copy.deepcopy(self._node.last_timestep.observation) 64 | new_observation['board'] = self._env_state.observation_distiller( 65 | self._env_state.current_game._board) 66 | new_timestep = dm_env.TimeStep( 67 | step_type=self._node.last_timestep.step_type, 68 | reward=self._node.last_timestep.reward, 69 | discount=self._node.last_timestep.discount, 70 | observation=new_observation) 71 | 72 | self.new_node = dataclasses.replace( 73 | self._node, env_state=self._env_state, last_timestep=new_timestep) 74 | 75 | 76 | class PycoworldInterventions(default_interventions.DefaultInterventions): 77 | """Object containing the pycolab interventions.""" 78 | 79 | # pylint: disable=protected-access 80 | def set_sprite_visibility( 81 | self, 82 | node: node_lib.Node, 83 | sprite_id: int, 84 | visibility: bool, 85 | ) -> node_lib.Node: 86 | """Intervenes on a pycolab environment state to make a sprite (in)visible. 87 | 88 | Args: 89 | node: The node to intervene on. 90 | sprite_id: The identifier of the sprite. A full list can be found in the 91 | pycoworld/default_constants.py file. 92 | visibility: Whether the sprite should be made visible or invisible. 93 | 94 | Returns: 95 | A new node, where the sprite is (in)visible. 96 | """ 97 | with InterventionContext(node) as context: 98 | sprite = context.engine._sprites_and_drapes[chr(sprite_id)] # pytype: disable=attribute-error 99 | sprite._visible = visibility 100 | return context.new_node # pytype: disable=attribute-error 101 | 102 | def move_sprite_to( 103 | self, 104 | node: node_lib.Node, 105 | sprite_id: int, 106 | dest_position: types.Position, 107 | ) -> node_lib.Node: 108 | """Intervenes on a pycolab environment state to move a sprite. 109 | 110 | Args: 111 | node: The node to intervene on. 112 | sprite_id: The identifier of the sprite. A full list can be found in the 113 | pycoworld/default_constants.py file. 114 | dest_position: The desired position of the sprite. 115 | 116 | Returns: 117 | A new node, where the sprite has been moved. 118 | """ 119 | with InterventionContext(node) as context: 120 | dest_position = types.to_pycolab_position(dest_position) 121 | _check_positive_position(dest_position) 122 | sprite = context.engine._sprites_and_drapes[chr(sprite_id)] # pytype: disable=attribute-error 123 | sprite._position = dest_position 124 | return context.new_node # pytype: disable=attribute-error 125 | 126 | def move_drape_element_to( 127 | self, 128 | node: node_lib.Node, 129 | drape_id: int, 130 | start_position: types.Position, 131 | dest_position: types.Position, 132 | ) -> node_lib.Node: 133 | """Intervenes on a pycolab environment state to move a drape. 134 | 135 | Args: 136 | node: The node to intervene on. 137 | drape_id: The identifier of the drape. A full list can be found in the 138 | pycoworld/default_constants.py file. 139 | start_position: The position of the element to replace, coordinates must 140 | be positive. 141 | dest_position: The destination position of the element, coordinates must 142 | be positive. 143 | 144 | Returns: 145 | A new node, which drape has been updated. 146 | """ 147 | with InterventionContext(node) as context: 148 | dest_position = types.to_pycolab_position(dest_position) 149 | start_position = types.to_pycolab_position(start_position) 150 | _check_positive_position(start_position) 151 | _check_positive_position(dest_position) 152 | drape = context.engine._sprites_and_drapes[chr(drape_id)] # pytype: disable=attribute-error 153 | drape._c_u_r_t_a_i_n[start_position] = False 154 | drape._c_u_r_t_a_i_n[tuple(dest_position)] = True 155 | return context.new_node # pytype: disable=attribute-error 156 | 157 | def remove_drape_element( 158 | self, 159 | node: node_lib.Node, 160 | drape_id: int, 161 | position: types.Position, 162 | ) -> node_lib.Node: 163 | """Intervenes to remove an element in a drape. 164 | 165 | Args: 166 | node: The node to intervene on. 167 | drape_id: The identifier of the drape. A full list can be found in the 168 | pycoworld/default_constants.py file. 169 | position: The position of the element to remove, coordinates must be 170 | positive. 171 | 172 | Returns: 173 | A new node, which drape has been updated. 174 | """ 175 | with InterventionContext(node) as context: 176 | position = types.to_pycolab_position(position) 177 | _check_positive_position(position) 178 | drape = context.engine._sprites_and_drapes[chr(drape_id)] # pytype: disable=attribute-error 179 | drape._c_u_r_t_a_i_n[position] = False 180 | return context.new_node # pytype: disable=attribute-error 181 | 182 | def add_drape_element( 183 | self, 184 | node: node_lib.Node, 185 | drape_id: int, 186 | position: types.Position, 187 | ) -> node_lib.Node: 188 | """Intervenes to add an element in a drape. 189 | 190 | Args: 191 | node: The node to intervene on. 192 | drape_id: The identifier of the drape. A full list can be found in the 193 | pycoworld/default_constants.py file. 194 | position: The position of the element to add, coordinates must be 195 | positive. 196 | 197 | Returns: 198 | A new node, which drape has been updated. 199 | """ 200 | with InterventionContext(node) as context: 201 | position = types.to_pycolab_position(position) 202 | _check_positive_position(position) 203 | drape = context.engine._sprites_and_drapes[chr(drape_id)] # pytype: disable=attribute-error 204 | drape._c_u_r_t_a_i_n[position] = True 205 | return context.new_node # pytype: disable=attribute-error 206 | 207 | def replace_backdrop_element( 208 | self, 209 | node: node_lib.Node, 210 | position: types.Position, 211 | new_element_id: int, 212 | ) -> node_lib.Node: 213 | """Intervenes to replace an element of the backdrop. 214 | 215 | Args: 216 | node: The node to intervene on. 217 | position: The position of the element to replace, coordinates must be 218 | positive. 219 | new_element_id: The new element id to put at this position. 220 | 221 | Returns: 222 | A new node, where the backdrop has been updated. 223 | """ 224 | with InterventionContext(node) as context: 225 | position = types.to_pycolab_position(position) 226 | _check_positive_position(position) 227 | context.engine._backdrop._p_a_l_e_t_t_e._legal_characters.add( # pytype: disable=attribute-error 228 | chr(new_element_id)) 229 | context.engine._backdrop.curtain[tuple(position)] = new_element_id # pytype: disable=attribute-error 230 | return context.new_node # pytype: disable=attribute-error 231 | 232 | def set_frame_count( 233 | self, 234 | node: node_lib.Node, 235 | frame_count: int, 236 | ) -> node_lib.Node: 237 | """Changes the internal frame number of a pycolab engine. 238 | 239 | This number influences the 'plot' object, which contains the frame number as 240 | an attribute. 241 | Args: 242 | node: The node to intervene on. 243 | frame_count: the frame countto set on the engine. It must be positive. 244 | 245 | Returns: 246 | A new node, where the backdrop has been updated. 247 | """ 248 | with InterventionContext(node) as context: 249 | if frame_count < 0: 250 | raise ValueError(f'The frame count must be positive. Got {frame_count}') 251 | context.engine._the_plot._frame = frame_count # pytype: disable=attribute-error 252 | return context.new_node # pytype: disable=attribute-error 253 | 254 | # pylint: enable=protected-access 255 | -------------------------------------------------------------------------------- /src/pycoworld/default_sprites_and_drapes.py: -------------------------------------------------------------------------------- 1 | # Copyright 2023 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | 16 | """This file contains the default Sprites and Drapes used in Pycoworld. 17 | 18 | Sprites and Drapes are concepts part of Pycolab, the game engine we use. Sprites 19 | are unique entities which can usually move and act in the environment. Drapes 20 | are sets of entities which are not interacting with the environment, mostly 21 | idle. 22 | """ 23 | 24 | import abc 25 | from typing import MutableMapping, Any 26 | 27 | import numpy as np 28 | from pycolab import engine 29 | from pycolab import things as plab_things 30 | 31 | from agent_debugger.src.pycoworld import default_constants 32 | 33 | Tile = default_constants.Tile 34 | 35 | 36 | class PlayerSprite(plab_things.Sprite): 37 | """A `Sprite` for our player. 38 | 39 | The player can move around freely, as long as it doesn't attempt to 40 | move into an `IMPASSABLE` tile. There's also additional effects for 41 | certain special tiles, such as a goal, lava, and water tile. 42 | """ 43 | 44 | def update( 45 | self, 46 | actions: int, 47 | board: np.ndarray, 48 | layers: MutableMapping[str, np.ndarray], 49 | backdrop: np.ndarray, 50 | things: MutableMapping[str, Any], 51 | the_plot: engine.plot.Plot 52 | ) -> None: 53 | ypos, xpos = self.position 54 | height, width = board.shape 55 | 56 | # Where does the agent want to move? 57 | if actions == 0 and ypos >= 1: # go upward? 58 | if board[ypos - 1, xpos] not in default_constants.IMPASSABLE: 59 | self._position = self.Position(ypos - 1, xpos) 60 | elif actions == 1 and ypos <= height - 2: # go downward? 61 | if board[ypos + 1, xpos] not in default_constants.IMPASSABLE: 62 | self._position = self.Position(ypos + 1, xpos) 63 | elif actions == 2 and xpos >= 1: # go leftward? 64 | if board[ypos, xpos - 1] not in default_constants.IMPASSABLE: 65 | self._position = self.Position(ypos, xpos - 1) 66 | elif actions == 3 and xpos <= width - 2: # go rightward? 67 | if board[ypos, xpos + 1] not in default_constants.IMPASSABLE: 68 | self._position = self.Position(ypos, xpos + 1) 69 | ypos, xpos = self.position 70 | 71 | floor_tile = backdrop.curtain[ypos, xpos] 72 | terminals = [ 73 | Tile.TERMINAL, Tile.TERMINAL_R, Tile.TERMINAL_G, Tile.TERMINAL_B 74 | ] 75 | if floor_tile in terminals: 76 | the_plot.terminate_episode() 77 | 78 | if floor_tile in [Tile.LAVA]: 79 | the_plot.add_reward(default_constants.REWARD_LAVA) 80 | the_plot.terminate_episode() 81 | 82 | if floor_tile in [Tile.WATER]: 83 | the_plot.add_reward(default_constants.REWARD_WATER) 84 | 85 | 86 | class RewardDrape(plab_things.Drape): 87 | """A drape for small rewards.""" 88 | 89 | def update( 90 | self, 91 | actions: int, 92 | board: np.ndarray, 93 | layers: MutableMapping[str, np.ndarray], 94 | backdrop: np.ndarray, 95 | things: MutableMapping[str, Any], 96 | the_plot: engine.plot.Plot 97 | ) -> None: 98 | ypos, xpos = things[chr(Tile.PLAYER)].position 99 | if self.curtain[ypos, xpos]: 100 | the_plot.add_reward(default_constants.REWARD_SMALL) 101 | self.curtain[ypos, xpos] = False 102 | 103 | 104 | class BigRewardDrape(plab_things.Drape): 105 | """A drape for big rewards.""" 106 | 107 | def update( 108 | self, 109 | actions: int, 110 | board: np.ndarray, 111 | layers: MutableMapping[str, np.ndarray], 112 | backdrop: np.ndarray, 113 | things: MutableMapping[str, Any], 114 | the_plot: engine.plot.Plot 115 | ) -> None: 116 | ypos, xpos = things[chr(Tile.PLAYER)].position 117 | if self.curtain[ypos, xpos]: 118 | the_plot.add_reward(default_constants.REWARD_BIG) 119 | self.curtain[ypos, xpos] = False 120 | 121 | 122 | class ObjectDrape(plab_things.Drape): 123 | """A drape with objects that the agent can carry in its inventory.""" 124 | 125 | def update( 126 | self, 127 | actions: int, 128 | board: np.ndarray, 129 | layers: MutableMapping[str, np.ndarray], 130 | backdrop: np.ndarray, 131 | things: MutableMapping[str, Any], 132 | the_plot: engine.plot.Plot 133 | ) -> None: 134 | ypos, xpos = things[chr(Tile.PLAYER)].position 135 | 136 | if self.character not in the_plot.keys(): 137 | # The inventory is empty by default. 138 | the_plot[self.character] = 0 139 | 140 | if self.curtain[ypos, xpos]: 141 | the_plot[self.character] = the_plot[self.character] + 1 142 | self.curtain[ypos, xpos] = False 143 | 144 | 145 | class WallDrape(plab_things.Drape): 146 | """A drape for walls, which does nothing.""" 147 | 148 | def update( 149 | self, 150 | actions: int, 151 | board: np.ndarray, 152 | layers: MutableMapping[str, Any], 153 | backdrop: np.ndarray, 154 | things: MutableMapping[str, Any], 155 | the_plot: engine.plot.Plot 156 | ) -> None: 157 | """Updates the environment with the actions.""" 158 | 159 | 160 | class SensorDrape(plab_things.Drape, abc.ABC): 161 | """A drape for plain sensors, triggering some function in the environment.""" 162 | 163 | @abc.abstractmethod 164 | def _trigger( 165 | self, 166 | board: np.ndarray, 167 | layers: MutableMapping[str, np.ndarray], 168 | backdrop: np.ndarray, 169 | things: MutableMapping[str, Any], 170 | the_plot: engine.plot.Plot 171 | ) -> None: 172 | """Triggers something in the environment as soon as the sensor is pushed.""" 173 | 174 | def update( 175 | self, 176 | actions: int, 177 | board: np.ndarray, 178 | layers: MutableMapping[str, np.ndarray], 179 | backdrop: np.ndarray, 180 | things: MutableMapping[str, Any], 181 | the_plot: engine.plot.Plot 182 | ) -> None: 183 | ypos, xpos = things[chr(default_constants.Tile.PLAYER)].position 184 | if self.curtain[ypos, xpos]: 185 | # As soon as the agent steps on the sensor, we trigger it. 186 | self._trigger(board, layers, backdrop, things, the_plot) 187 | 188 | 189 | class DoorDrape(plab_things.Drape): 190 | """A drape with doors tiles. 191 | 192 | Doors come in different colors. They act as impassable tiles, unless the 193 | player has a key of the associated color in its inventory: in this case, 194 | the key is consumed and the door disappears. 195 | """ 196 | 197 | def update( 198 | self, 199 | actions: int, 200 | board: np.ndarray, 201 | layers: MutableMapping[str, np.ndarray], 202 | backdrop: np.ndarray, 203 | things: MutableMapping[str, Any], 204 | the_plot: engine.plot.Plot 205 | ) -> None: 206 | ypos, xpos = things[chr(Tile.PLAYER)].position 207 | height, width = self.curtain.shape 208 | 209 | # What's the required key? 210 | key = chr(ord(self.character) - 4) 211 | 212 | # Where does the agent want to move? 213 | # If the agent wants to move into a door, then check whether 214 | # the corresponding key is available. If it is, then remove the door. 215 | if actions == 0 and ypos >= 1: # go upward? 216 | if self.curtain[ypos - 1, xpos] and the_plot[key] > 0: 217 | self.curtain[ypos - 1, xpos] = False 218 | the_plot[key] = the_plot[key] - 1 219 | elif actions == 1 and ypos <= height - 2: # go downward? 220 | if self.curtain[ypos + 1, xpos] and the_plot[key] > 0: 221 | self.curtain[ypos + 1, xpos] = False 222 | the_plot[key] = the_plot[key] - 1 223 | elif actions == 2 and xpos >= 1: # go leftward? 224 | if self.curtain[ypos, xpos - 1] and the_plot[key] > 0: 225 | self.curtain[ypos, xpos - 1] = False 226 | the_plot[key] = the_plot[key] - 1 227 | elif actions == 3 and xpos <= width - 2: # go rightward? 228 | if self.curtain[ypos, xpos + 1] and the_plot[key] > 0: 229 | self.curtain[ypos, xpos + 1] = False 230 | the_plot[key] = the_plot[key] - 1 231 | 232 | 233 | class BoxDrape(plab_things.Drape): 234 | """A drape for pushable blocks. 235 | 236 | These blocks can be pushed by the player if there is a free space behind. 237 | """ 238 | 239 | def update( 240 | self, 241 | actions: int, 242 | board: np.ndarray, 243 | layers: MutableMapping[str, np.ndarray], 244 | backdrop: np.ndarray, 245 | things: MutableMapping[str, Any], 246 | the_plot: engine.plot.Plot 247 | ) -> None: 248 | ypos, xpos = things[chr(Tile.PLAYER)].position 249 | height, width = self.curtain.shape 250 | 251 | # Where does the agent want to move? 252 | if actions == 0 and ypos >= 2: # go upward? 253 | check_impassable = board[ypos - 2, 254 | xpos] not in default_constants.IMPASSABLE 255 | if self.curtain[ypos - 1, xpos] and check_impassable: 256 | self.curtain[ypos - 1, xpos] = False 257 | self.curtain[ypos - 2, xpos] = True 258 | elif actions == 1 and ypos <= height - 3: # go downward? 259 | check_impassable = board[ypos + 2, 260 | xpos] not in default_constants.IMPASSABLE 261 | if self.curtain[ypos + 1, xpos] and check_impassable: 262 | self.curtain[ypos + 1, xpos] = False 263 | self.curtain[ypos + 2, xpos] = True 264 | elif actions == 2 and xpos >= 2: # go leftward? 265 | check_impassable = board[ypos, 266 | xpos - 2] not in default_constants.IMPASSABLE 267 | if self.curtain[ypos, xpos - 1] and check_impassable: 268 | self.curtain[ypos, xpos - 1] = False 269 | self.curtain[ypos, xpos - 2] = True 270 | elif actions == 3 and xpos <= width - 3: # go rightward? 271 | check_impassable = board[ypos, 272 | xpos + 2] not in default_constants.IMPASSABLE 273 | if self.curtain[ypos, xpos + 1] and check_impassable: 274 | self.curtain[ypos, xpos + 1] = False 275 | self.curtain[ypos, xpos + 2] = True 276 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | 2 | Apache License 3 | Version 2.0, January 2004 4 | http://www.apache.org/licenses/ 5 | 6 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 7 | 8 | 1. Definitions. 9 | 10 | "License" shall mean the terms and conditions for use, reproduction, 11 | and distribution as defined by Sections 1 through 9 of this document. 12 | 13 | "Licensor" shall mean the copyright owner or entity authorized by 14 | the copyright owner that is granting the License. 15 | 16 | "Legal Entity" shall mean the union of the acting entity and all 17 | other entities that control, are controlled by, or are under common 18 | control with that entity. For the purposes of this definition, 19 | "control" means (i) the power, direct or indirect, to cause the 20 | direction or management of such entity, whether by contract or 21 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 22 | outstanding shares, or (iii) beneficial ownership of such entity. 23 | 24 | "You" (or "Your") shall mean an individual or Legal Entity 25 | exercising permissions granted by this License. 26 | 27 | "Source" form shall mean the preferred form for making modifications, 28 | including but not limited to software source code, documentation 29 | source, and configuration files. 30 | 31 | "Object" form shall mean any form resulting from mechanical 32 | transformation or translation of a Source form, including but 33 | not limited to compiled object code, generated documentation, 34 | and conversions to other media types. 35 | 36 | "Work" shall mean the work of authorship, whether in Source or 37 | Object form, made available under the License, as indicated by a 38 | copyright notice that is included in or attached to the work 39 | (an example is provided in the Appendix below). 40 | 41 | "Derivative Works" shall mean any work, whether in Source or Object 42 | form, that is based on (or derived from) the Work and for which the 43 | editorial revisions, annotations, elaborations, or other modifications 44 | represent, as a whole, an original work of authorship. For the purposes 45 | of this License, Derivative Works shall not include works that remain 46 | separable from, or merely link (or bind by name) to the interfaces of, 47 | the Work and Derivative Works thereof. 48 | 49 | "Contribution" shall mean any work of authorship, including 50 | the original version of the Work and any modifications or additions 51 | to that Work or Derivative Works thereof, that is intentionally 52 | submitted to Licensor for inclusion in the Work by the copyright owner 53 | or by an individual or Legal Entity authorized to submit on behalf of 54 | the copyright owner. For the purposes of this definition, "submitted" 55 | means any form of electronic, verbal, or written communication sent 56 | to the Licensor or its representatives, including but not limited to 57 | communication on electronic mailing lists, source code control systems, 58 | and issue tracking systems that are managed by, or on behalf of, the 59 | Licensor for the purpose of discussing and improving the Work, but 60 | excluding communication that is conspicuously marked or otherwise 61 | designated in writing by the copyright owner as "Not a Contribution." 62 | 63 | "Contributor" shall mean Licensor and any individual or Legal Entity 64 | on behalf of whom a Contribution has been received by Licensor and 65 | subsequently incorporated within the Work. 66 | 67 | 2. Grant of Copyright License. Subject to the terms and conditions of 68 | this License, each Contributor hereby grants to You a perpetual, 69 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 70 | copyright license to reproduce, prepare Derivative Works of, 71 | publicly display, publicly perform, sublicense, and distribute the 72 | Work and such Derivative Works in Source or Object form. 73 | 74 | 3. Grant of Patent License. Subject to the terms and conditions of 75 | this License, each Contributor hereby grants to You a perpetual, 76 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 77 | (except as stated in this section) patent license to make, have made, 78 | use, offer to sell, sell, import, and otherwise transfer the Work, 79 | where such license applies only to those patent claims licensable 80 | by such Contributor that are necessarily infringed by their 81 | Contribution(s) alone or by combination of their Contribution(s) 82 | with the Work to which such Contribution(s) was submitted. If You 83 | institute patent litigation against any entity (including a 84 | cross-claim or counterclaim in a lawsuit) alleging that the Work 85 | or a Contribution incorporated within the Work constitutes direct 86 | or contributory patent infringement, then any patent licenses 87 | granted to You under this License for that Work shall terminate 88 | as of the date such litigation is filed. 89 | 90 | 4. Redistribution. You may reproduce and distribute copies of the 91 | Work or Derivative Works thereof in any medium, with or without 92 | modifications, and in Source or Object form, provided that You 93 | meet the following conditions: 94 | 95 | (a) You must give any other recipients of the Work or 96 | Derivative Works a copy of this License; and 97 | 98 | (b) You must cause any modified files to carry prominent notices 99 | stating that You changed the files; and 100 | 101 | (c) You must retain, in the Source form of any Derivative Works 102 | that You distribute, all copyright, patent, trademark, and 103 | attribution notices from the Source form of the Work, 104 | excluding those notices that do not pertain to any part of 105 | the Derivative Works; and 106 | 107 | (d) If the Work includes a "NOTICE" text file as part of its 108 | distribution, then any Derivative Works that You distribute must 109 | include a readable copy of the attribution notices contained 110 | within such NOTICE file, excluding those notices that do not 111 | pertain to any part of the Derivative Works, in at least one 112 | of the following places: within a NOTICE text file distributed 113 | as part of the Derivative Works; within the Source form or 114 | documentation, if provided along with the Derivative Works; or, 115 | within a display generated by the Derivative Works, if and 116 | wherever such third-party notices normally appear. The contents 117 | of the NOTICE file are for informational purposes only and 118 | do not modify the License. You may add Your own attribution 119 | notices within Derivative Works that You distribute, alongside 120 | or as an addendum to the NOTICE text from the Work, provided 121 | that such additional attribution notices cannot be construed 122 | as modifying the License. 123 | 124 | You may add Your own copyright statement to Your modifications and 125 | may provide additional or different license terms and conditions 126 | for use, reproduction, or distribution of Your modifications, or 127 | for any such Derivative Works as a whole, provided Your use, 128 | reproduction, and distribution of the Work otherwise complies with 129 | the conditions stated in this License. 130 | 131 | 5. Submission of Contributions. Unless You explicitly state otherwise, 132 | any Contribution intentionally submitted for inclusion in the Work 133 | by You to the Licensor shall be under the terms and conditions of 134 | this License, without any additional terms or conditions. 135 | Notwithstanding the above, nothing herein shall supersede or modify 136 | the terms of any separate license agreement you may have executed 137 | with Licensor regarding such Contributions. 138 | 139 | 6. Trademarks. This License does not grant permission to use the trade 140 | names, trademarks, service marks, or product names of the Licensor, 141 | except as required for reasonable and customary use in describing the 142 | origin of the Work and reproducing the content of the NOTICE file. 143 | 144 | 7. Disclaimer of Warranty. Unless required by applicable law or 145 | agreed to in writing, Licensor provides the Work (and each 146 | Contributor provides its Contributions) on an "AS IS" BASIS, 147 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 148 | implied, including, without limitation, any warranties or conditions 149 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 150 | PARTICULAR PURPOSE. You are solely responsible for determining the 151 | appropriateness of using or redistributing the Work and assume any 152 | risks associated with Your exercise of permissions under this License. 153 | 154 | 8. Limitation of Liability. In no event and under no legal theory, 155 | whether in tort (including negligence), contract, or otherwise, 156 | unless required by applicable law (such as deliberate and grossly 157 | negligent acts) or agreed to in writing, shall any Contributor be 158 | liable to You for damages, including any direct, indirect, special, 159 | incidental, or consequential damages of any character arising as a 160 | result of this License or out of the use or inability to use the 161 | Work (including but not limited to damages for loss of goodwill, 162 | work stoppage, computer failure or malfunction, or any and all 163 | other commercial damages or losses), even if such Contributor 164 | has been advised of the possibility of such damages. 165 | 166 | 9. Accepting Warranty or Additional Liability. While redistributing 167 | the Work or Derivative Works thereof, You may choose to offer, 168 | and charge a fee for, acceptance of support, warranty, indemnity, 169 | or other liability obligations and/or rights consistent with this 170 | License. However, in accepting such obligations, You may act only 171 | on Your own behalf and on Your sole responsibility, not on behalf 172 | of any other Contributor, and only if You agree to indemnify, 173 | defend, and hold each Contributor harmless for any liability 174 | incurred by, or claims asserted against, such Contributor by reason 175 | of your accepting any such warranty or additional liability. 176 | 177 | END OF TERMS AND CONDITIONS 178 | 179 | APPENDIX: How to apply the Apache License to your work. 180 | 181 | To apply the Apache License to your work, attach the following 182 | boilerplate notice, with the fields enclosed by brackets "[]" 183 | replaced with your own identifying information. (Don't include 184 | the brackets!) The text should be enclosed in the appropriate 185 | comment syntax for the file format. We also recommend that a 186 | file or class name and description of purpose be included on the 187 | same "printed page" as the copyright notice for easier 188 | identification within third-party archives. 189 | 190 | Copyright [yyyy] [name of copyright owner] 191 | 192 | Licensed under the Apache License, Version 2.0 (the "License"); 193 | you may not use this file except in compliance with the License. 194 | You may obtain a copy of the License at 195 | 196 | http://www.apache.org/licenses/LICENSE-2.0 197 | 198 | Unless required by applicable law or agreed to in writing, software 199 | distributed under the License is distributed on an "AS IS" BASIS, 200 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 201 | See the License for the specific language governing permissions and 202 | limitations under the License. 203 | -------------------------------------------------------------------------------- /src/pycoworld/environment.py: -------------------------------------------------------------------------------- 1 | # Copyright 2023 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | 16 | """Pycoworld environment. 17 | 18 | Pycolab only provides a game engine, i.e. a way to create a game, with the 19 | different objects and their dynamics. However, we'd like to use a more formal 20 | interface which is an environment, with a reset(), step(), action_spec() and 21 | observation_spec() methods. We need to create an object to wrap the engine 22 | and create a true environment, inheriting dm_env.Environment. 23 | 24 | This file also contains a build_environment method, which takes a pycoworld 25 | level name, and directly creates the associated environment, ready to be 26 | used in association with the agent. 27 | """ 28 | 29 | import copy 30 | import dataclasses 31 | import itertools 32 | from typing import Any, Callable, Optional, Union 33 | 34 | import dm_env 35 | import jax.numpy as jnp 36 | import jax.random as jrandom 37 | import numpy as np 38 | from pycolab import cropping as pycolab_cropping 39 | from pycolab import engine as pycolab_engine 40 | from pycolab import rendering 41 | import tree 42 | 43 | from agent_debugger.src.pycoworld import default_constants 44 | from agent_debugger.src.pycoworld import serializable_environment 45 | from agent_debugger.src.pycoworld.levels import base_level 46 | from agent_debugger.src.pycoworld.levels import level_names 47 | 48 | _TILES = list(range(len(default_constants.Tile))) 49 | _PLAYER_TILE = chr(8) 50 | 51 | _RGB_MAPPING = {chr(n): (4 * n, 4 * n, 4 * n) for n in _TILES} 52 | 53 | 54 | class ObservationDistiller: 55 | """This class modifies the observations before they are sent to the agent.""" 56 | 57 | def __init__( 58 | self, 59 | array_converter: rendering.ObservationToArray, 60 | cropping: Optional[pycolab_cropping.ObservationCropper] = None 61 | ) -> None: 62 | """Initializes a Distiller.""" 63 | self._array_converter = array_converter 64 | cropper = pycolab_cropping.ObservationCropper 65 | self._cropping = tree.map_structure( 66 | lambda c: cropper() if c is None else c, cropping) 67 | 68 | def set_engine(self, engine: pycolab_engine.Engine) -> None: 69 | """Informs the Distiller of the current game engine in use. 70 | 71 | Args: 72 | engine: current engine in use by the `Environment` adapter. 73 | """ 74 | tree.map_structure(lambda c: c.set_engine(engine), self._cropping) 75 | 76 | def __call__(self, observation: Any) -> Any: 77 | """Distills a pycolab observation into the format required by the agent. 78 | 79 | Args: 80 | observation: observation to distill. 81 | 82 | Returns: 83 | the observation distilled for supplying to the agent. 84 | """ 85 | return tree.map_structure( 86 | lambda c: self._array_converter(c.crop(observation)), self._cropping) 87 | 88 | 89 | @dataclasses.dataclass 90 | class EnvState: 91 | """The Pycoworld environment state.""" 92 | state: Any 93 | current_game: pycolab_engine.Engine 94 | game_over: bool 95 | observation_distiller: ObservationDistiller 96 | seed: int 97 | 98 | 99 | class PycoworldEnvironment(serializable_environment.SerializableEnvironment): 100 | """A wrapper around the for_python.Environment to get and set states. 101 | 102 | The serialization is just a copy of the wrapped environment. 103 | """ 104 | 105 | def __init__( 106 | self, 107 | seed: int, 108 | game_factory: Callable[[jnp.ndarray], pycolab_engine.Engine], 109 | possible_actions: set[int], 110 | default_reward: Optional[float], 111 | observation_distiller: ObservationDistiller, 112 | max_iterations: Optional[int] = None, 113 | max_iterations_discount: Optional[float] = None, 114 | ) -> None: 115 | """Initializes the pycoworld environment.""" 116 | self._game_factory = game_factory 117 | self._default_reward = default_reward 118 | self._observation_distiller = observation_distiller 119 | self._max_iterations = max_iterations 120 | self._max_iterations_discount = max_iterations_discount 121 | 122 | self._state = None 123 | self._current_game = None 124 | self._game_over = None 125 | self._seed = seed 126 | 127 | # Compute action and observation spec. 128 | self._action_spec = dm_env.specs.DiscreteArray( 129 | len(possible_actions), dtype='int32', name='discrete') 130 | self._observation_spec = self._compute_observation_spec() 131 | 132 | def reset(self) -> dm_env.TimeStep: 133 | """Starts a new episode.""" 134 | self._current_game = self._game_factory(jrandom.PRNGKey(self._seed)) 135 | self._state = dm_env.StepType.FIRST 136 | self._observation_distiller.set_engine(self._current_game) 137 | # Collect environment returns from starting the game and update state. 138 | observations, reward, discount = self._current_game.its_showtime() 139 | 140 | self._game_over = self._is_game_over() 141 | observations, _, _ = self._standardise_timestep_values( 142 | observations, reward, discount) 143 | return dm_env.TimeStep( 144 | step_type=self._state, 145 | reward=None, 146 | discount=None, 147 | observation=observations) 148 | 149 | def step(self, action: int) -> dm_env.TimeStep: 150 | """Applies action, steps the world forward, and returns observations.""" 151 | # Clear episode internals and start a new episode, if episode ended or if 152 | # the game was not already underway. 153 | if self._state == dm_env.StepType.LAST: 154 | self._drop_last_episode() 155 | if self._current_game is None: 156 | return self.reset() 157 | 158 | # Execute the action in pycolab. 159 | observations, reward, discount = self._current_game.play(action) 160 | 161 | self._game_over = self._is_game_over() 162 | observations, reward, discount = self._standardise_timestep_values( 163 | observations, reward, discount) 164 | 165 | # Check the current status of the game. 166 | if self._game_over: 167 | self._state = dm_env.StepType.LAST 168 | else: 169 | self._state = dm_env.StepType.MID 170 | 171 | return dm_env.TimeStep( 172 | step_type=self._state, 173 | reward=reward, 174 | discount=discount, 175 | observation=observations) 176 | 177 | def observation_spec(self) -> dm_env.specs.Array: 178 | """Returns the observation specifications of the environment.""" 179 | return self._observation_spec 180 | 181 | def action_spec(self) -> dm_env.specs.Array: 182 | """Returns the action specifications of the environment.""" 183 | return self._action_spec 184 | 185 | def get_state(self) -> EnvState: 186 | """Returns the state of the pycoworld environment.""" 187 | env_state = EnvState( 188 | state=self._state, 189 | current_game=self._current_game, 190 | game_over=self._game_over, 191 | observation_distiller=self._observation_distiller, 192 | seed=self._seed) 193 | return copy.deepcopy(env_state) 194 | 195 | def set_state(self, env_state: EnvState) -> None: 196 | """Sets the state of the pycoworld environment.""" 197 | env_state = copy.deepcopy(env_state) 198 | self._state = env_state.state 199 | self._current_game = env_state.current_game 200 | self._game_over = env_state.game_over 201 | self._observation_distiller = env_state.observation_distiller 202 | self._seed = env_state.seed 203 | 204 | def _compute_observation_spec(self) -> Any: 205 | """Returns after compute the observation spec of the environment. 206 | 207 | We need a special method to do this, as we are retrieving the specs by 208 | launching a game. This is an internal and private method only. 209 | """ 210 | 211 | def obs_names_that_count_up(): 212 | for i in itertools.count(): 213 | yield 'obs{}'.format(i) 214 | 215 | names = obs_names_that_count_up() 216 | 217 | timestep = self.reset() 218 | spec_array = dm_env.specs.Array 219 | observation_spec = tree.map_structure( 220 | lambda a: spec_array(shape=a.shape, dtype=a.dtype, name=next(names)), 221 | timestep.observation) 222 | self._drop_last_episode() 223 | return observation_spec 224 | 225 | def _standardise_timestep_values( 226 | self, 227 | observations: Any, 228 | reward: Optional[float], 229 | discount: Optional[float], 230 | ) -> tuple[Any, Optional[float], Optional[float]]: 231 | """Applies defaults and standard packaging to timestep values if needed.""" 232 | observations = copy.deepcopy(self._observation_distiller(observations)) 233 | if isinstance(observations, np.ndarray): 234 | observations = {'board': observations} 235 | reward = reward if reward is not None else self._default_reward 236 | if self._max_iterations is not None and ( 237 | self._current_game.the_plot.frame >= 238 | self._max_iterations) and (not self._current_game.game_over): 239 | if self._max_iterations_discount is not None: 240 | discount = self._max_iterations_discount 241 | return observations, reward, discount 242 | 243 | def _is_game_over(self) -> bool: 244 | """Returns whether the game is over.""" 245 | # If we've reached the maximum number of game iterations, terminate the 246 | # current game. 247 | if self._max_iterations is not None and (self._current_game.the_plot.frame 248 | >= self._max_iterations): 249 | return True 250 | return self._current_game.game_over 251 | 252 | def _drop_last_episode(self) -> None: 253 | """Clears all the internal information about the game.""" 254 | self._state = None 255 | self._current_game = None 256 | self._game_over = None 257 | 258 | 259 | def build_environment( 260 | level: Union[base_level.PycoworldLevel, str], 261 | seed: int = 1, 262 | episode_length: Optional[int] = None, 263 | observation_encoding: str = 'feature_map', 264 | egocentric_horizon: Optional[int] = None, 265 | ) -> PycoworldEnvironment: 266 | """Builds and returns the environment for pycoworld. 267 | 268 | The observation is either a feature map (one hot encoding), a raw (the tile id 269 | directly) or rgb map of the board of the game. The actions are discrete and 270 | you can get the bounds via the env specs. 271 | Args: 272 | level: The level to build. 273 | seed: The seed used by the level 274 | episode_length: Number of steps in the episode 275 | observation_encoding: Must be in {feature_map, board, rgb} 276 | egocentric_horizon: The sight distance of the agent. None means the agent 277 | sees the complete board. 278 | 279 | Returns: 280 | The environment for the game 281 | """ 282 | if isinstance(level, str): 283 | level = level_names.level_by_name(level) 284 | 285 | # The distiller is used to processed the raw observation provided by pycolab 286 | # to an observation fed to the agent, here a numpy binary array. 287 | if observation_encoding == 'feature_map': 288 | array_converter = rendering.ObservationToFeatureArray( 289 | layers=list(map(chr, _TILES)), permute=(1, 2, 0)) 290 | elif observation_encoding == 'board': 291 | no_mapping = {chr(x): x for x in _TILES} 292 | array_converter = rendering.ObservationToArray(no_mapping, dtype=np.uint8) 293 | elif observation_encoding == 'rgb': 294 | array_converter = rendering.ObservationToArray( 295 | _RGB_MAPPING, permute=(1, 2, 0), dtype=np.uint8) 296 | else: 297 | raise ValueError('Observation encoding must be in {feature_map, rgb},' 298 | f'got {observation_encoding}.') 299 | 300 | if egocentric_horizon is not None: 301 | square_side_length = 2 * egocentric_horizon + 1 302 | cropper = pycolab_cropping.ScrollingCropper( 303 | rows=square_side_length, 304 | cols=square_side_length, 305 | to_track=[_PLAYER_TILE], 306 | pad_char=chr(0), 307 | scroll_margins=(1, 1)) 308 | else: 309 | cropper = None 310 | 311 | observation_distiller = ObservationDistiller( 312 | array_converter=array_converter, cropping=cropper) 313 | 314 | return PycoworldEnvironment( 315 | seed=seed, 316 | game_factory=level.sample_game, 317 | possible_actions=set(level.actions), 318 | observation_distiller=observation_distiller, 319 | default_reward=0., 320 | max_iterations=episode_length, 321 | max_iterations_discount=1.) 322 | -------------------------------------------------------------------------------- /colabs/experiments.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "markdown", 5 | "metadata": { 6 | "id": "2xxP2w6jcuk7" 7 | }, 8 | "source": [ 9 | "Copyright 2022 DeepMind Technologies Limited\n", 10 | "\n", 11 | "All software is licensed under the Apache License, Version 2.0 (Apache 2.0);\n", 12 | "you may not use this file except in compliance with the Apache 2.0 license.\n", 13 | "You may obtain a copy of the Apache 2.0 license at:\n", 14 | "https://www.apache.org/licenses/LICENSE-2.0\n", 15 | "\n", 16 | "All other materials are licensed under the Creative Commons Attribution 4.0\n", 17 | "International License (CC-BY). You may obtain a copy of the CC-BY license at:\n", 18 | "https://creativecommons.org/licenses/by/4.0/legalcode\n", 19 | "\n", 20 | "Unless required by applicable law or agreed to in writing, all software and\n", 21 | "materials distributed here under the Apache 2.0 or CC-BY licenses are\n", 22 | "distributed on an \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,\n", 23 | "either express or implied. See the licenses for the specific language governing\n", 24 | "permissions and limitations under those licenses.\n", 25 | "\n", 26 | "This is not an official Google product." 27 | ] 28 | }, 29 | { 30 | "cell_type": "markdown", 31 | "metadata": { 32 | "id": "d88AJSu2dBDr" 33 | }, 34 | "source": [ 35 | "# **Causal analysis of agents using the Agent Debugger**" 36 | ] 37 | }, 38 | { 39 | "cell_type": "markdown", 40 | "metadata": { 41 | "id": "d9rnDfPOyv9V" 42 | }, 43 | "source": [ 44 | "This colab contains the experiments presented in the paper [Causal Analysis of Agent Behavior for AI Safety](https://arxiv.org/pdf/2103.03938.pdf).\n", 45 | "\n", 46 | "It uses agents which have been previously trained using [Impala](https://arxiv.org/pdf/1802.01561.pdf), and we download their parameters (i.e. neural networks weights) from local files, also open sourced. The environments are based on [Pycolab](https://github.com/deepmind/pycolab), an open source library to build 2D gridworlds." 47 | ] 48 | }, 49 | { 50 | "cell_type": "markdown", 51 | "metadata": { 52 | "id": "4G_B4ySezog7" 53 | }, 54 | "source": [ 55 | "The main tool used is the \"Agent Debugger\", which allows us to perform interventions easily and in a standardized way on the agent and the environment." 56 | ] 57 | }, 58 | { 59 | "cell_type": "code", 60 | "execution_count": null, 61 | "metadata": { 62 | "id": "D07vGGnMQ4tc" 63 | }, 64 | "outputs": [], 65 | "source": [ 66 | "!git clone https://github.com/deepmind/agent_debugger.git\n", 67 | "!pip install -r agent_debugger/requirements.txt" 68 | ] 69 | }, 70 | { 71 | "cell_type": "code", 72 | "execution_count": null, 73 | "metadata": { 74 | "cellView": "form", 75 | "id": "VK2RNVV2lCmo" 76 | }, 77 | "outputs": [], 78 | "source": [ 79 | "# python3\n", 80 | "# @title Imports\n", 81 | "\n", 82 | "import collections\n", 83 | "import dill\n", 84 | "import functools\n", 85 | "import itertools\n", 86 | "import random\n", 87 | "import requests\n", 88 | "from typing import Any\n", 89 | "\n", 90 | "import dm_env\n", 91 | "import haiku as hk\n", 92 | "import numpy as np\n", 93 | "import tensorflow as tf\n", 94 | "\n", 95 | "from agent_debugger.src import impala_net\n", 96 | "from agent_debugger.src import impala_agent\n", 97 | "\n", 98 | "from agent_debugger.src.agent_debugger import node as node_lib\n", 99 | "from agent_debugger.src.agent_debugger.pycoworld import debugger as pcw_dbg\n", 100 | "from agent_debugger.src.agent_debugger.pycoworld import interventions as pcw_interv\n", 101 | "from agent_debugger.src.pycoworld import default_constants as pcw_constants\n", 102 | "from agent_debugger.src.pycoworld import environment as pcw_env" 103 | ] 104 | }, 105 | { 106 | "cell_type": "code", 107 | "execution_count": null, 108 | "metadata": { 109 | "cellView": "form", 110 | "id": "pJhjvVVRszr5" 111 | }, 112 | "outputs": [], 113 | "source": [ 114 | "# @title Constants used throughout.\n", 115 | "\n", 116 | "n_rollouts = 100\n", 117 | "\n", 118 | "# Useful shortcut.\n", 119 | "Tile = pcw_constants.Tile" 120 | ] 121 | }, 122 | { 123 | "cell_type": "code", 124 | "execution_count": null, 125 | "metadata": { 126 | "cellView": "form", 127 | "id": "WTQznPzM6pef" 128 | }, 129 | "outputs": [], 130 | "source": [ 131 | "# @title Utils used throughout.\n", 132 | "\n", 133 | "def prepare_init_nodes(debugger: pcw_dbg.PycoworldDebugger) -\u003e list[node_lib.Node]:\n", 134 | " \"\"\"Returns a list of nodes, which agent and env seeds have been changed.\"\"\"\n", 135 | " root_node = debugger.get_root_node()\n", 136 | " init_nodes = []\n", 137 | " for seed in range(n_rollouts):\n", 138 | " node = debugger.interventions.change_env_seed(root_node, seed)\n", 139 | " node = debugger.interventions.change_agent_seed(node, 2 * seed)\n", 140 | " init_nodes.append(node)\n", 141 | " return init_nodes\n", 142 | "\n", 143 | "def load_params(agent_name: str) -\u003e Any:\n", 144 | " \"\"\"Returns a set of parameters downloaded from google storage.\"\"\"\n", 145 | " url_base = \"https://storage.googleapis.com/dm_agent_debugger/{}.dill\"\n", 146 | " url = url_base.format(agent_name)\n", 147 | " file = requests.get(url, allow_redirects=True, stream=True)\n", 148 | " return dill.load(file.raw)\n", 149 | "\n", 150 | "def make_agent_from_params(\n", 151 | " params: hk.Params, with_lstm: bool\n", 152 | ") -\u003e impala_agent.ImpalaAgent:\n", 153 | " \"\"\"Returns an Impala agent, adapted for the Agent Debugger.\n", 154 | "\n", 155 | " Args:\n", 156 | " params: The parameters of the agent (neural network weights).\n", 157 | " with_lstm: Whether the agent contains an LSTM cell or not.\n", 158 | " \"\"\"\n", 159 | " net_factory = functools.partial(\n", 160 | " impala_net.RecurrentConvNet,\n", 161 | " conv_widths=(128, 128, 128),\n", 162 | " conv_kernels=(3, 3, 3),\n", 163 | " padding='SAME',\n", 164 | " torso_widths=(128,),\n", 165 | " lstm_width=128 if with_lstm else 0,\n", 166 | " head_widths=(128,),\n", 167 | " num_actions=4\n", 168 | " )\n", 169 | "\n", 170 | " # Changing the keys of the parameters to match the module's name.\n", 171 | " new_params = {}\n", 172 | " for key in params:\n", 173 | " values = key.split('/')\n", 174 | " new_key = '/'.join(['recurrent_conv_net']+values[1:])\n", 175 | " new_params[new_key] = params[key]\n", 176 | "\n", 177 | " return impala_agent.ImpalaAgent(net_factory=net_factory, params=new_params)" 178 | ] 179 | }, 180 | { 181 | "cell_type": "markdown", 182 | "metadata": { 183 | "id": "2RQgTCHKkwbM" 184 | }, 185 | "source": [ 186 | "## 1 - Confounders" 187 | ] 188 | }, 189 | { 190 | "cell_type": "code", 191 | "execution_count": null, 192 | "metadata": { 193 | "cellView": "form", 194 | "id": "lCEGWaVq5ZUz" 195 | }, 196 | "outputs": [], 197 | "source": [ 198 | "# @title Create interventions.\n", 199 | "# @markdown An intervention acts on a node (see node.py in the code) and changes some internal states (agent or environment).\n", 200 | "# @markdown In this case we change the floor type and pill position, which are attributes of the environment.\n", 201 | "\n", 202 | "def change_floor_type(node: node_lib.Node, floor_type: str) -\u003e node_lib.Node:\n", 203 | " new_floor_type = Tile.SAND if floor_type==\"sand\" else Tile.GRASS\n", 204 | " # These positions are hardcoded for the environment we use, ie \"grass_sand\".\n", 205 | " for position in [(3, 2), (3, 3), (3, 4), (2, 4), (4, 4)]:\n", 206 | " node = debugger.interventions.replace_backdrop_element(\n", 207 | " node, position=position, new_element_id=new_floor_type)\n", 208 | " return node\n", 209 | "\n", 210 | "def change_pill_position(node: node_lib.Node, pill_pos: str) -\u003e node_lib.Node:\n", 211 | " new_position = (1, 4) if pill_pos==\"right\" else (5, 4)\n", 212 | " position = debugger.extractors.get_element_positions(node, element_id=Tile.REWARD)[0]\n", 213 | " # Check that the position is actually new, otherwise there is nothing to do.\n", 214 | " if new_position != position:\n", 215 | " node = debugger.interventions.move_drape_element_to(\n", 216 | " node, drape_id=Tile.REWARD, start_position=position, dest_position=new_position)\n", 217 | " # Don't forget to change the terminal states too, which lie underneath the tiles.\n", 218 | " node = debugger.interventions.replace_backdrop_element(node, new_position, Tile.TERMINAL)\n", 219 | " floor_type = debugger.extractors.get_backdrop_curtain(node)[3, 2]\n", 220 | " node = debugger.interventions.replace_backdrop_element(node, position, floor_type)\n", 221 | " return node\n" 222 | ] 223 | }, 224 | { 225 | "cell_type": "code", 226 | "execution_count": null, 227 | "metadata": { 228 | "cellView": "form", 229 | "id": "qz8JvK1A5aW9" 230 | }, 231 | "outputs": [], 232 | "source": [ 233 | "# @title Create variable extractors.\n", 234 | "# @markdown This tells us what is the reward's postion or the agent's position for a given rollout. We use this information to compute the probabilities.\n", 235 | "\n", 236 | "def get_agent_final_position(rollout: list[node_lib.Node]) -\u003e str:\n", 237 | " \"\"\"Returns the agent final position, in {\"right\", \"left\"}.\"\"\"\n", 238 | " final_position = debugger.extractors.get_element_positions(\n", 239 | " rollout[-1], Tile.PLAYER)[0]\n", 240 | " return \"right\" if final_position[0] \u003c= 2 else \"left\"\n", 241 | "\n", 242 | "def get_reward_position(rollout: list[node_lib.Node]) -\u003e str:\n", 243 | " \"\"\"Returns the reward position, in {\"right\", \"left\"}.\"\"\"\n", 244 | " reward_position = debugger.extractors.get_element_positions(\n", 245 | " rollout[0], Tile.REWARD)[0]\n", 246 | " return \"right\" if reward_position[0] \u003c= 2 else \"left\"" 247 | ] 248 | }, 249 | { 250 | "cell_type": "code", 251 | "execution_count": null, 252 | "metadata": { 253 | "cellView": "form", 254 | "id": "Vi6Qu3pbaeZ3" 255 | }, 256 | "outputs": [], 257 | "source": [ 258 | "# @title Create the debugger for our agent and environment.\n", 259 | "# @markdown A different agent implies a different debugger, because it acts on the coupled agent+environment dynamical system.\n", 260 | "\n", 261 | "agent = \"B\" #@param[\"A\", \"B\"]\n", 262 | "\n", 263 | "# Agent A - trained on correlated reward+floor\n", 264 | "if agent == \"A\":\n", 265 | " params = load_params('confounders_a')\n", 266 | " env = pcw_env.build_environment(level='grass_sand')\n", 267 | "\n", 268 | "# Agent B - trained on uncorrelated reward+floor\n", 269 | "if agent == \"B\":\n", 270 | " params = load_params('confounders_b')\n", 271 | " env = pcw_env.build_environment(level='grass_sand_uncorrelated')\n", 272 | "\n", 273 | "trained_agent = make_agent_from_params(params, with_lstm=True)\n", 274 | "\n", 275 | "debugger = pcw_dbg.PycoworldDebugger(trained_agent, env)\n", 276 | "root_node = debugger.get_root_node()" 277 | ] 278 | }, 279 | { 280 | "cell_type": "code", 281 | "execution_count": null, 282 | "metadata": { 283 | "cellView": "form", 284 | "id": "-N4jtCqPti8f" 285 | }, 286 | "outputs": [], 287 | "source": [ 288 | "# @title Prepare nodes with different agent and environment seeds.\n", 289 | "# @markdown This cell creates nodes with different environment and agent seeds. We'll use these nodes later, but intervene on them before creating the rollouts.\n", 290 | "\n", 291 | "init_nodes = prepare_init_nodes(debugger)" 292 | ] 293 | }, 294 | { 295 | "cell_type": "code", 296 | "execution_count": null, 297 | "metadata": { 298 | "cellView": "form", 299 | "id": "0oPoYj27updg" 300 | }, 301 | "outputs": [], 302 | "source": [ 303 | "# @title Compute conditional probabilities.\n", 304 | "# @markdown \u003cstrong\u003eNote that, in all the colab, we compute the conditional probabilities using P(A|B) = P(A and B) / P(B).\u003c/strong\u003e\n", 305 | "\n", 306 | "reward_position_count = collections.defaultdict(lambda: 0)\n", 307 | "reward_match_agent_count = collections.defaultdict(lambda: 0)\n", 308 | "for node in init_nodes:\n", 309 | " rollout = debugger.get_rollout(node, maximum_length=15)\n", 310 | " reward_position = get_reward_position(rollout)\n", 311 | " agent_final_position = get_agent_final_position(rollout)\n", 312 | " reward_position_count[reward_position] += 1\n", 313 | " if reward_position == agent_final_position:\n", 314 | " reward_match_agent_count[reward_position] += 1\n", 315 | "print(\"P(T=l | R=l) = \" + str(reward_match_agent_count[\"left\"] / (reward_position_count[\"left\"])))\n", 316 | "print(\"P(T=r | R=r) = \" + str(reward_match_agent_count[\"right\"] / (reward_position_count[\"right\"])))" 317 | ] 318 | }, 319 | { 320 | "cell_type": "code", 321 | "execution_count": null, 322 | "metadata": { 323 | "cellView": "form", 324 | "id": "ASVK7QeN0wfr" 325 | }, 326 | "outputs": [], 327 | "source": [ 328 | "# @title Compute interventional probabilities - reward position interventions.\n", 329 | "\n", 330 | "agent_right_count = 0\n", 331 | "for node in init_nodes:\n", 332 | " do_reward_right = change_pill_position(node, pill_pos='right')\n", 333 | " rollout = debugger.get_rollout(do_reward_right, maximum_length=15)\n", 334 | " agent_right_count += int(get_agent_final_position(rollout) == 'right')\n", 335 | "print(\"P(T=r | do(R=r)) = \" + str(agent_right_count / (len(init_nodes))))\n", 336 | "\n", 337 | "agent_left_count = 0\n", 338 | "for node in init_nodes:\n", 339 | " do_reward_left = change_pill_position(node, pill_pos='left')\n", 340 | " rollout = debugger.get_rollout(do_reward_left, maximum_length=15)\n", 341 | " agent_left_count += int(get_agent_final_position(rollout) == 'left')\n", 342 | "print(\"P(T=l | do(R=l)) = \" + str(agent_left_count / (len(init_nodes))))" 343 | ] 344 | }, 345 | { 346 | "cell_type": "code", 347 | "execution_count": null, 348 | "metadata": { 349 | "cellView": "form", 350 | "id": "g3uA3Na71cUV" 351 | }, 352 | "outputs": [], 353 | "source": [ 354 | "# @title Compute interventional probabilities - floor type interventions.\n", 355 | "\n", 356 | "agent_right_count = 0\n", 357 | "for node in init_nodes:\n", 358 | " do_floor_sand = change_floor_type(node, \"sand\")\n", 359 | " rollout = debugger.get_rollout(do_floor_sand, maximum_length=15)\n", 360 | " agent_right_count += int(get_agent_final_position(rollout) == 'right')\n", 361 | "print(\"P(T=r | do(F=s)) = \" + str(agent_right_count / (len(init_nodes))))\n", 362 | "\n", 363 | "agent_left_count = 0\n", 364 | "for node in init_nodes:\n", 365 | " do_floor_grass = change_floor_type(node, \"grass\")\n", 366 | " rollout = debugger.get_rollout(do_floor_grass, maximum_length=15)\n", 367 | " agent_left_count += int(get_agent_final_position(rollout) == 'left')\n", 368 | "print(\"P(T=l | do(F=g)) = \" + str(agent_left_count / (len(init_nodes))))" 369 | ] 370 | }, 371 | { 372 | "cell_type": "markdown", 373 | "metadata": { 374 | "id": "GQaEwvB83OhV" 375 | }, 376 | "source": [ 377 | "# 2 - Memory" 378 | ] 379 | }, 380 | { 381 | "cell_type": "code", 382 | "execution_count": null, 383 | "metadata": { 384 | "cellView": "form", 385 | "id": "XRY0rWWTsHmb" 386 | }, 387 | "outputs": [], 388 | "source": [ 389 | "# @title Create variable extractors.\n", 390 | "\n", 391 | "def floor_type(rollout):\n", 392 | " is_sand = bool(debugger.extractors.get_element_positions(rollout[0], Tile.SAND))\n", 393 | " return 'sand' if is_sand else 'grass'\n", 394 | "\n", 395 | "def agent_pos_after_interv(rollout):\n", 396 | " pos = debugger.extractors.get_element_positions(rollout[5], Tile.PLAYER)[0]\n", 397 | " return 'right' if pos.row \u003c= 3 else 'left'\n", 398 | "\n", 399 | "def agent_pos_end(rollout):\n", 400 | " pos = debugger.extractors.get_element_positions(rollout[-1], Tile.PLAYER)[0]\n", 401 | " return 'right' if pos.row \u003c= 3 else 'left'\n" 402 | ] 403 | }, 404 | { 405 | "cell_type": "code", 406 | "execution_count": null, 407 | "metadata": { 408 | "cellView": "form", 409 | "id": "5r-EIby03Rq3" 410 | }, 411 | "outputs": [], 412 | "source": [ 413 | "# @title Create the debugger.\n", 414 | "\n", 415 | "agent = \"B\" #@param[\"A\", \"B\"]\n", 416 | "\n", 417 | "env = pcw_env.build_environment(\n", 418 | " level='large_color_memory',\n", 419 | " egocentric_horizon=1\n", 420 | ")\n", 421 | "\n", 422 | "# Agent A - with memory\n", 423 | "if agent == \"A\":\n", 424 | " params = load_params('memory_a')\n", 425 | " with_lstm = True\n", 426 | "\n", 427 | "# Agent B - without memory\n", 428 | "if agent == \"B\":\n", 429 | " params = load_params('memory_b')\n", 430 | " with_lstm = False\n", 431 | "\n", 432 | "trained_agent = make_agent_from_params(params, with_lstm=with_lstm)\n", 433 | "\n", 434 | "debugger = pcw_dbg.PycoworldDebugger(trained_agent, env)\n", 435 | "init_nodes = prepare_init_nodes(debugger)" 436 | ] 437 | }, 438 | { 439 | "cell_type": "code", 440 | "execution_count": null, 441 | "metadata": { 442 | "cellView": "form", 443 | "id": "N7OTQP0Jr3n1" 444 | }, 445 | "outputs": [], 446 | "source": [ 447 | "# @title Compute conditional probabilities.\n", 448 | "\n", 449 | "floor_type_l = []\n", 450 | "agent_pos_after_interv_l = []\n", 451 | "agent_pos_end_l = []\n", 452 | "for node in init_nodes:\n", 453 | " rollout = debugger.get_rollout(node, maximum_length=15)\n", 454 | " floor_type_l.append(floor_type(rollout))\n", 455 | " agent_pos_after_interv_l.append(agent_pos_after_interv(rollout))\n", 456 | " agent_pos_end_l.append(agent_pos_end(rollout))\n", 457 | "\n", 458 | "print(\"P(T=l | F=g) = \"+str(np.mean(np.logical_and(np.array(agent_pos_end_l)==\"left\", np.array(floor_type_l)==\"grass\"))/np.mean(np.array(floor_type_l)==\"grass\")))\n", 459 | "print(\"P(T=r | F=s) = \"+str(np.mean(np.logical_and(np.array(agent_pos_end_l)==\"right\", np.array(floor_type_l)==\"sand\"))/np.mean(np.array(floor_type_l)==\"sand\")))\n", 460 | "print(\"P(P=l | F=g) = \"+str(np.mean(np.logical_and(np.array(agent_pos_after_interv_l)==\"left\", np.array(floor_type_l)==\"grass\"))/np.mean(np.array(floor_type_l)==\"grass\")))\n", 461 | "print(\"P(P=r | F=s) = \"+str(np.mean(np.logical_and(np.array(agent_pos_after_interv_l)==\"right\", np.array(floor_type_l)==\"sand\"))/np.mean(np.array(floor_type_l)==\"sand\")))" 462 | ] 463 | }, 464 | { 465 | "cell_type": "code", 466 | "execution_count": null, 467 | "metadata": { 468 | "cellView": "form", 469 | "id": "gNcBgnHkuVXe" 470 | }, 471 | "outputs": [], 472 | "source": [ 473 | "# @title Compute interventional probabilities.\n", 474 | "\n", 475 | "UP = 0\n", 476 | "DOWN = 1\n", 477 | "LEFT = 2\n", 478 | "RIGHT = 3\n", 479 | "\n", 480 | "floor_type_l = []\n", 481 | "agent_pos_end_l = []\n", 482 | "for node in init_nodes:\n", 483 | " breakpoint = lambda node: node.episode_step == 3\n", 484 | " intervention_at_breakpoint = functools.partial(\n", 485 | " debugger.interventions.change_agent_next_actions,\n", 486 | " forced_next_actions=([UP, UP]))\n", 487 | " rollout = debugger.get_intervened_rollout(\n", 488 | " node, 15,\n", 489 | " breakpoint, intervention_at_breakpoint)\n", 490 | " floor_type_l.append(floor_type(rollout))\n", 491 | " agent_pos_end_l.append(agent_pos_end(rollout))\n", 492 | "\n", 493 | "print(\"P(T=l | do(P=r), F=g) = \"+str(np.mean(np.logical_and(np.array(agent_pos_end_l)==\"left\", np.array(floor_type_l)==\"grass\"))/np.mean(np.array(floor_type_l)==\"grass\")))\n", 494 | "\n", 495 | "\n", 496 | "floor_type_l = []\n", 497 | "agent_pos_end_l = []\n", 498 | "for node in init_nodes:\n", 499 | " breakpoint = lambda node: node.episode_step == 3\n", 500 | " intervention_at_breakpoint = functools.partial(\n", 501 | " debugger.interventions.change_agent_next_actions,\n", 502 | " forced_next_actions=([DOWN, DOWN]))\n", 503 | " rollout = debugger.get_intervened_rollout(\n", 504 | " node, 15,\n", 505 | " breakpoint, intervention_at_breakpoint)\n", 506 | " floor_type_l.append(floor_type(rollout))\n", 507 | " agent_pos_end_l.append(agent_pos_end(rollout))\n", 508 | "\n", 509 | "print(\"P(T=r | do(P=l), F=s) = \"+str(np.mean(np.logical_and(np.array(agent_pos_end_l)==\"right\", np.array(floor_type_l)==\"sand\"))/np.mean(np.array(floor_type_l)==\"sand\")))" 510 | ] 511 | }, 512 | { 513 | "cell_type": "markdown", 514 | "metadata": { 515 | "id": "TAOYm9E_5MMN" 516 | }, 517 | "source": [ 518 | "# 3 - Robust generalization" 519 | ] 520 | }, 521 | { 522 | "cell_type": "code", 523 | "execution_count": null, 524 | "metadata": { 525 | "cellView": "form", 526 | "id": "recDkNO25lao" 527 | }, 528 | "outputs": [], 529 | "source": [ 530 | "# @title Create interventions and variable extractors.\n", 531 | "\n", 532 | "def move_reward_to_quadrant(node, quadrant: str):\n", 533 | " all_positions = itertools.product(range(1, 7), range(1, 7))\n", 534 | " if quadrant == 'south':\n", 535 | " filter_quadrant = lambda pos: pos[0] \u003e= 4 and pos[1] \u003e= 4\n", 536 | " elif quadrant == 'north':\n", 537 | " filter_quadrant = lambda pos: pos[0] \u003c 4 and pos[1] \u003c 4\n", 538 | " elif quadrant == 'east':\n", 539 | " filter_quadrant = lambda pos: pos[0] \u003c 4 and pos[1] \u003e= 4\n", 540 | " else:\n", 541 | " filter_quadrant = lambda pos: pos[0] \u003e= 4 and pos[1] \u003c 4\n", 542 | " positions = list(filter(filter_quadrant, all_positions))\n", 543 | " agent_pos = debugger.extractors.get_element_positions(node, Tile.PLAYER)[0]\n", 544 | " agent_pos = (agent_pos.row, agent_pos.col)\n", 545 | " if agent_pos in positions:\n", 546 | " positions.remove(agent_pos)\n", 547 | " new_reward_pos = random.choice(positions)\n", 548 | " reward_pos = debugger.extractors.get_element_positions(node, Tile.REWARD)[0]\n", 549 | " node = debugger.interventions.move_drape_element_to(\n", 550 | " node, Tile.REWARD, start_position=reward_pos, dest_position=new_reward_pos)\n", 551 | " node = debugger.interventions.replace_backdrop_element(node, reward_pos, Tile.FLOOR)\n", 552 | " node = debugger.interventions.replace_backdrop_element(node, new_reward_pos, Tile.TERMINAL_R)\n", 553 | " return node\n", 554 | "\n", 555 | "def reward_taken(rollout):\n", 556 | " return (rollout[-1].last_timestep.reward == 1)\n" 557 | ] 558 | }, 559 | { 560 | "cell_type": "code", 561 | "execution_count": null, 562 | "metadata": { 563 | "cellView": "form", 564 | "id": "Vk1WlWHd5UvK" 565 | }, 566 | "outputs": [], 567 | "source": [ 568 | "# @title Create the debugger\n", 569 | "\n", 570 | "agent = \"B\" #@param[\"A\", \"B\"]\n", 571 | "\n", 572 | "# Agent A - trained on the full environment\n", 573 | "if agent == \"A\":\n", 574 | " params = load_params('generalization_a')\n", 575 | " env = pcw_env.build_environment(level='apples_full')\n", 576 | "\n", 577 | "# Agent B - trained only on part of the environment distribution\n", 578 | "if agent == \"B\":\n", 579 | " params = load_params('generalization_b')\n", 580 | " env = pcw_env.build_environment(level='apples_corner')\n", 581 | "\n", 582 | "trained_agent = make_agent_from_params(params, with_lstm=True)\n", 583 | "\n", 584 | "debugger = pcw_dbg.PycoworldDebugger(trained_agent, env)\n", 585 | "init_nodes = prepare_init_nodes(debugger)" 586 | ] 587 | }, 588 | { 589 | "cell_type": "code", 590 | "execution_count": null, 591 | "metadata": { 592 | "cellView": "form", 593 | "id": "46ppkAQf8VAL" 594 | }, 595 | "outputs": [], 596 | "source": [ 597 | "# @title Compute conditional probs.\n", 598 | "\n", 599 | "reward_taken_count = 0\n", 600 | "for node in init_nodes:\n", 601 | " rollout = debugger.get_rollout(node, maximum_length=15)\n", 602 | " reward_taken_count += int(reward_taken(rollout))\n", 603 | "print(\"P(R=l) = \" + str(reward_taken_count / len(init_nodes)))" 604 | ] 605 | }, 606 | { 607 | "cell_type": "code", 608 | "execution_count": null, 609 | "metadata": { 610 | "cellView": "form", 611 | "id": "uYagr17BweZ5" 612 | }, 613 | "outputs": [], 614 | "source": [ 615 | "# @title Compute interventional probs.\n", 616 | "\n", 617 | "for quadrant in ['south', 'north', 'east', 'west']:\n", 618 | " reward_taken_count = 0\n", 619 | " for node in init_nodes:\n", 620 | " node = move_reward_to_quadrant(node, quadrant)\n", 621 | " rollout = debugger.get_rollout(node, maximum_length=15)\n", 622 | " reward_taken_count += int(reward_taken(rollout))\n", 623 | " print(\"P(R=1 | G=\"+quadrant+\") = \" + str(reward_taken_count / len(init_nodes)))" 624 | ] 625 | }, 626 | { 627 | "cell_type": "markdown", 628 | "metadata": { 629 | "id": "pMne5HKZxy8f" 630 | }, 631 | "source": [ 632 | "# 4 - Counterfactuals" 633 | ] 634 | }, 635 | { 636 | "cell_type": "code", 637 | "execution_count": null, 638 | "metadata": { 639 | "cellView": "form", 640 | "id": "cjOe_6lr4ReZ" 641 | }, 642 | "outputs": [], 643 | "source": [ 644 | "# @title Create interventions and variable extractors.\n", 645 | "\n", 646 | "def move_door_to(node, door_pos: str):\n", 647 | " new_position = (3, 1) if door_pos == \"left\" else (3, 5)\n", 648 | " position = debugger.extractors.get_element_positions(node, Tile.DOOR_B)[0]\n", 649 | " return debugger.interventions.move_drape_element_to(\n", 650 | " node, Tile.DOOR_B, start_position=position, dest_position=new_position)\n", 651 | " \n", 652 | "def reward_taken(rollout):\n", 653 | " agent_pos = debugger.extractors.get_element_positions(\n", 654 | " rollout[-1], Tile.PLAYER)[0]\n", 655 | " if agent_pos in [(2, 2), (2, 4)]:\n", 656 | " return \"green\"\n", 657 | " return \"red\"\n", 658 | "\n", 659 | "def door_position(node):\n", 660 | " door_pos = debugger.extractors.get_element_positions(rollout[0], Tile.DOOR_B)[0]\n", 661 | " door_pos = (door_pos.row, door_pos.col)\n", 662 | " return \"left\" if door_pos == (3, 1) else \"right\"" 663 | ] 664 | }, 665 | { 666 | "cell_type": "code", 667 | "execution_count": null, 668 | "metadata": { 669 | "cellView": "form", 670 | "id": "vFQx-ZZsx2kQ" 671 | }, 672 | "outputs": [], 673 | "source": [ 674 | "# @title Compute variable values for both agents.\n", 675 | "# @markdown We compute the probabilities slightly differently than before here since we consider probabilities involving the ID of the agent (variable 'A' in the paper), and therefore compute values which depend on the metrics for agents 're' and 'gr' simultaneously.\n", 676 | "\n", 677 | "reward_taken_l = {\"A\": [], \"B\": []}\n", 678 | "door_pos_l = {\"A\": [], \"B\": []}\n", 679 | "\n", 680 | "for (agent_id, agent_name) in [(\"A\", \"counterfactuals_a\"), (\"B\", \"counterfactuals_b\")]:\n", 681 | " env = pcw_env.build_environment(level='red_green_apples')\n", 682 | "\n", 683 | " params = load_params(agent_name)\n", 684 | " trained_agent = make_agent_from_params(params, with_lstm=True)\n", 685 | "\n", 686 | " debugger = pcw_dbg.PycoworldDebugger(trained_agent, env)\n", 687 | " init_nodes = prepare_init_nodes(debugger)\n", 688 | " \n", 689 | " # Conditional regime.\n", 690 | " for node in init_nodes:\n", 691 | " rollout = debugger.get_rollout(node, maximum_length=15)\n", 692 | " reward_taken_l[agent_id].append(reward_taken(rollout))\n", 693 | " door_pos_l[agent_id].append(door_position(rollout))\n", 694 | " \n", 695 | " # Interventional regime.\n", 696 | " reward_taken_count = collections.defaultdict(lambda: 0)\n", 697 | " for node in init_nodes:\n", 698 | " node = move_door_to(node, \"right\")\n", 699 | " rollout = debugger.get_rollout(node, maximum_length=15)\n", 700 | " reward_taken_count[reward_taken(rollout)] += 1\n", 701 | " if agent_id == \"A\":\n", 702 | " print(\"P(R_{D=r}=gr | D=l, R=gr) = \" + str(reward_taken_count[\"green\"] / len(init_nodes)))\n", 703 | " if agent_id == \"B\":\n", 704 | " print(\"P(R_{D=r}=re | D=l, R=re) = \" + str(reward_taken_count[\"red\"] / len(init_nodes)))\n" 705 | ] 706 | }, 707 | { 708 | "cell_type": "code", 709 | "execution_count": null, 710 | "metadata": { 711 | "cellView": "form", 712 | "id": "Pw6o9hSS5NKs" 713 | }, 714 | "outputs": [], 715 | "source": [ 716 | "# @title Compute counterfactual probs.\n", 717 | "\n", 718 | "reward_a_red = np.array(reward_taken_l[\"A\"])==\"red\"\n", 719 | "reward_b_red = np.array(reward_taken_l[\"B\"])==\"red\"\n", 720 | "door_b_left = np.array(door_pos_l[\"B\"])==\"left\"\n", 721 | "door_a_left = np.array(door_pos_l[\"A\"])==\"left\"\n", 722 | "\n", 723 | "print(\"P(R=re) = \"+str((np.sum(reward_a_red) + np.sum(reward_b_red))/(2*n_rollouts)))\n", 724 | "print(\"P(A=re | R=re) = \"+str(np.mean(reward_b_red)/np.mean(reward_a_red + reward_b_red)))\n", 725 | "\n", 726 | "p_intersect = np.mean(np.logical_and(reward_b_red, door_b_left))\n", 727 | "p_evidence = np.mean(np.logical_and(reward_a_red, door_a_left) + np.logical_and(reward_b_red, door_b_left))\n", 728 | "print(\"P(A=re | D=l, R=re) = \"+str(p_intersect / p_evidence))" 729 | ] 730 | }, 731 | { 732 | "cell_type": "markdown", 733 | "metadata": { 734 | "id": "EbvBgKlf70CH" 735 | }, 736 | "source": [ 737 | "# 5 - Causal induction" 738 | ] 739 | }, 740 | { 741 | "cell_type": "markdown", 742 | "metadata": { 743 | "id": "Jcl2r4_k9qOi" 744 | }, 745 | "source": [ 746 | "No need to use pycoworld here, as the setup is very simple." 747 | ] 748 | }, 749 | { 750 | "cell_type": "code", 751 | "execution_count": null, 752 | "metadata": { 753 | "cellView": "form", 754 | "id": "eggfoUCv9dh4" 755 | }, 756 | "outputs": [], 757 | "source": [ 758 | "# @title Creating leader and follower actions, no interventions.\n", 759 | "\n", 760 | "episode_length = 1\n", 761 | "n_episodes = 10000\n", 762 | "\n", 763 | "# Agent 1 is red, 2 is blue.\n", 764 | "# Action 0 is right, action 1 is left.\n", 765 | "agent_1_actions, agent_2_actions = [], []\n", 766 | "leader_bool = []\n", 767 | "for episode in range(n_episodes):\n", 768 | " leader = np.random.randint(0, 2)\n", 769 | " leader_bool.append(leader)\n", 770 | "\n", 771 | " leader_action = np.random.randint(0, 2)\n", 772 | " follower_action = leader_action\n", 773 | " if np.random.rand() \u003c 0.1:\n", 774 | " follower_action = np.random.randint(0, 2)\n", 775 | " if leader == 0:\n", 776 | " agent_1_actions.append(leader_action)\n", 777 | " agent_2_actions.append(follower_action)\n", 778 | " else:\n", 779 | " agent_1_actions.append(follower_action)\n", 780 | " agent_2_actions.append(leader_action)\n", 781 | "leader_bool = np.array(leader_bool)\n", 782 | "agent_1_actions = np.array(agent_1_actions)\n", 783 | "agent_2_actions = np.array(agent_2_actions)\n", 784 | "\n", 785 | "print(\"P(L=b) = \"+str(np.mean(leader_bool)))\n", 786 | "print(\"P(L=b | R=l, B=l) = \"+str(np.mean(leader_bool * agent_1_actions * agent_2_actions) / np.mean(agent_1_actions * agent_2_actions)))\n", 787 | "print(\"P(L=b | R=l, B=r) = \"+str(np.mean(leader_bool * agent_1_actions * (1-agent_2_actions)) / np.mean(agent_1_actions * (1-agent_2_actions))))" 788 | ] 789 | }, 790 | { 791 | "cell_type": "code", 792 | "execution_count": null, 793 | "metadata": { 794 | "cellView": "form", 795 | "id": "nRg4gLxcBS4x" 796 | }, 797 | "outputs": [], 798 | "source": [ 799 | "# @title Creating leader and follower actions, with intervention do(R=r)\n", 800 | "\n", 801 | "forced_action = 0\n", 802 | "agent_1_actions, agent_2_actions = [], []\n", 803 | "leader_bool = []\n", 804 | "for episode in range(n_episodes):\n", 805 | " leader = np.random.randint(0, 2)\n", 806 | " leader_bool.append(leader)\n", 807 | "\n", 808 | " agent_1_actions.append(forced_action)\n", 809 | " leader_action = np.random.randint(0, 2) if leader == 1 else forced_action\n", 810 | " if leader == 1:\n", 811 | " agent_2_actions.append(leader_action)\n", 812 | " else:\n", 813 | " follower_action = leader_action\n", 814 | " if np.random.rand() \u003c 0.1:\n", 815 | " follower_action = np.random.randint(0, 2)\n", 816 | " agent_2_actions.append(follower_action)\n", 817 | "leader_bool = np.array(leader_bool)\n", 818 | "agent_1_actions = np.array(agent_1_actions)\n", 819 | "agent_2_actions = np.array(agent_2_actions)\n", 820 | "\n", 821 | "print(\"P(L=b | do(R=r), B=l) = \"+str(np.mean(leader_bool * agent_2_actions) / np.mean(agent_2_actions)))" 822 | ] 823 | }, 824 | { 825 | "cell_type": "code", 826 | "execution_count": null, 827 | "metadata": { 828 | "cellView": "form", 829 | "id": "-iZ6Pk6RGY4g" 830 | }, 831 | "outputs": [], 832 | "source": [ 833 | "# @title Creating leader and follower actions, with intervention do(R=l)\n", 834 | "\n", 835 | "forced_action = 1\n", 836 | "agent_1_actions, agent_2_actions = [], []\n", 837 | "leader_bool = []\n", 838 | "for episode in range(n_episodes):\n", 839 | " leader = np.random.randint(0, 2)\n", 840 | " leader_bool.append(leader)\n", 841 | "\n", 842 | " agent_1_actions.append(forced_action)\n", 843 | " leader_action = np.random.randint(0, 2) if leader == 1 else forced_action\n", 844 | " if leader == 1:\n", 845 | " agent_2_actions.append(leader_action)\n", 846 | " else:\n", 847 | " follower_action = leader_action\n", 848 | " if np.random.rand() \u003c 0.1:\n", 849 | " follower_action = np.random.randint(0, 2)\n", 850 | " agent_2_actions.append(follower_action)\n", 851 | "leader_bool = np.array(leader_bool)\n", 852 | "agent_1_actions = np.array(agent_1_actions)\n", 853 | "agent_2_actions = np.array(agent_2_actions)\n", 854 | "\n", 855 | "print(\"P(L=b | do(R=l), B=l) = \"+str(np.mean(leader_bool * agent_2_actions) / np.mean(agent_2_actions)))" 856 | ] 857 | }, 858 | { 859 | "cell_type": "markdown", 860 | "metadata": { 861 | "id": "W9dojnwUHdfI" 862 | }, 863 | "source": [ 864 | "# 6 - Causal pathways" 865 | ] 866 | }, 867 | { 868 | "cell_type": "code", 869 | "execution_count": null, 870 | "metadata": { 871 | "cellView": "form", 872 | "id": "Yaf0uSsMKIbc" 873 | }, 874 | "outputs": [], 875 | "source": [ 876 | "# @title Create interventions.\n", 877 | "\n", 878 | "def change_door_state(node, open: bool):\n", 879 | " door_pos = debugger.extractors.get_element_positions(node, Tile.DOOR_R)\n", 880 | " if door_pos == [] and open == True:\n", 881 | " return node\n", 882 | " if door_pos != [] and open == False:\n", 883 | " return node\n", 884 | " \n", 885 | " if open:\n", 886 | " return debugger.interventions.remove_drape_element(\n", 887 | " node, Tile.DOOR_R, position=(3, 4))\n", 888 | " return debugger.interventions.add_drape_element(\n", 889 | " node, Tile.DOOR_R, position=(3, 4))\n", 890 | " \n", 891 | "def change_internal_key_state(node: node_lib.Node, value: int) -\u003e node_lib.Node:\n", 892 | " with pcw_interv.InterventionContext(node) as context:\n", 893 | " context.engine.the_plot[chr(Tile.KEY_R)] = value\n", 894 | " return context.new_node\n", 895 | "\n", 896 | "def remove_key(node):\n", 897 | " node = change_internal_key_state(node=node, value=0)\n", 898 | " key_pos = debugger.extractors.get_element_positions(node, Tile.KEY_R)[0]\n", 899 | " return debugger.interventions.remove_drape_element(\n", 900 | " node, Tile.KEY_R, position=key_pos)\n", 901 | " \n", 902 | "def add_key_to_agent(node):\n", 903 | " node = remove_key(node)\n", 904 | " return change_internal_key_state(node, value=1)\n", 905 | " \n", 906 | "def door_state(node):\n", 907 | " open = debugger.extractors.get_element_positions(rollout[0], Tile.DOOR_R) == []\n", 908 | " return \"open\" if open else \"closed\"\n", 909 | "\n", 910 | "def key_taken(rollout):\n", 911 | " return (debugger.extractors.get_element_positions(rollout[-1], Tile.KEY_R) == [])\n", 912 | "\n", 913 | "def reward_taken(rollout):\n", 914 | " return (rollout[-1].last_timestep.reward == 1)" 915 | ] 916 | }, 917 | { 918 | "cell_type": "code", 919 | "execution_count": null, 920 | "metadata": { 921 | "cellView": "form", 922 | "id": "K08kGzskJoOQ" 923 | }, 924 | "outputs": [], 925 | "source": [ 926 | "# @title Create the debugger.\n", 927 | "\n", 928 | "agent = \"A\" #@param[\"A\", \"B\"]\n", 929 | "\n", 930 | "\n", 931 | "# Agent A - trained on all the environment distribution\n", 932 | "if agent == \"A\":\n", 933 | " params = load_params('pathways_a')\n", 934 | " env = pcw_env.build_environment(level='key_door')\n", 935 | "\n", 936 | "# Agent B - trained only when the door is closed\n", 937 | "if agent == \"B\":\n", 938 | " params = load_params('pathways_b')\n", 939 | " env = pcw_env.build_environment(level='key_door_closed')\n", 940 | "\n", 941 | "trained_agent = make_agent_from_params(params, with_lstm=True)\n", 942 | "\n", 943 | "debugger = pcw_dbg.PycoworldDebugger(trained_agent, env)\n", 944 | "init_nodes = prepare_init_nodes(debugger)" 945 | ] 946 | }, 947 | { 948 | "cell_type": "code", 949 | "execution_count": null, 950 | "metadata": { 951 | "id": "3jvHEdGKQIPE" 952 | }, 953 | "outputs": [], 954 | "source": [ 955 | "# @title Compute conditional probs.\n", 956 | "# @markdown Note that for agent B some probs are Nans, since the door is **always** closed in the environment, implying that the event that we condition on can have zero probability.\n", 957 | "\n", 958 | "reward_taken_l = []\n", 959 | "key_taken_l = []\n", 960 | "door_open_l = []\n", 961 | "for node in init_nodes:\n", 962 | " rollout = debugger.get_rollout(node, maximum_length=15)\n", 963 | " reward_taken_l.append(int(reward_taken(rollout)))\n", 964 | " key_taken_l.append(int(key_taken(rollout)))\n", 965 | " door_open_l.append(int(door_state(rollout) == \"open\"))\n", 966 | "key_taken_l = np.array(key_taken_l)\n", 967 | "reward_taken_l = np.array(reward_taken_l)\n", 968 | "door_open_l = np.array(door_open_l)\n", 969 | "\n", 970 | "print(\"P(R=1) = \"+str(np.mean(reward_taken_l)))\n", 971 | "print(\"P(R=1 | K=y) = \"+str(np.mean(reward_taken_l * key_taken_l) / np.mean(key_taken_l)))\n", 972 | "print(\"P(R=1 | K=n) = \"+str(np.mean(reward_taken_l * (1-key_taken_l)) / np.mean(1-key_taken_l)))\n", 973 | "print(\"P(R=1 | D=o) = \"+str(np.mean(reward_taken_l * door_open_l) / np.mean(door_open_l)))\n", 974 | "print(\"P(R=1 | D=c) = \"+str(np.mean(reward_taken_l * (1-door_open_l)) / np.mean(1-door_open_l)))" 975 | ] 976 | }, 977 | { 978 | "cell_type": "code", 979 | "execution_count": null, 980 | "metadata": { 981 | "cellView": "form", 982 | "id": "VvYrXU9HRbI2" 983 | }, 984 | "outputs": [], 985 | "source": [ 986 | "# @title Compute interventional probs.\n", 987 | "\n", 988 | "reward_taken_l = []\n", 989 | "for node in init_nodes:\n", 990 | " node = remove_key(node)\n", 991 | " rollout = debugger.get_rollout(node, maximum_length=15)\n", 992 | " reward_taken_l.append(int(reward_taken(rollout)))\n", 993 | "reward_taken_l = np.array(reward_taken_l)\n", 994 | "print(\"P(R=1 | do(K=n)) = \"+str(np.mean(reward_taken_l)))\n", 995 | "\n", 996 | "reward_taken_l = []\n", 997 | "for node in init_nodes:\n", 998 | " node = add_key_to_agent(node)\n", 999 | " rollout = debugger.get_rollout(node, maximum_length=15)\n", 1000 | " reward_taken_l.append(int(reward_taken(rollout)))\n", 1001 | "reward_taken_l = np.array(reward_taken_l)\n", 1002 | "print(\"P(R=1 | do(K=y)) = \"+str(np.mean(reward_taken_l)))\n", 1003 | "\n", 1004 | "key_taken_l = []\n", 1005 | "for node in init_nodes:\n", 1006 | " node = change_door_state(node, open=True)\n", 1007 | " rollout = debugger.get_rollout(node, maximum_length=15)\n", 1008 | " key_taken_l.append(int(key_taken(rollout)))\n", 1009 | "key_taken_l = np.array(key_taken_l)\n", 1010 | "print(\"P(K=y | do(D=o)) = \"+str(np.mean(key_taken_l)))\n", 1011 | "\n", 1012 | "key_taken_l = []\n", 1013 | "for node in init_nodes:\n", 1014 | " node = change_door_state(node, open=False)\n", 1015 | " rollout = debugger.get_rollout(node, maximum_length=15)\n", 1016 | " key_taken_l.append(int(key_taken(rollout)))\n", 1017 | "key_taken_l = np.array(key_taken_l)\n", 1018 | "print(\"P(K=y | do(D=c)) = \"+str(np.mean(key_taken_l)))" 1019 | ] 1020 | }, 1021 | { 1022 | "cell_type": "markdown", 1023 | "metadata": { 1024 | "id": "WbYAFNqGFtVN" 1025 | }, 1026 | "source": [ 1027 | "The causal response can be calculated manually using the above." 1028 | ] 1029 | } 1030 | ], 1031 | "metadata": { 1032 | "colab": { 1033 | "collapsed_sections": [], 1034 | "name": "Agent Debugger Experiments", 1035 | "private_outputs": true, 1036 | "provenance": [] 1037 | }, 1038 | "kernelspec": { 1039 | "display_name": "Python 3", 1040 | "name": "python3" 1041 | } 1042 | }, 1043 | "nbformat": 4, 1044 | "nbformat_minor": 0 1045 | } 1046 | --------------------------------------------------------------------------------