Planning module

class mbrl.planning.Agent

Bases: object

Abstract class for all agents.

abstract act(obs: numpy.ndarray, **_kwargs) → numpy.ndarray

Issues an action given an observation.

Parameters

obs (np.ndarray) – the observation for which the action is needed.

Returns

the action.

Return type

(np.ndarray)

plan(obs: numpy.ndarray, **_kwargs) → numpy.ndarray

Issues a sequence of actions given an observation.

Unless overridden by a child class, this will be equivalent to act().

Parameters

obs (np.ndarray) – the observation for which the sequence is needed.

Returns

a sequence of actions.

Return type

(np.ndarray)

reset()

Resets any internal state of the agent.

class mbrl.planning.CEMOptimizer(num_iterations: int, elite_ratio: float, population_size: int, lower_bound: Sequence[float], upper_bound: Sequence[float], alpha: float, device: torch.device, return_mean_elites: bool = False)

Bases: mbrl.planning.trajectory_opt.Optimizer

Implements the Cross-Entropy Method optimization algorithm.

A good description of CEM [1] can be found at https://arxiv.org/pdf/2008.06389.pdf. This code implements the version described in Section 2.1, labeled CEM_PETS (but note that the shift-initialization between planning time steps is handled outside of this class by TrajectoryOptimizer).

This implementation also returns the best solution found as opposed to the mean of the last generation.

Parameters

  • num_iterations (int) – the number of iterations (generations) to perform.

  • elite_ratio (float) – the proportion of the population that will be kept as elite (rounds up).

  • population_size (int) – the size of the population.

  • lower_bound (sequence of floats) – the lower bound for the optimization variables.

  • upper_bound (sequence of floats) – the upper bound for the optimization variables.

  • alpha (float) – momentum term.

  • device (torch.device) – device where computations will be performed.

  • return_mean_elites (bool) – if True returns the mean of the elites of the last iteration. Otherwise, it returns the max solution found over all iterations.

[1] R. Rubinstein and W. Davidson. “The cross-entropy method for combinatorial and continuous optimization”. Methodology and Computing in Applied Probability, 1999.

optimize(obj_fun: Callable[[torch.Tensor], torch.Tensor], x0: Optional[torch.Tensor] = None, callback: Optional[Callable[[torch.Tensor, torch.Tensor, int], None]] = None, **kwargs) → torch.Tensor

Runs the optimization using CEM.

Parameters

  • obj_fun (callable(tensor) -> tensor) – objective function to maximize.

  • x0 (tensor, optional) – initial mean for the population. Must be consistent with lower/upper bounds.

  • callback (callable(tensor, tensor, int) -> any, optional) – if given, this function will be called after every iteration, passing it as input the full population tensor, its corresponding objective function values, and the index of the current iteration. This can be used for logging and plotting purposes.

Returns

the best solution found.

Return type

(torch.Tensor)

class mbrl.planning.RandomAgent(env: gym.core.Env)

Bases: mbrl.planning.core.Agent

An agent that samples action from the environments action space.

Parameters

env (gym.Env) – the environment on which the agent will act.

act(*_args, **_kwargs) → numpy.ndarray

Issues an action given an observation.

Returns

an action sampled from the environment’s action space.

Return type

(np.ndarray)

class mbrl.planning.TrajectoryOptimizer(optimizer_cfg: omegaconf.dictconfig.DictConfig, action_lb: numpy.ndarray, action_ub: numpy.ndarray, planning_horizon: int, replan_freq: int = 1, keep_last_solution: bool = True)

Bases: object

Class for using generic optimizers on trajectory optimization problems.

This is a convenience class that sets up optimization problem for trajectories, given only action bounds and the length of the horizon. Using this class, the concern of handling appropriate tensor shapes for the optimization problem is hidden from the users, which only need to provide a function that is capable of evaluating trajectories of actions. It also takes care of shifting previous solution for the next optimization call, if the user desires.

The optimization variables for the problem will have shape H x A, where H and A represent planning horizon and action dimension, respectively. The initial solution for the optimizer will be computed as (action_ub - action_lb) / 2, for each time step.

Parameters

  • optimizer_cfg (omegaconf.DictConfig) – the configuration of the optimizer to use.

  • action_lb (np.ndarray) – the lower bound for actions.

  • action_ub (np.ndarray) – the upper bound for actions.

  • planning_horizon (int) – the length of the trajectories that will be optimized.

  • replan_freq (int) – the frequency of re-planning. This is used for shifting the previous

  • for the next time step (solution) –

  • keep_last_solution == True. Defaults to 1. (when) –

  • keep_last_solution (bool) – if True, the last solution found by a call to optimize() is kept as the initial solution for the next step. This solution is shifted replan_freq time steps, and the new entries are filled using th3 initial solution. Defaults to True.

optimize(trajectory_eval_fn: Callable[[torch.Tensor], torch.Tensor], callback: Optional[Callable] = None) → numpy.ndarray

Runs the trajectory optimization.

Parameters

  • trajectory_eval_fn (callable(tensor) -> tensor) – A function that receives a batch of action sequences and returns a batch of objective function values (e.g., accumulated reward for each sequence). The shape of the action sequence tensor will be B x H x A, where B, H, and A represent batch size, planning horizon, and action dimension, respectively.

  • callback (callable, optional) – a callback function to pass to the optimizer.

Returns

the best action sequence.

Return type

(tuple of np.ndarray and float)

reset()

Resets the previous solution cache to the initial solution.

class mbrl.planning.TrajectoryOptimizerAgent(optimizer_cfg: omegaconf.dictconfig.DictConfig, action_lb: Sequence[float], action_ub: Sequence[float], planning_horizon: int = 1, replan_freq: int = 1, verbose: bool = False)

Bases: mbrl.planning.core.Agent

Agent that performs trajectory optimization on a given objective function for each action.

This class uses an internal TrajectoryOptimizer object to generate sequence of actions, given a user-defined trajectory optimization function.

Parameters

  • optimizer_cfg (omegaconf.DictConfig) – the configuration of the base optimizer to pass to the trajectory optimizer.

  • action_lb (sequence of floats) – the lower bound of the action space.

  • action_ub (sequence of floats) – the upper bound of the action space.

  • planning_horizon (int) – the length of action sequences to evaluate. Defaults to 1.

  • replan_freq (int) – the frequency of re-planning. The agent will keep a cache of the generated sequences an use it for replan_freq number of act() calls. Defaults to 1.

  • verbose (bool) – if True, prints the planning time on the console.

Note

After constructing an agent of this type, the user must call set_trajectory_eval_fn(). This is not passed to the constructor so that the agent can be automatically instantiated with Hydra (which in turn makes it easy to replace this agent with an agent of another type via config-only changes).

act(obs: numpy.ndarray, **_kwargs) → numpy.ndarray

Issues an action given an observation.

This method optimizes a full sequence of length self.planning_horizon and returns the first action in the sequence. If self.replan_freq > 1, future calls will use subsequent actions in the sequence, for self.replan_freq number of steps. After that, the method will plan again, and repeat this process.

Parameters

obs (np.ndarray) – the observation for which the action is needed.

Returns

the action.

Return type

(np.ndarray)

plan(obs: numpy.ndarray, **_kwargs) → numpy.ndarray

Issues a sequence of actions given an observation.

Returns s sequence of length self.planning_horizon.

Parameters

obs (np.ndarray) – the observation for which the sequence is needed.

Returns

a sequence of actions.

Return type

(np.ndarray)

reset(planning_horizon: Optional[int] = None)

Resets the underlying trajectory optimizer.

set_trajectory_eval_fn(trajectory_eval_fn: Callable[[Union[torch.Tensor, numpy.ndarray], torch.Tensor], torch.Tensor])

Sets the trajectory evaluation function.

Parameters

trajectory_eval_fn (callable) – a trajectory evaluation function, as described in TrajectoryOptimizer.

mbrl.planning.complete_agent_cfg(env: Union[gym.core.Env, mbrl.models.model_env.ModelEnv], agent_cfg: omegaconf.dictconfig.DictConfig)

Completes an agent’s configuration given information from the environment.

The goal of this function is to completed information about state and action shapes and ranges, without requiring the user to manually enter this into the Omegaconf configuration object.

It will check for and complete any of the following keys:

  • “obs_dim”: set to env.observation_space.shape

  • “action_dim”: set to env.action_space.shape

  • “action_range”: set to max(env.action_space.high) - min(env.action_space.low)

  • “action_lb”: set to env.action_space.low

  • “action_ub”: set to env.action_space.high

Note

If the user provides any of these values in the Omegaconf configuration object, these will not be overridden by this function.

mbrl.planning.create_trajectory_optim_agent_for_model(model_env: mbrl.models.model_env.ModelEnv, agent_cfg: omegaconf.dictconfig.DictConfig, num_particles: int = 1) → mbrl.planning.trajectory_opt.TrajectoryOptimizerAgent

Utility function for creating a trajectory optimizer agent for a model environment.

This is a convenience function for creating a TrajectoryOptimizerAgent, using mbrl.models.ModelEnv.evaluate_action_sequences() as its objective function.

Parameters

  • model_env (mbrl.models.ModelEnv) – the model environment.

  • agent_cfg (omegaconf.DictConfig) – the agent’s configuration.

  • num_particles (int) – the number of particles for taking averages of action sequences’ total rewards.

Returns

the agent.

Return type

(TrajectoryOptimizerAgent)

mbrl.planning.load_agent(agent_path: Union[str, pathlib.Path], env: gym.core.Env) → mbrl.planning.core.Agent

Loads an agent from a Hydra config file at the given path.

For agent of type “pytorch_sac.agent.sac.SACAgent”, the directory must contain the following files:

  • “.hydra/config.yaml”: the Hydra configuration for the agent.

  • “critic.pth”: the saved checkpoint for the critic.

  • “actor.pth”: the saved checkpoint for the actor.

Parameters

  • agent_path (str or pathlib.Path) – a path to the directory where the agent is saved.

  • env (gym.Env) – the environment on which the agent will operate (only used to complete the agent’s configuration).

Returns

the new agent.

Return type

(Agent)