torchrl 0.11.0__cp314-cp314t-macosx_11_0_arm64.whl

torchrl/envs/llm/libs/mlgym.py

@@ -0,0 +1,869 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from __future__ import annotations
+
+import functools
+
+import json
+import os
+import re
+from contextlib import contextmanager
+from dataclasses import asdict
+from pathlib import Path
+from typing import Any, Literal, TYPE_CHECKING
+
+import numpy as np
+
+import torch
+from tensordict import NestedKey, NonTensorData, TensorDict, TensorDictBase
+from tensordict.tensorclass import is_non_tensor
+
+from torchrl._utils import logger as torchrl_logger
+from torchrl.data import Choice, Composite, NonTensor
+from torchrl.data.llm import History
+from torchrl.envs import ConditionalSkip, GymWrapper, Transform, TransformedEnv
+
+if TYPE_CHECKING:
+    import mlgym
+    import transformers
+
+# Inv transforms:
+#  Transforms to apply prior to pass the model output to the env
+
+
+@contextmanager
+def _temp_cwd_mlgym():
+    """Temporarily change the current working directory to mlgym."""
+    import mlgym
+
+    path = Path(mlgym.__spec__.submodule_search_locations[0]).parent
+    old_pwd = os.getcwd()
+    os.chdir(str(path))
+    # sys.path.insert(-1, "mlgym")
+    try:
+        yield
+    finally:
+        # sys.path.pop()
+        os.chdir(old_pwd)
+
+
+class MLGymBaseTransform(Transform):
+    """Base class for all MLGym transforms."""
+
+    @property
+    def config(self):
+        return self.parent.base_env.config
+
+    @property
+    def system_args(self):
+        return {
+            "command_docs": self.config.tools_handler.command_docs,
+            **self.config.tools_handler.env_variables,
+        }
+
+    @property
+    def task_args(self):
+        # Placeholder
+        task_args = getattr(self, "_task_args", None)
+        if task_args is None:
+            return self.parent.base_env.task.args
+        return task_args
+
+    @task_args.setter
+    def task_args(self, task_args):
+        self._task_args = task_args
+
+    @property
+    def name(self):
+        return "torchrl"
+
+    @property
+    def state_command(self):
+        return self.config.state_command.name
+
+    @property
+    def agent_args(self):
+        return self.parent.base_env.agent_args
+
+    @property
+    def model_name(self) -> Literal["human", "human_thought"]:
+        return self.agent_args.model.model_name
+
+
+#######################################################
+# Forward transforms: Format the env output
+
+
+# Transform #0: Resets the env
+class ResetModule(MLGymBaseTransform):
+    """Runs setup pipeline and enables multi-resets.
+
+    The reset method reads the 'system' initial input from the config and parses it to a History
+    object.
+
+    """
+
+    response_key: NestedKey = "text_response"
+
+    def __init__(self):
+        super().__init__(in_keys=[], out_keys=["history"])
+
+    @_temp_cwd_mlgym()
+    def _reset_env_preprocess(self, tensordict: TensorDictBase) -> TensorDictBase:
+        base_env = self.parent.base_env._env
+        if tensordict is not None and "task" in tensordict:
+            import gymnasium as gym
+
+            task = tensordict["task"]
+            torchrl_logger.info(f"Resetting with {task=}")
+            if is_non_tensor(task):
+                task = task.data
+            task_id, agent_args = _TASK_IDS[task]
+            try:
+                base_env.close()
+            except Exception:
+                torchrl_logger.info(f"Failed to close {base_env=}")
+            base_env = gym.make(
+                f"mlgym/{task}",
+                devices=["cpu_0"],
+            ).unwrapped
+            base_env.config = agent_args.config
+            self.parent.base_env.set_env(base_env)
+        base_env.reset_container()
+        base_env.communicate(f"cd {Path(base_env.task_workspace).parent}")
+        return tensordict
+
+    @_temp_cwd_mlgym()
+    def _reset(
+        self, tensordict: TensorDictBase, tensordict_reset: TensorDictBase
+    ) -> TensorDictBase:
+        # TODO: what to do with this?
+        # reset model stats
+        # self.model.reset_stats(init_model_stats)
+        # env = self.parent.base_env._env
+
+        env = self.parent.base_env._env
+        self.set_environment_vars(env, self.config.env_variables)
+
+        system_msg = self.config.system_template.format(
+            **self.system_args, **asdict(self.task_args)
+        )
+        # self.logger.log(self._default_logging_level, f"SYSTEM ({self.name})\n{system_msg}")
+        history = History(
+            role="system",
+            content=system_msg,  # agent=self.name,
+            batch_size=(1,),
+            device=self.parent.device,
+        )
+        tensordict_reset["history"] = history
+
+        return tensordict_reset
+
+    def _step(
+        self, tensordict: TensorDictBase, next_tensordict: TensorDictBase
+    ) -> TensorDictBase:
+        # Placeholder
+        if "history" not in next_tensordict:
+            if "local_history" in tensordict:
+                local_history = tensordict["local_history"]
+            else:
+                local_history = None
+            history = tensordict["history"]
+            if local_history is not None:
+                history = history.append(local_history, inplace=False)
+                tensordict["history"] = history
+            next_tensordict["history"] = history
+        return next_tensordict
+
+    def set_environment_vars(
+        self, env: MLGymWrapper, env_variables: dict[str, Any]
+    ) -> None:
+        commands_to_execute = (
+            [self.config.state_command.code]
+            +  # [code for code in self.config.util_functions] +
+            # [command.code for command in self.config._commands] +
+            [f"{k}={v}" for k, v in env_variables.items()]
+        )
+        commands = "\n".join(commands_to_execute)
+        try:
+            output = env.communicate(commands)
+            if env.returncode != 0:
+                msg = f"Nonzero return code: {env.returncode}\nOutput: {output}"
+                raise RuntimeError(msg)
+        except KeyboardInterrupt:
+            raise
+        except Exception as e:
+            raise e
+        command_files = []
+        for file in self.config.command_files:
+            datum = {}
+            with open(file) as f:
+                contents = f.read()
+            datum["contents"] = contents
+            filename = Path(file).name
+            if not contents.strip().startswith("#!"):
+                if filename.endswith(".sh"):
+                    # files are sourced, so they are not executable
+                    datum["name"] = Path(file).name
+                    datum["type"] = "source_file"
+                elif filename.startswith("_"):
+                    # files are sourced, so they are not executable
+                    datum["name"] = Path(file).name
+                    datum["type"] = "utility"
+                else:
+                    msg = (
+                        f"Non-shell script file {file} does not start with shebang.\n"
+                        "Either add a shebang (#!) or change the file extension to .sh if you want to source it.\n"
+                        "You can override this behavior by adding an underscore to the file name (e.g. _utils.py)."
+                    )
+                    raise ValueError(msg)
+            else:
+                # scripts are made executable
+                datum["name"] = Path(file).name.rsplit(".", 1)[0]
+                datum["type"] = "script"
+            command_files.append(datum)
+        # TODO: implement add commands method in environment
+        env.add_commands(command_files)
+
+    def transform_observation_spec(self, observation_spec: Composite) -> Composite:
+        observation_spec["history"] = History.default_spec()
+        return observation_spec
+
+    def transform_action_spec(self, action_spec: Composite) -> Composite:
+        if isinstance(action_spec, Composite):
+            action_spec[self.response_key] = self.transform_action_spec(
+                action_spec[self.response_key]
+            )
+            return action_spec
+        # make the "random" action just a choice between innocuous bash commands
+        return Choice(
+            [
+                NonTensor(example_data="ls -rtlh", shape=action_spec.shape),
+                NonTensor(example_data="pwd", shape=action_spec.shape),
+            ]
+        )
+
+    def transform_state_spec(self, state_spec: Composite) -> Composite:
+        state_spec["history"] = History.default_spec()
+        return state_spec
+
+
+class TaskSampler(Transform):
+    """A sampler for tasks in a certain task set."""
+
+    def __init__(self, tasks: list[str]):
+        super().__init__()
+        self.tasks = tasks
+
+    def transform_observation_spec(self, observation_spec: Composite) -> Composite:
+        observation_spec["task"] = NonTensor(example_data="<a task>", shape=())
+        return observation_spec
+
+    @_temp_cwd_mlgym()
+    def _reset_env_preprocess(
+        self, tensordict: TensorDictBase | None
+    ) -> TensorDictBase:
+        if tensordict is None:
+            tensordict = TensorDict(batch_size=self.parent.batch_size)
+        # Sample a task
+        task = np.random.choice(self.tasks)
+        tensordict["task"] = NonTensorData(task)
+        self._current_task = task
+        return tensordict
+
+    def _call(self, next_tensordict: TensorDictBase) -> TensorDictBase:
+        next_tensordict["task"] = self._current_task
+        return next_tensordict
+
+
+# Transform #1: env -> state
+class ReadState(MLGymBaseTransform):
+    """Reads current state and writes it as a parsable str in the tensordict."""
+
+    # from mlgym/agent/base.py:BaseAgent:forward_model
+    def _step(
+        self, tensordict: TensorDictBase, next_tensordict: TensorDictBase
+    ) -> TensorDictBase:
+        base_mlgym_env = self.parent.base_env  # getattr is forwarded
+
+        command = self.state_command
+        state = base_mlgym_env.communicate(command) if self.state_command else None
+
+        next_tensordict["state"] = state
+        return next_tensordict
+
+    def _reset(
+        self, tensordict: TensorDictBase, tensordict_reset: TensorDictBase
+    ) -> TensorDictBase:
+        # tensordict_reset.setdefault("message", NonTensorData(""))
+        # tensordict_reset.setdefault("state", NonTensorData(""))
+        return self._step(tensordict_reset, tensordict_reset)
+
+    def transform_observation_spec(self, observation_spec):
+        observation_spec.set(
+            "state",
+            NonTensor(
+                example_data="a string",
+                device=observation_spec.device,
+                shape=observation_spec.shape,
+            ),
+        )
+        return observation_spec
+
+
+# Transform #2: state -> message
+class StateToMessage(MLGymBaseTransform):
+    """Parses the string using json to a given template.
+
+    Requires:
+        - a 'state' key from the ReadState transform
+        - an 'observation' key from the base environment
+
+    """
+
+    def _step(
+        self, tensordict: TensorDictBase, next_tensordict: TensorDictBase
+    ) -> TensorDictBase:
+        base_mlgym_env = self.parent.base_env  # getattr is forwarded
+        observation = tensordict["observation"]
+        state = tensordict["state"]
+        config = self.config
+
+        current_step = base_mlgym_env.current_step
+        max_steps = base_mlgym_env.max_steps
+        try:
+            state_vars = json.loads(state)
+        except json.JSONDecodeError as e:
+            msg = f"State {state!r} is not valid json. This is an internal error, please report it."
+            raise ValueError(msg) from e
+        # add step information to state_vars
+        state_vars["current_step"] = current_step
+        state_vars["remaining_steps"] = max_steps - current_step
+
+        # FIXME: we don't need to do this, we have our own observation space
+        # Determine observation template based on what prior observation was
+
+        history: History = tensordict["history"]
+        if history[..., -1].role == "system":
+            # Show task template if prev. obs. was initial system message
+            templates = [config.task_template]
+            if config.strategy_template is not None:
+                templates.append(config.strategy_template)
+        elif observation is None or observation.strip() == "":
+            # Show no output template if observation content was empty
+            assert config.next_step_no_output_template is not None  # linting
+            templates = [config.next_step_no_output_template]
+        else:
+            # Show standard output template if there is observation content
+            assert config.next_step_template is not None  # linting
+            templates = [config.next_step_template]
+
+        # Format selected template(s) with information
+        messages = []
+        assert self.task_args is not None
+        for template in templates:
+            messages.append(
+                template.format(
+                    **asdict(self.task_args),
+                    **self.system_args,
+                    **state_vars,
+                    observation=(observation if observation is not None else ""),
+                    # missing forwarded_vars because no attempts
+                ),
+            )
+
+        message = "\n".join(messages)
+        next_tensordict["message"] = message
+        # model query hooks here
+        return next_tensordict
+
+    def _reset(
+        self, tensordict: TensorDictBase, tensordict_reset: TensorDictBase
+    ) -> TensorDictBase:
+        # tensordict_reset.setdefault("message", NonTensorData(""))
+        # tensordict_reset.setdefault("state", NonTensorData(""))
+        return self._step(tensordict_reset, tensordict_reset)
+
+    def transform_observation_spec(self, observation_spec):
+        observation_spec.set(
+            "message",
+            NonTensor(
+                example_data="a string",
+                device=observation_spec.device,
+                shape=observation_spec.shape,
+            ),
+        )
+        return observation_spec
+
+
+# Transform #3: Append message to history
+class MessageToHistory(MLGymBaseTransform):
+    """Parses the message string to a History object, then reparses the history to a complete message.
+
+    .. seealso:: HistoryToMessage
+
+    """
+
+    def __init__(self):
+        super().__init__(in_keys=["message", "history"], out_keys=["history", "chat"])
+
+    # from mlgym/agent/base.py:BaseAgent:local_history
+    # from mlgym/agent/base.py:BaseAgent:_append_history
+    def _step(
+        self, tensordict: TensorDictBase, next_tensordict: TensorDictBase
+    ) -> TensorDictBase:
+        # From PrepareDataForModel
+        message: str = next_tensordict["message"]
+        # from mlgym/agent/base.py:BaseAgent:forward_model
+        history = tensordict["history"]
+        cur_history = History(
+            role="user", content=message, batch_size=(), device=self.parent.device
+        )
+        # This is the basic thing our transform does: append the history to the existing one.
+        # (We should be able to extend the lazy stack directly)
+        history = history.append(cur_history, inplace=False)
+
+        next_tensordict["history"] = history
+        return next_tensordict
+
+    def _reset(
+        self, tensordict: TensorDictBase, tensordict_reset: TensorDictBase
+    ) -> TensorDictBase:
+        return self._step(tensordict_reset, tensordict_reset)
+
+
+# Inverse transforms:
+#  Format the action from the model for the env
+
+
+class TemplateTransform(MLGymBaseTransform):
+    """A transform to apply the chat template to the History."""
+
+    response_key: NestedKey = "text_response"
+    prompt_key: NestedKey = "text"
+
+    # alternative to DummyFormat, wip
+    def __init__(
+        self,
+        in_keys=None,
+        out_keys=None,
+        in_keys_inv=None,
+        out_keys_inv=None,
+        tokenizer=None,
+        chat_template_name: Literal["chatml_format"] | None = None,
+        continue_final_message: bool = False,
+        tokenize: bool = False,
+        return_tensors: str = "pt",
+        return_dict: bool = False,
+        padding: bool | str = False,
+        truncation: bool | str = False,
+    ):
+        super().__init__(
+            in_keys=["history"] if in_keys is None else in_keys,
+            out_keys=[self.prompt_key] if out_keys is None else out_keys,
+            in_keys_inv=[self.prompt_key, self.response_key]
+            if in_keys_inv is None
+            else in_keys_inv,
+            # TODO: we should not use the response key here but another dedicated entry, like "action_parsed"
+            out_keys_inv=[self.response_key] if out_keys_inv is None else out_keys_inv,
+        )
+        self.chat_template_name = chat_template_name
+        self.tokenizer = tokenizer
+        self.tokenize = tokenize
+        self.continue_final_message = continue_final_message
+        self.return_tensors = return_tensors
+        self.return_dict = return_dict
+        self.padding = padding
+        self.truncation = truncation
+
+    def transform_observation_spec(self, observation_spec: Composite):
+        observation_spec[self.prompt_key] = NonTensor(
+            example_data="<some chat string>",
+            shape=observation_spec.shape,
+            device=observation_spec.device,
+        )
+        return observation_spec
+
+    @property
+    def _chat_template(self):
+        chat_template = None
+        if self.chat_template_name:
+            from torchrl.data.llm.datatypes.chat import _CHAT_TEMPLATES
+
+            chat_template = _CHAT_TEMPLATES[self.chat_template_name]
+        elif self.tokenizer.chat_template is not None:
+            chat_template = self.tokenizer.chat_template
+        elif chat_template is None:
+            raise ValueError("Failed to determine chat template.")
+        return chat_template
+
+    def _apply_transform(self, history: History) -> NonTensorData:
+        if self.tokenizer is None:
+            raise RuntimeError("Cannot apply chat template without a tokenizer.")
+        result = history.apply_chat_template(
+            tokenizer=self.tokenizer,
+            add_generation_prompt=True,
+            chat_template=self._chat_template,
+            continue_final_message=self.continue_final_message,
+            tokenize=self.tokenize,
+            padding=self.padding,
+            truncation=self.truncation,
+            return_tensors=self.return_tensors,
+        )
+        return result
+
+    def _reset(self, tensordict, tensordict_reset):
+        return self._call(tensordict_reset)
+
+    def _inv_call(self, tensordict: TensorDictBase) -> TensorDictBase:
+        if self.in_keys_inv:
+            prompt = tensordict[self.prompt_key]
+            response = tensordict[self.response_key]
+            if isinstance(prompt, list):
+                action = [
+                    prompt + response for prompt, response in zip(prompt, response)
+                ]
+            else:
+                action = prompt + response
+            try:
+                history, action = self._inv_apply_transform(action)
+                tensordict["local_history"] = history
+                tensordict[self.response_key] = action
+            except RuntimeError as e:
+                if "Expected assistant role" in str(e):
+                    tensordict["local_history"] = History(role="assistant", content="")
+                    tensordict[self.response_key] = ""
+        return tensordict
+
+    def _inv_apply_transform(self, action):
+        if self.tokenize:
+            action = self.tokenizer.decode(action)
+
+        if not isinstance(action, (str, list)):
+            action = action.data
+            history, action = self._inv_apply_transform(action)
+            action = NonTensorData(
+                action, batch_size=action.batch_size, device=action.device
+            )
+            return history, action
+
+        history = History.from_text(
+            action,
+            # chat_template=self._chat_template,
+        )[..., -1]
+        if history.role != "assistant":
+            raise RuntimeError(f"Expected assistant role, got {history.role=}")
+        action = history.get("content")
+        return history, action
+
+
+class IsolateCodeBlock(MLGymBaseTransform):
+    """A transform that isolates the code block in the action generated by the LLM.
+
+    Optionally, wrongly formatted actions are assigned a negative reward.
+    """
+
+    response_key: NestedKey = "text_response"
+
+    def __init__(self, reward_wrong_format: float | None = None):
+        super().__init__(
+            in_keys_inv=[self.response_key], out_keys_inv=[self.response_key]
+        )
+        from mlgym.agent.parsing import ThoughtActionParser
+
+        self.parser = ThoughtActionParser()
+        self.reward_wrong_format = reward_wrong_format
+        self._assign_reward = False
+
+    def _inv_call(self, tensordict: TensorDictBase) -> TensorDictBase:
+        torchrl_logger.info("inv call with IsolateCodeBlock")
+        action = tensordict[self.response_key]
+        # if we didn't find an action, the action is empty
+        if not action:
+            torchrl_logger.info(
+                "Did not find a suitable action, skipping the call to step."
+            )
+            tensordict["retry"] = torch.ones(tensordict.shape, dtype=torch.bool)
+            self._assign_reward = True
+        else:
+            from mlgym.exceptions import FormatError
+
+            try:
+                action = self._inv_apply_transform(action)
+                tensordict[self.response_key] = action
+                torchrl_logger.info(f"Code block: {action}")
+                tensordict["retry"] = torch.zeros(tensordict.shape, dtype=torch.bool)
+                self._assign_reward = False
+            except FormatError:
+                tensordict["retry"] = torch.ones(tensordict.shape, dtype=torch.bool)
+                self._assign_reward = True
+        return tensordict
+
+    def _call(self, tensordict: TensorDictBase) -> TensorDictBase:
+        if self._assign_reward:
+            torchrl_logger.info(
+                f"Assigning penalty for unsuitable action: {self.reward_wrong_format}"
+            )
+            if self.reward_wrong_format is not None:
+                tensordict[self.parent.reward_key] += self.reward_wrong_format
+        return tensordict
+
+    def _inv_apply_transform(self, action):
+        if not isinstance(action, (str, list)):
+            return NonTensorData(
+                self._inv_apply_transform(action.tolist()),
+                batch_size=action.batch_size,
+                device=action.device,
+            )
+        if isinstance(action, list):
+            return [self._inv_apply_transform(action) for action in action]
+        thought, action = self.parser(action, None)
+        return action
+
+
+class EvaluationOutputParser:
+    """Parser for the reward transform in MLGym.
+
+    .. seealso:: :class:`~torchrl.envs.llm.libs.mlgym.MLGymRewardAssignment`
+
+    """
+
+    def __init__(self):
+        # Regular expressions to match the required fields
+        self.patterns = {
+            "submission_artefact_path": r"valid submission artefact at (.*)\.",
+            "baseline_score": r"Baseline Score: \{'Score': (.*)\}",
+            "evaluation_score": r"Evaluation Score: \{'Score': (.*)\}",
+            "current_step": r"\(Current Step: (\d+),",
+            "remaining_steps": r"Remaining Steps: (\d+)\)",
+            "open_file": r"\(Open file: (.*)\)",
+            "current_directory": r"\(Current directory: (.*)\)",
+        }
+
+    def __call__(self, output_string):
+
+        parsed_data = {}
+
+        for key, pattern in self.patterns.items():
+            match = re.search(pattern, output_string)
+            if match:
+                parsed_data[key] = match.group(1).strip()
+        if "baseline_score" in parsed_data:
+            parsed_data["baseline_score"] = float(parsed_data["baseline_score"])
+
+        if "evaluation_score" in parsed_data:
+            parsed_data["evaluation_score"] = float(parsed_data["evaluation_score"])
+        if "current_step" in parsed_data:
+            parsed_data["current_step"] = int(parsed_data["current_step"])
+        if "remaining_steps" in parsed_data:
+            parsed_data["remaining_steps"] = int(parsed_data["remaining_steps"])
+
+        return parsed_data
+
+
+class MLGymRewardAssignment(MLGymBaseTransform):
+    """Reward assignment through parsing of the last item in history.
+
+    By default, the :class:`~torchrl.envs.llm.libs.mlgym.EvaluationOutputParser` class is used as parser.
+
+    """
+
+    def __init__(self):
+        super().__init__(in_keys=["reward", "history"], out_keys=["reward"])
+        self.parser = EvaluationOutputParser()
+
+    def _call(self, tensordict):
+        history = tensordict.get("history")
+        if history is None:
+            raise KeyError(f"History is missing in tensordict {tensordict}")
+        if history.ndim != 1:
+            raise ValueError(f"History shape must be 1D, got {history.shape}")
+        content = history[-1].content
+        torchrl_logger.info(f"Parsing reward from: {content}")
+        parsed = self.parser(content)
+        reward = parsed.get("evaluation_score", 0.0) - parsed.get("baseline_score", 0.0)
+        torchrl_logger.info(f"Parsed reward: {reward}")
+        tensordict["reward"] = tensordict["reward"] + reward
+        return tensordict
+
+
+class _add_info_to_reset:
+    def __init__(self, func):
+        functools.update_wrapper(self, func)
+        self.func = func
+
+    def __call__(self, *args, **kwargs):
+        return self.func(*args, **kwargs), {}
+
+
+class _add_truncated_to_step:
+    def __init__(self, func):
+        functools.update_wrapper(self, func)
+        self.func = func
+
+    @_temp_cwd_mlgym()
+    def __call__(self, *args, **kwargs):
+        obs, r, done, info = self.func(*args, **kwargs)
+        return obs, r, done, False, info
+
+
+class MLGymWrapper(GymWrapper):
+    """A thin wrapper for MLGym environments.
+
+    This specialized :class:`~torchrl.envs.GymWrapper` subclass defines the observation space with `observation=NonTensor()`
+    and the action space with `text_response=NonTensor()`, according to the :class:`~torchrl.envs.llm.ChatEnv` API.
+
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.full_action_spec = Composite(
+            text_response=NonTensor(example_data="<a string>", shape=())
+        )
+        self.full_observation_spec = Composite(
+            observation=NonTensor(example_data="<a string>", shape=())
+        )
+        self.set_env()
+
+    def set_env(self, env: Any = None):
+        if env is not None:
+            self._env = env
+        self._patch_reset()
+        self._patch_step()
+
+    def _patch_reset(self):
+        if not isinstance(self._env.reset, _add_info_to_reset):
+            self._env.reset = _add_info_to_reset(self._env.reset)
+
+    def _patch_step(self):
+        if not isinstance(self._env.reset, _add_truncated_to_step):
+            self._env.step = _add_truncated_to_step(self._env.step)
+
+    @_temp_cwd_mlgym()
+    def _reset(
+        self, tensordict: TensorDictBase | None = None, **kwargs
+    ) -> TensorDictBase:
+        return super()._reset(tensordict=tensordict, **kwargs)
+
+
+_TASK_IDS = {}
+
+
+def get_args(
+    task: Literal["prisonersDilemma"] = "prisonersDilemma",
+) -> tuple[
+    mlgym.environment.env.EnvironmentArguments,  # noqa
+    mlgym.agent.base.AgentArguments,  # noqa
+]:  # noqa
+    """Parse command line arguments and return a ScriptArguments object.
+
+    Args:
+        args: Optional list of arguments to parse. If not provided, uses sys.argv.
+    """
+    import mlgym.environment.registration  # noqa
+    from mlgym import CONFIG_DIR
+    from mlgym.agent.base import AgentArguments
+    from mlgym.backend.base import ModelArguments
+    from mlgym.environment.env import EnvironmentArguments
+    from mlgym.environment.registration import register_task
+
+    environment_args = EnvironmentArguments(
+        task_config_path=f"tasks/{task}.yaml",
+        max_steps=10,
+        seed=42,
+        container_type="podman",
+        verbose=False,
+        aliases_file="docker/aliases.sh",
+    )
+
+    agent_args = AgentArguments(
+        # placeholder
+        model=ModelArguments(""),
+        # Despite using torchrl as an agent, we still need the agent config - see StateToMessage parser
+        agent_config_path=CONFIG_DIR / "agents" / "default.yaml",
+    )
+
+    register_task(environment_args)
+
+    _TASK_IDS[task] = (environment_args.task.id, agent_args)
+
+    return environment_args, agent_args
+
+
+def make_mlgym(
+    *,
+    task: Literal["prisonersDilemma"] | None = None,
+    tasks: list[Literal["prisonersDilemma"]] | None = None,
+    tokenizer: transformers.AutoTokenizer | str | None = None,  # noqa
+    device="cpu",
+    reward_wrong_format: float | None = None,
+) -> TransformedEnv:
+    """Wraps an MLGymEnv in a TorchRL Environment.
+
+    The appended transforms will make sure that the data is formatted for the LLM during (for the outputs of `env.step`)
+    and for the MLGym API (for inputs to `env.step`).
+
+    Keyword Args:
+        task (str): The task to wrap. Exclusive with `tasks` argument.
+
+            .. note:: The correct format is simply the task name, e.g., `"prisonersDilemma"`.
+
+        tasks (List[str]): The tasks available for the env. Exclusive with `task` argument.
+
+            .. note:: The correct format is simply the task name, e.g., `"prisonersDilemma"`.
+
+        tokenizer (transformers.AutoTokenizer or str, optional): A transformer that tokenizes the data.
+            If a string is passed, it will be converted to a `transformers.AutoTokenizer`.
+        device (str, optional): The device to set to the env. Defaults to "cpu".
+        reward_wrong_format (float, optional): The reward (negative penalty) for wrongly formatted actions.
+            Defaults to `None` (no penalty).
+
+    """
+    import gymnasium as gym
+
+    if isinstance(tokenizer, str):
+        import transformers
+
+        tokenizer = transformers.AutoTokenizer.from_pretrained(tokenizer)
+
+    with _temp_cwd_mlgym():
+
+        if task and not tasks:
+            environment_args, agent_args = get_args(task=task)
+        elif tasks and not task:
+            for task in tasks:
+                environment_args, agent_args = get_args(task=task)
+        else:
+            raise ValueError(
+                f"Either task or tasks should be provided, not both and not none. Got {task=} and {tasks=}."
+            )
+
+        base_env = gym.make(
+            f"mlgym/{_TASK_IDS[task][0]}",
+            devices=["cpu_0"],
+        ).unwrapped
+        # we need the env to have access to the config
+        base_env.config = agent_args.config
+        env = TransformedEnv(
+            MLGymWrapper(base_env, auto_reset=False, device=device), auto_unwrap=False
+        )
+
+        env.append_transform(ConditionalSkip(lambda td: td["retry"]))
+        env.append_transform(IsolateCodeBlock(reward_wrong_format=reward_wrong_format))
+
+        env.append_transform(ResetModule())
+        if tasks:
+            # Add a task sampler
+            env.append_transform(TaskSampler(tasks))
+        env.append_transform(ReadState())
+        env.append_transform(StateToMessage())
+        env.append_transform(MessageToHistory())
+        env.append_transform(TemplateTransform(tokenizer=tokenizer))
+        env.append_transform(MLGymRewardAssignment())
+        # # We want the env to have a batch-size of (1,) because it will be easier to interact with
+        # #  LLMs
+        # env.append_transform(BatchSizeTransform(batch_size=(1,)))
+        return env