Cross-Entropy Method (CEM)



Algorithm


Algorithm implementation

Main notation/symbols:
- policy function approximator (\(\pi_\theta\))
- states (\(s\)), actions (\(a\)), rewards (\(r\)), next states (\(s'\)), terminated (\(d_{_{end}}\)), truncated (\(d_{_{timeout}}\))
- loss (\(L\))

Decision making


act(...)
\(a \leftarrow \pi_\theta(s)\)

Learning algorithm


_update(...)
# sample all memory
\(s, a, r \leftarrow\) states, actions, rewards
# compute discounted return threshold
\([G] \leftarrow \sum_{t=0}^{E-1}\) discount_factor\(^{t} \, r_t\) for each episode
\(G_{_{bound}} \leftarrow q_{th_{quantile}}([G])\) at the given percentile
# get elite states and actions
\(s_{_{elite}} \leftarrow s[G \geq G_{_{bound}}]\)
\(a_{_{elite}} \leftarrow a[G \geq G_{_{bound}}]\)
# compute scores for the elite states
\(scores \leftarrow \theta(s_{_{elite}})\)
# compute policy loss
\(L_{\pi_\theta} \leftarrow -\sum_{i=1}^{N} a_{_{elite}} \log(scores)\)
# optimization step
reset \(\text{optimizer}_\theta\)
\(\nabla_{\theta} L_{\pi_\theta}\)
step \(\text{optimizer}_\theta\)
# update learning rate
IF there is a learning_rate_scheduler THEN
step \(\text{scheduler}_\theta (\text{optimizer}_\theta)\)

Usage

# import the agent and its default configuration
from skrl.agents.torch.cem import CEM, CEM_CFG

# instantiate the agent's models
models = {}
models["policy"] = ...

# adjust some configuration if necessary
cfg_agent = CEM_CFG()
cfg_agent.KEY = ...

# instantiate the agent
# (assuming a defined environment <env> and memory <memory>)
agent = CEM(
    models=models,
    memory=memory,  # only required during training
    cfg=cfg_agent,
    observation_space=env.observation_space,
    state_space=env.state_space,
    action_space=env.action_space,
    device=env.device,
)

Configuration and hyperparameters

Dataclass

    pytorch    

    jax    

    warp    

CEM_CFG

CEM_CFG

CEM_CFG


Spaces

The implementation supports the following Gymnasium spaces:

Gymnasium spaces

Observation

Action

Discrete

\(\square\)

\(\blacksquare\)

MultiDiscrete

\(\square\)

\(\blacksquare\)

Box

\(\blacksquare\)

\(\square\)

Dict

\(\blacksquare\)

\(\square\)


Models

The implementation uses 1 discrete function approximator. This function approximator (model) must be collected in a dictionary and passed to the constructor of the class under the argument models.

Notation

Concept

Key

Input shape

Output shape

Type

\(\pi(s)\)

Policy

"policy"

observation

action

Categorical /
Multi-Categorical


Features

Support for advanced features is described in the following table:

Feature

Support and remarks

    pytorch    

    jax    

    warp    

RNN support

-

\(\square\)

\(\square\)

\(\square\)

Mixed precision

Automatic mixed precision

\(\blacksquare\)

\(\square\)

\(\square\)

Distributed

-

\(\square\)

\(\square\)

\(\square\)


API


PyTorch

CEM_CFG

Configuration for the CEM agent.

CEM

Cross-Entropy Method (CEM).

class skrl.agents.torch.cem.CEM_CFG(*, experiment: ExperimentCfg = <factory>, rollouts: int = 16, percentile: float = 0.7, discount_factor: float = 0.99, learning_rate: float = 0.01, learning_rate_scheduler: type | None = None, learning_rate_scheduler_kwargs: dict = <factory>, observation_preprocessor: type | None = None, observation_preprocessor_kwargs: dict = <factory>, state_preprocessor: type | None = None, state_preprocessor_kwargs: dict = <factory>, random_timesteps: int = 0, learning_starts: int = 0, rewards_shaper: Callable | None = None, mixed_precision: bool = False)[source]

Bases: AgentCfg

Configuration for the CEM agent.

Methods:

expand()

Expand the configuration.

validate()

Validate the configuration.

Attributes:

discount_factor

Parameter that balances the importance of future rewards (close to 1.0) versus immediate rewards (close to 0.0).

experiment

Experiment settings.

learning_rate

Learning rate for the policy network.

learning_rate_scheduler

Learning rate scheduler class for the policy network.

learning_rate_scheduler_kwargs

Keyword arguments for the learning rate scheduler's constructor.

learning_starts

Number of steps to perform before calling the algorithm update function.

mixed_precision

Whether to enable automatic mixed precision for higher performance.

observation_preprocessor

Preprocessor class to process the environment's observations.

observation_preprocessor_kwargs

Keyword arguments for the observation preprocessor's constructor.

percentile

Percentile to compute the reward bound.

random_timesteps

Number of random exploration (sampling random actions) steps to perform before sampling actions from the policy.

rewards_shaper

Rewards shaping function.

rollouts

Number of collection steps to perform between updates.

state_preprocessor

Preprocessor class to process the environment's states.

state_preprocessor_kwargs

Keyword arguments for the state preprocessor's constructor.

expand() None[source]

Expand the configuration.

validate() bool[source]

Validate the configuration.

discount_factor: float = 0.99

Parameter that balances the importance of future rewards (close to 1.0) versus immediate rewards (close to 0.0).

Range: [0.0, 1.0].

experiment: ExperimentCfg

Experiment settings.

learning_rate: float = 0.01

Learning rate for the policy network.

learning_rate_scheduler: type | None = None

Learning rate scheduler class for the policy network.

See Learning rate schedulers for more details.

learning_rate_scheduler_kwargs: dict

Keyword arguments for the learning rate scheduler’s constructor.

See Learning rate schedulers for more details.

Warning

The optimizer argument is automatically passed to the learning rate scheduler’s constructor. Therefore, it must not be provided in the keyword arguments.

learning_starts: int = 0

Number of steps to perform before calling the algorithm update function.

mixed_precision: bool = False

Whether to enable automatic mixed precision for higher performance.

observation_preprocessor: type | None = None

Preprocessor class to process the environment’s observations.

See Preprocessors for more details.

observation_preprocessor_kwargs: dict

Keyword arguments for the observation preprocessor’s constructor.

See Preprocessors for more details.

percentile: float = 0.7

Percentile to compute the reward bound.

Range: [0.0, 1.0].

random_timesteps: int = 0

Number of random exploration (sampling random actions) steps to perform before sampling actions from the policy.

rewards_shaper: Callable | None = None

Rewards shaping function.

rollouts: int = 16

Number of collection steps to perform between updates.

state_preprocessor: type | None = None

Preprocessor class to process the environment’s states.

See Preprocessors for more details.

state_preprocessor_kwargs: dict

Keyword arguments for the state preprocessor’s constructor.

See Preprocessors for more details.

class skrl.agents.torch.cem.CEM(*, models: dict[str, Model], memory: Memory | None = None, observation_space: gymnasium.Space | None = None, state_space: gymnasium.Space | None = None, action_space: gymnasium.Space | None = None, device: str | torch.device | None = None, cfg: CEM_CFG | dict = {})[source]

Bases: Agent

Cross-Entropy Method (CEM).

https://ieeexplore.ieee.org/abstract/document/6796865/

Parameters:
  • models – Agent’s models.

  • memory – Memory to storage agent’s data and environment transitions.

  • observation_space – Observation space.

  • state_space – State space.

  • action_space – Action space.

  • device – Data allocation and computation device. If not specified, the default device will be used.

  • cfg – Agent’s configuration.

Raises:

KeyError – If a configuration key is missing.

Methods:

act(observations, states, *, timestep, timesteps)

Process the environment's observations/states to make a decision (actions) using the main policy.

enable_models_training_mode([enabled])

Set the training mode of all the agent's models: enabled (training) or disabled (evaluation).

enable_training_mode([enabled, apply_to_models])

Set the training mode of the agent: enabled (training) or disabled (evaluation).

init(*[, trainer_cfg])

Initialize the agent.

load(path)

Load the agent from the specified path.

post_interaction(*, timestep, timesteps)

Method called after the interaction with the environment.

pre_interaction(*, timestep, timesteps)

Method called before the interaction with the environment.

record_transition(*, observations, states, ...)

Record an environment transition in memory.

save(path)

Save the agent to the specified path.

track_data(tag, value)

Track data to TensorBoard.

update(*, timestep, timesteps)

Algorithm's main update step.

write_checkpoint(*, timestep, timesteps)

Write checkpoint (modules) to persistent storage.

write_tracking_data(*, timestep, timesteps)

Write tracking data to TensorBoard.

act(observations: torch.Tensor, states: torch.Tensor | None, *, timestep: int, timesteps: int) tuple[torch.Tensor, dict[str, Any]][source]

Process the environment’s observations/states to make a decision (actions) using the main policy.

Parameters:
  • observations – Environment observations.

  • states – Environment states.

  • timestep – Current timestep.

  • timesteps – Number of timesteps.

Returns:

Agent output. The first component is the expected action/value returned by the agent. The second component is a dictionary containing extra output values according to the model.

enable_models_training_mode(enabled: bool = True) None[source]

Set the training mode of all the agent’s models: enabled (training) or disabled (evaluation).

Parameters:

enabled – True to enable the training mode, False to enable the evaluation mode.

enable_training_mode(enabled: bool = True, *, apply_to_models: bool = False) None[source]

Set the training mode of the agent: enabled (training) or disabled (evaluation).

The training mode can be queried by the training property.

Parameters:
  • enabled – True to enable the training mode, False to enable the evaluation mode.

  • apply_to_models – Whether to apply the training mode to all the agent’s models.

init(*, trainer_cfg: dict[str, Any] | None = None) None[source]

Initialize the agent.

Parameters:

trainer_cfg – Trainer configuration.

load(path: str) None[source]

Load the agent from the specified path.

Note

The final storage device is determined by the constructor of the agent.

Parameters:

path – Path to load the agent from.

post_interaction(*, timestep: int, timesteps: int) None[source]

Method called after the interaction with the environment.

Parameters:
  • timestep – Current timestep.

  • timesteps – Number of timesteps.

pre_interaction(*, timestep: int, timesteps: int) None[source]

Method called before the interaction with the environment.

Parameters:
  • timestep – Current timestep.

  • timesteps – Number of timesteps.

record_transition(*, observations: torch.Tensor, states: torch.Tensor, actions: torch.Tensor, rewards: torch.Tensor, next_observations: torch.Tensor, next_states: torch.Tensor, terminated: torch.Tensor, truncated: torch.Tensor, infos: Any, timestep: int, timesteps: int) None[source]

Record an environment transition in memory.

Parameters:
  • observations – Environment observations.

  • states – Environment states.

  • actions – Actions taken by the agent.

  • rewards – Instant rewards achieved by the current actions.

  • next_observations – Next environment observations.

  • next_states – Next environment states.

  • terminated – Signals that indicate episodes have terminated.

  • truncated – Signals that indicate episodes have been truncated.

  • infos – Additional information about the environment.

  • timestep – Current timestep.

  • timesteps – Number of timesteps.

save(path: str) None[source]

Save the agent to the specified path.

Parameters:

path – Path to save the agent to.

track_data(tag: str, value: float) None[source]

Track data to TensorBoard.

Note

Currently only scalar data is supported.

Parameters:
  • tag – Data identifier (e.g. ‘Loss/Policy loss’).

  • value – Value to track.

update(*, timestep: int, timesteps: int) None[source]

Algorithm’s main update step.

Parameters:
  • timestep – Current timestep.

  • timesteps – Number of timesteps.

write_checkpoint(*, timestep: int, timesteps: int) None[source]

Write checkpoint (modules) to persistent storage.

Note

The checkpoints are stored in the subdirectory checkpoints within the experiment directory. The checkpoint name is the timestep argument value (if it is not None), or the current system date-time otherwise.

Parameters:
  • timestep – Current timestep.

  • timesteps – Number of timesteps.

write_tracking_data(*, timestep: int, timesteps: int) None[source]

Write tracking data to TensorBoard.

Parameters:
  • timestep – Current timestep.

  • timesteps – Number of timesteps.


JAX

CEM_CFG

Configuration for the CEM agent.

CEM

Cross-Entropy Method (CEM).

class skrl.agents.jax.cem.CEM_CFG(*, experiment: ExperimentCfg = <factory>, rollouts: int = 16, percentile: float = 0.7, discount_factor: float = 0.99, learning_rate: float = 0.01, learning_rate_scheduler: type | None = None, learning_rate_scheduler_kwargs: dict = <factory>, observation_preprocessor: type | None = None, observation_preprocessor_kwargs: dict = <factory>, state_preprocessor: type | None = None, state_preprocessor_kwargs: dict = <factory>, random_timesteps: int = 0, learning_starts: int = 0, rewards_shaper: Callable | None = None)[source]

Bases: AgentCfg

Configuration for the CEM agent.

Methods:

expand()

Expand the configuration.

validate()

Validate the configuration.

Attributes:

discount_factor

Parameter that balances the importance of future rewards (close to 1.0) versus immediate rewards (close to 0.0).

experiment

Experiment settings.

learning_rate

Learning rate for the policy network.

learning_rate_scheduler

Learning rate scheduler class for the policy network.

learning_rate_scheduler_kwargs

Keyword arguments for the learning rate scheduler's constructor.

learning_starts

Number of steps to perform before calling the algorithm update function.

observation_preprocessor

Preprocessor class to process the environment's observations.

observation_preprocessor_kwargs

Keyword arguments for the observation preprocessor's constructor.

percentile

Percentile to compute the reward bound.

random_timesteps

Number of random exploration (sampling random actions) steps to perform before sampling actions from the policy.

rewards_shaper

Rewards shaping function.

rollouts

Number of collection steps to perform between updates.

state_preprocessor

Preprocessor class to process the environment's states.

state_preprocessor_kwargs

Keyword arguments for the state preprocessor's constructor.

expand() None[source]

Expand the configuration.

validate() bool[source]

Validate the configuration.

discount_factor: float = 0.99

Parameter that balances the importance of future rewards (close to 1.0) versus immediate rewards (close to 0.0).

Range: [0.0, 1.0].

experiment: ExperimentCfg

Experiment settings.

learning_rate: float = 0.01

Learning rate for the policy network.

learning_rate_scheduler: type | None = None

Learning rate scheduler class for the policy network.

See Learning rate schedulers for more details.

learning_rate_scheduler_kwargs: dict

Keyword arguments for the learning rate scheduler’s constructor.

See Learning rate schedulers for more details.

Warning

The optimizer argument is automatically passed to the learning rate scheduler’s constructor. Therefore, it must not be provided in the keyword arguments.

learning_starts: int = 0

Number of steps to perform before calling the algorithm update function.

observation_preprocessor: type | None = None

Preprocessor class to process the environment’s observations.

See Preprocessors for more details.

observation_preprocessor_kwargs: dict

Keyword arguments for the observation preprocessor’s constructor.

See Preprocessors for more details.

percentile: float = 0.7

Percentile to compute the reward bound.

Range: [0.0, 1.0].

random_timesteps: int = 0

Number of random exploration (sampling random actions) steps to perform before sampling actions from the policy.

rewards_shaper: Callable | None = None

Rewards shaping function.

rollouts: int = 16

Number of collection steps to perform between updates.

state_preprocessor: type | None = None

Preprocessor class to process the environment’s states.

See Preprocessors for more details.

state_preprocessor_kwargs: dict

Keyword arguments for the state preprocessor’s constructor.

See Preprocessors for more details.

class skrl.agents.jax.cem.CEM(*, models: dict[str, Model], memory: Memory | None = None, observation_space: gymnasium.Space | None = None, state_space: gymnasium.Space | None = None, action_space: gymnasium.Space | None = None, device: str | jax.Device | None = None, cfg: CEM_CFG | dict = {})[source]

Bases: Agent

Cross-Entropy Method (CEM).

https://ieeexplore.ieee.org/abstract/document/6796865/

Parameters:
  • models – Agent’s models.

  • memory – Memory to storage agent’s data and environment transitions.

  • observation_space – Observation space.

  • state_space – State space.

  • action_space – Action space.

  • device – Data allocation and computation device. If not specified, the default device will be used.

  • cfg – Agent’s configuration.

Raises:

KeyError – If a configuration key is missing.

Methods:

act(observations, states, *, timestep, timesteps)

Process the environment's observations/states to make a decision (actions) using the main policy.

enable_models_training_mode([enabled])

Set the training mode of all the agent's models: enabled (training) or disabled (evaluation).

enable_training_mode([enabled, apply_to_models])

Set the training mode of the agent: enabled (training) or disabled (evaluation).

init(*[, trainer_cfg])

Initialize the agent.

load(path)

Load the agent from the specified path.

post_interaction(*, timestep, timesteps)

Method called after the interaction with the environment.

pre_interaction(*, timestep, timesteps)

Method called before the interaction with the environment.

record_transition(*, observations, states, ...)

Record an environment transition in memory.

save(path)

Save the agent to the specified path.

track_data(tag, value)

Track data to TensorBoard.

update(*, timestep, timesteps)

Algorithm's main update step.

write_checkpoint(*, timestep, timesteps)

Write checkpoint (modules) to persistent storage.

write_tracking_data(*, timestep, timesteps)

Write tracking data to TensorBoard.

act(observations: jax.Array, states: jax.Array | None, *, timestep: int, timesteps: int) tuple[jax.Array, dict[str, Any]][source]

Process the environment’s observations/states to make a decision (actions) using the main policy.

Parameters:
  • observations – Environment observations.

  • states – Environment states.

  • timestep – Current timestep.

  • timesteps – Number of timesteps.

Returns:

Agent output. The first component is the expected action/value returned by the agent. The second component is a dictionary containing extra output values according to the model.

enable_models_training_mode(enabled: bool = True) None[source]

Set the training mode of all the agent’s models: enabled (training) or disabled (evaluation).

Parameters:

enabled – True to enable the training mode, False to enable the evaluation mode.

enable_training_mode(enabled: bool = True, *, apply_to_models: bool = False) None[source]

Set the training mode of the agent: enabled (training) or disabled (evaluation).

The training mode can be queried by the training property.

Parameters:
  • enabled – True to enable the training mode, False to enable the evaluation mode.

  • apply_to_models – Whether to apply the training mode to all the agent’s models.

init(*, trainer_cfg: dict[str, Any] | None = None) None[source]

Initialize the agent.

Parameters:

trainer_cfg – Trainer configuration.

load(path: str) None[source]

Load the agent from the specified path.

Parameters:

path – Path to load the agent from.

post_interaction(*, timestep: int, timesteps: int) None[source]

Method called after the interaction with the environment.

Parameters:
  • timestep – Current timestep.

  • timesteps – Number of timesteps.

pre_interaction(*, timestep: int, timesteps: int) None[source]

Method called before the interaction with the environment.

Parameters:
  • timestep – Current timestep.

  • timesteps – Number of timesteps.

record_transition(*, observations: jax.Array, states: jax.Array, actions: jax.Array, rewards: jax.Array, next_observations: jax.Array, next_states: jax.Array, terminated: jax.Array, truncated: jax.Array, infos: Any, timestep: int, timesteps: int) None[source]

Record an environment transition in memory.

Parameters:
  • observations – Environment observations.

  • states – Environment states.

  • actions – Actions taken by the agent.

  • rewards – Instant rewards achieved by the current actions.

  • next_observations – Next environment observations.

  • next_states – Next environment states.

  • terminated – Signals that indicate episodes have terminated.

  • truncated – Signals that indicate episodes have been truncated.

  • infos – Additional information about the environment.

  • timestep – Current timestep.

  • timesteps – Number of timesteps.

save(path: str) None[source]

Save the agent to the specified path.

Parameters:

path – Path to save the agent to.

track_data(tag: str, value: float) None[source]

Track data to TensorBoard.

Note

Currently only scalar data is supported.

Parameters:
  • tag – Data identifier (e.g. ‘Loss/Policy loss’).

  • value – Value to track.

update(*, timestep: int, timesteps: int) None[source]

Algorithm’s main update step.

Parameters:
  • timestep – Current timestep.

  • timesteps – Number of timesteps.

write_checkpoint(*, timestep: int, timesteps: int) None[source]

Write checkpoint (modules) to persistent storage.

Note

The checkpoints are stored in the subdirectory checkpoints within the experiment directory. The checkpoint name is the timestep argument value (if it is not None), or the current system date-time otherwise.

Parameters:
  • timestep – Current timestep.

  • timesteps – Number of timesteps.

write_tracking_data(*, timestep: int, timesteps: int) None[source]

Write tracking data to TensorBoard.

Parameters:
  • timestep – Current timestep.

  • timesteps – Number of timesteps.