salina.agents.brax

salina.agents.brax.BraxAgent (TAgent)

An agent based on a brax environment, with autoreset

The agent reads action at t-1 and outputs env/env_obs,env/reward,env/initial_state,env/done,env/timestep,env/cumulated_reward

__init__(self, n_envs, env_name, input='action', output='env/', **kwargs) special

Initialize the agent

Parameters:

Name	Type	Description	Default
n_envs	[int]	number of envs (batch dimension)	required
env_name	[str]	[the BRAX environment name	required
input	str	Defaults to "action".	'action'
output	str	Defaults to "env/".	'env/'

Source code in salina/agents/brax.py

def __init__(self, n_envs, env_name, input="action", output="env/", **kwargs):
    """ Initialize the agent

    Args:
        n_envs ([int]): number of envs (batch dimension)
        env_name ([str]): [the BRAX environment name
        input (str, optional):  Defaults to "action".
        output (str, optional): Defaults to "env/".
    """
    super().__init__()
    self.args = kwargs
    self.brax_env_name = env_name
    self.gym_env = None
    self._seed = None
    self.n_envs = n_envs
    self.output = output
    self.input = input
    self.brax_device = None
    self.ghost_params = torch.nn.Parameter(torch.randn(()))

salina.agents.brax.AutoResetBraxAgent (BraxAgent)

The same than BraxAgent

salina.agents.brax.NoAutoResetBraxAgent (Agents)

A BraxAgent without auto-reset