Skip to content

Agent

Agent

Abstract base class for agents.

This class provides a template for creating agents that can optionally record their actions and observations.

Attributes:

Name Type Description
recorder Recorder

The recorder to record observations and actions.
actor Backend

The backend actor to perform actions.
kwargs dict

Additional arguments to pass to the recorder.
Source code in mbodied/agents/agent.py 
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
class Agent:
    """Abstract base class for agents.

    This class provides a template for creating agents that can
    optionally record their actions and observations.

    Attributes:
        recorder (Recorder): The recorder to record observations and actions.
        actor (Backend): The backend actor to perform actions.
        kwargs (dict): Additional arguments to pass to the recorder.
    """

    ACTOR_MAP = {
        "openai": OpenAIBackend,
        "anthropic": AnthropicBackend,
        "ollama": OllamaBackend,
        "gradio": GradioBackend,
        "http": HttpxBackend,
    }

    @staticmethod
    def init_backend(model_src: str, model_kwargs: dict, api_key: str) -> type:
        """Initialize the backend based on the model source.

        Args:
            model_src: The model source to use.
            model_kwargs: The additional arguments to pass to the model.
            api_key: The API key to use for the remote actor.

        Returns:
            type: The backend class to use.
        """
        if model_src in Agent.ACTOR_MAP:
            if model_src == "gradio":
                # Gradio doesn't take api_key.
                return Agent.ACTOR_MAP[model_src](**model_kwargs)
            else:
                return Agent.ACTOR_MAP[model_src](**model_kwargs, api_key=api_key)
        return Agent.handle_default(model_src, model_kwargs)

    @staticmethod
    def handle_default(model_src: str, model_kwargs: dict) -> None:
        """Default to gradio then httpx backend if the model source is not recognized.

        Args:
            model_src: The model source to use.
            model_kwargs: The additional arguments to pass to the model.
        """
        try:
            return GradioBackend(endpoint=model_src, **model_kwargs)
        except Exception as e:
            logging.error(
                f"Failed to initialize Gradio backend: {e}. Defaulting to Httpx backend. Ensure that the source is a valid http endpoint.",
            )
            try:
                return HttpxBackend(endpoint=model_src, **model_kwargs)
            except Exception as e:
                logging.error(f"Failed to initialize Httpx backend: {e}.")
                raise ValueError(
                    f"Failed to initialize backend for model source: {model_src}. Pleases select one of {Agent.ACTOR_MAP.keys()} or valid huggingface space or http endpoint.",
                )

    def __init__(
        self,
        recorder: Literal["omit", "auto"] | str = "omit",
        recorder_kwargs=None,
        api_key: str = None,
        model_src=None,
        model_kwargs=None,
    ):
        """Initialize the agent, optionally setting up a recorder, remote actor, or loading a local model.

        Args:
            recorder: The recorder config or name to use for recording observations and actions.
            recorder_kwargs: Additional arguments to pass to the recorder.
            api_key: The API key to use for the remote actor (if applicable).
            model_src: The model or inference client or weights path to setup and preload if applicable.
                       You can pass in for example, "openai", "anthropic", "gradio", or a gradio endpoint,
                       or a path to a weights file.
            model_kwargs: Additional arguments to pass to the remote actor.
        """
        if model_src is None:
            raise ValueError("Model source must be provided.")
        if not isinstance(model_src, str):
            raise ValueError("Model source must be a string.")

        self.recorder = None
        recorder_kwargs = recorder_kwargs or {}
        if recorder == "auto":
            self.recorder = Recorder("base_agent", out_dir="outs", **recorder_kwargs)
        elif recorder != "omit":
            self.recorder = Recorder(recorder, out_dir="outs", **recorder_kwargs)

        model_kwargs = model_kwargs or {}
        self.actor = None
        if isinstance(model_src, str) and Path(model_src[:120]).exists():
            self.load_model(model_src, **model_kwargs)
        else:
            self.actor: Backend = self.init_backend(model_src, model_kwargs, api_key)

    def load_model(self, model: str) -> None:
        """Load a model from a file or path. Required if the model is a weights path.

        Args:
            model: The path to the model file.
        """
        pass

    def act(self, *args, **kwargs) -> Sample:
        """Act based on the observation.

        Subclass should implement this method.

        For remote actors, this method should call actor.act() correctly to perform the actions.
        """
        raise NotImplementedError("Subclass should implement this method.")

    async def async_act(self, *args, **kwargs) -> Sample:
        """Act asynchronously based on the observation.

        Subclass should implement this method.

        For remote actors, this method should call actor.async_act() correctly to perform the actions.
        """
        return await asyncio.to_thread(self.act, *args, **kwargs)

    def act_and_record(self, *args, **kwargs) -> Sample:
        """Peform action based on the observation and record the action, if applicable.

        Args:
            *args: Additional arguments to customize the action.
            **kwargs: Additional arguments to customize the action.

        Returns:
            Sample: The action sample created by the agent.
        """
        action = self.act(*args, **kwargs)
        if self.recorder is not None:
            observation = self.create_observation_from_args(
                self.recorder.observation_space,
                self.act_and_record,
                args,
                kwargs,
            )
            self.recorder.record(observation=observation, action=action)
        return action

    async def async_act_and_record(self, *args, **kwargs) -> Sample:
        """Act asynchronously based on the observation.

        Subclass should implement this method.

        For remote actors, this method should call actor.async_act() correctly to perform the actions.
        """
        return await asyncio.to_thread(self.act_and_record, *args, **kwargs)

    @staticmethod
    def create_observation_from_args(observation_space, function, args, kwargs) -> dict:
        """Helper method to create an observation from the arguments of a function."""
        param_names = list(signature(function).parameters.keys())

        # Create the observation from the arguments
        params = {**kwargs}
        for arg, val in zip(param_names, args, strict=False):
            params[arg] = val
        if observation_space is not None:
            observation = observation_space.sample()
            return {k: v for k, v in params.items() if k in observation}

        return {k: v for k, v in params.items() if v is not None and k not in ["self", "kwargs"]}

__init__(recorder='omit', recorder_kwargs=None, api_key=None, model_src=None, model_kwargs=None)

Initialize the agent, optionally setting up a recorder, remote actor, or loading a local model.

Parameters:

Name Type Description Default
recorder Literal['omit', 'auto'] | str

The recorder config or name to use for recording observations and actions.

 'omit'
recorder_kwargs

Additional arguments to pass to the recorder.

 None
api_key str

The API key to use for the remote actor (if applicable).

 None
model_src

The model or inference client or weights path to setup and preload if applicable. You can pass in for example, "openai", "anthropic", "gradio", or a gradio endpoint, or a path to a weights file.

 None
model_kwargs

Additional arguments to pass to the remote actor.

 None
Source code in mbodied/agents/agent.py 
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
def __init__(
    self,
    recorder: Literal["omit", "auto"] | str = "omit",
    recorder_kwargs=None,
    api_key: str = None,
    model_src=None,
    model_kwargs=None,
):
    """Initialize the agent, optionally setting up a recorder, remote actor, or loading a local model.

    Args:
        recorder: The recorder config or name to use for recording observations and actions.
        recorder_kwargs: Additional arguments to pass to the recorder.
        api_key: The API key to use for the remote actor (if applicable).
        model_src: The model or inference client or weights path to setup and preload if applicable.
                   You can pass in for example, "openai", "anthropic", "gradio", or a gradio endpoint,
                   or a path to a weights file.
        model_kwargs: Additional arguments to pass to the remote actor.
    """
    if model_src is None:
        raise ValueError("Model source must be provided.")
    if not isinstance(model_src, str):
        raise ValueError("Model source must be a string.")

    self.recorder = None
    recorder_kwargs = recorder_kwargs or {}
    if recorder == "auto":
        self.recorder = Recorder("base_agent", out_dir="outs", **recorder_kwargs)
    elif recorder != "omit":
        self.recorder = Recorder(recorder, out_dir="outs", **recorder_kwargs)

    model_kwargs = model_kwargs or {}
    self.actor = None
    if isinstance(model_src, str) and Path(model_src[:120]).exists():
        self.load_model(model_src, **model_kwargs)
    else:
        self.actor: Backend = self.init_backend(model_src, model_kwargs, api_key)

act(*args, **kwargs)

Act based on the observation.

Subclass should implement this method.

For remote actors, this method should call actor.act() correctly to perform the actions.

Source code in mbodied/agents/agent.py 
137
138
139
140
141
142
143
144
def act(self, *args, **kwargs) -> Sample:
    """Act based on the observation.

    Subclass should implement this method.

    For remote actors, this method should call actor.act() correctly to perform the actions.
    """
    raise NotImplementedError("Subclass should implement this method.")

act_and_record(*args, **kwargs)

Peform action based on the observation and record the action, if applicable.

Parameters:

Name Type Description Default
*args

Additional arguments to customize the action.

 ()
**kwargs

Additional arguments to customize the action.

 {}

Returns:

Name Type Description
Sample Sample

The action sample created by the agent.
Source code in mbodied/agents/agent.py 
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
def act_and_record(self, *args, **kwargs) -> Sample:
    """Peform action based on the observation and record the action, if applicable.

    Args:
        *args: Additional arguments to customize the action.
        **kwargs: Additional arguments to customize the action.

    Returns:
        Sample: The action sample created by the agent.
    """
    action = self.act(*args, **kwargs)
    if self.recorder is not None:
        observation = self.create_observation_from_args(
            self.recorder.observation_space,
            self.act_and_record,
            args,
            kwargs,
        )
        self.recorder.record(observation=observation, action=action)
    return action

async_act(*args, **kwargs) async

Act asynchronously based on the observation.

Subclass should implement this method.

For remote actors, this method should call actor.async_act() correctly to perform the actions.

Source code in mbodied/agents/agent.py 
146
147
148
149
150
151
152
153
async def async_act(self, *args, **kwargs) -> Sample:
    """Act asynchronously based on the observation.

    Subclass should implement this method.

    For remote actors, this method should call actor.async_act() correctly to perform the actions.
    """
    return await asyncio.to_thread(self.act, *args, **kwargs)

async_act_and_record(*args, **kwargs) async

Act asynchronously based on the observation.

Subclass should implement this method.

For remote actors, this method should call actor.async_act() correctly to perform the actions.

Source code in mbodied/agents/agent.py 
176
177
178
179
180
181
182
183
async def async_act_and_record(self, *args, **kwargs) -> Sample:
    """Act asynchronously based on the observation.

    Subclass should implement this method.

    For remote actors, this method should call actor.async_act() correctly to perform the actions.
    """
    return await asyncio.to_thread(self.act_and_record, *args, **kwargs)

create_observation_from_args(observation_space, function, args, kwargs) staticmethod

Helper method to create an observation from the arguments of a function.

Source code in mbodied/agents/agent.py 
185
186
187
188
189
190
191
192
193
194
195
196
197
198
@staticmethod
def create_observation_from_args(observation_space, function, args, kwargs) -> dict:
    """Helper method to create an observation from the arguments of a function."""
    param_names = list(signature(function).parameters.keys())

    # Create the observation from the arguments
    params = {**kwargs}
    for arg, val in zip(param_names, args, strict=False):
        params[arg] = val
    if observation_space is not None:
        observation = observation_space.sample()
        return {k: v for k, v in params.items() if k in observation}

    return {k: v for k, v in params.items() if v is not None and k not in ["self", "kwargs"]}

handle_default(model_src, model_kwargs) staticmethod

Default to gradio then httpx backend if the model source is not recognized.

Parameters:

Name Type Description Default
model_src str

The model source to use.

 required
model_kwargs dict

The additional arguments to pass to the model.

 required
Source code in mbodied/agents/agent.py 
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
@staticmethod
def handle_default(model_src: str, model_kwargs: dict) -> None:
    """Default to gradio then httpx backend if the model source is not recognized.

    Args:
        model_src: The model source to use.
        model_kwargs: The additional arguments to pass to the model.
    """
    try:
        return GradioBackend(endpoint=model_src, **model_kwargs)
    except Exception as e:
        logging.error(
            f"Failed to initialize Gradio backend: {e}. Defaulting to Httpx backend. Ensure that the source is a valid http endpoint.",
        )
        try:
            return HttpxBackend(endpoint=model_src, **model_kwargs)
        except Exception as e:
            logging.error(f"Failed to initialize Httpx backend: {e}.")
            raise ValueError(
                f"Failed to initialize backend for model source: {model_src}. Pleases select one of {Agent.ACTOR_MAP.keys()} or valid huggingface space or http endpoint.",
            )

init_backend(model_src, model_kwargs, api_key) staticmethod

Initialize the backend based on the model source.

Parameters:

Name Type Description Default
model_src str

The model source to use.

 required
model_kwargs dict

The additional arguments to pass to the model.

 required
api_key str

The API key to use for the remote actor.

 required

Returns:

Name Type Description
type type

The backend class to use.
Source code in mbodied/agents/agent.py 
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
@staticmethod
def init_backend(model_src: str, model_kwargs: dict, api_key: str) -> type:
    """Initialize the backend based on the model source.

    Args:
        model_src: The model source to use.
        model_kwargs: The additional arguments to pass to the model.
        api_key: The API key to use for the remote actor.

    Returns:
        type: The backend class to use.
    """
    if model_src in Agent.ACTOR_MAP:
        if model_src == "gradio":
            # Gradio doesn't take api_key.
            return Agent.ACTOR_MAP[model_src](**model_kwargs)
        else:
            return Agent.ACTOR_MAP[model_src](**model_kwargs, api_key=api_key)
    return Agent.handle_default(model_src, model_kwargs)

load_model(model)

Load a model from a file or path. Required if the model is a weights path.

Parameters:

Name Type Description Default
model str

The path to the model file.

 required
Source code in mbodied/agents/agent.py 
129
130
131
132
133
134
135
def load_model(self, model: str) -> None:
    """Load a model from a file or path. Required if the model is a weights path.

    Args:
        model: The path to the model file.
    """
    pass