imitation.testing.reward_nets

Utility functions for testing reward nets.

Functions

make_ensemble(obs_space, action_space[, ...])

Create a simple reward ensemble.

Classes

MockRewardNet(observation_space, action_space)

A mock reward net for testing.

class imitation.testing.reward_nets.MockRewardNet(observation_space, action_space, value=0.0)

Bases: RewardNet

A mock reward net for testing.

__init__(observation_space, action_space, value=0.0)

Create mock reward.

Parameters

• observation_space (Space) – observation space of the env

• action_space (Space) – action space of the env

• value (float) – The reward to always return. Defaults to 0.0.

forward(state, action, next_state, done)

Compute rewards for a batch of transitions and keep gradients.

Return type

Tensor

training: bool

imitation.testing.reward_nets.make_ensemble(obs_space, action_space, num_members=2, **kwargs)

Create a simple reward ensemble.