Skip to content

A PyTorch Library for Reinforcement Learning Research

License

Notifications You must be signed in to change notification settings

learnables/cherry

Repository files navigation


Build Status

Cherry is a reinforcement learning framework for researchers built on top of PyTorch.

Unlike other reinforcement learning implementations, cherry doesn't implement a single monolithic interface to existing algorithms. Instead, it provides you with low-level, common tools to write your own algorithms. Drawing from the UNIX philosophy, each tool strives to be as independent from the rest of the framework as possible. So if you don't like a specific tool, you don’t need to use it.

Features

Cherry extends PyTorch with only a handful of new core concepts.

To learn more about the tools and philosophy behind cherry, check out our Getting Started tutorial.

Overview and Examples

The following snippet showcases a few of the tools offered by cherry. Many more high-quality examples are available in the examples/ folder.

Defining a cherry.nn.Policy

class VisionPolicy(cherry.nn.Policy):  # inherits from torch.nn.Module

   def __init__(self, feature_extractor, actor):
      super(VisionGaussianPolicy, self).__init__()
      self.feature_extractor = feature_extractor
      self.actor = actor

   def forward(self, obs):
      mean = self.actor(self.feature_extractor(obs))
      std = 0.1 * torch.ones_like(mean)
      return cherry.distributions.TanhNormal(mean, std)  # policies always return a distribution

policy = VisionPolicy(MyResnetExtractor(), MyMLPActor())
action = policy.act(obs)  # sampled from policy's distribution
deterministic_action = policy.act(obs, deterministic=True)  # distribution's mode
action_distribution = policy(obs)  # work with the policy's distribution
# building the replay
replay = cherry.ExperienceReplay()
state = env.reset()
for t in range(1000):
   action = policy.act(state)
   next_state, reward, done, info = env.step(action)
   replay.append(state, action, reward, next_state, done)
   next_state = state

# manipulating the replay
replay = replay[-256:]  # indexes like a list
batch = replay.sample(32, contiguous=True)  # sample transitions into a replay
batch = batch.to('cuda') # move replay to device
for transition in reversed(batch): # iterate over a replay
   transition.reward *= 0.99

# get all states, actions, and rewards as PyTorch tensors.
reinforce_loss = - torch.sum(policy(batch.state()).log_prob(batch.action()) * batch.reward())

Designing algorithms with cherry.td, cherry.pg, and cherry.algorithms

# defining a new algorithm
@dataclasses.dataclass
class MyA2C:
   discount: float = 0.99
   
   def update(self, replay, policy, state_value, optimizer):
      # discount rewards
      values = state_value(replay.action())
      discounted_rewards = cherry.td.discount(
         self.discount, replay.reward(), replay.done(), bootstrap=values[-1].detach()
      )

      # Compute losses
      policy_loss = cherry.algorithms.A2C.policy_loss(
         log_probs=policy(replay.state()).log_prob(replay.action()),
         advantages=discounted_rewards - values.detach(),
      )
      value_loss = cherry.algorithms.A2C.state_value_loss(values, discounted_rewards)

      # Optimization step
      optimizer.zero_grad()
      (policy_loss + value_loss).backward()
      optimizer.step()
      return {'a2c/policy_loss': policy_loss, 'a2c/value_loss': value_loss}

# using MyA2C
my_a2c = MyA2C(discount=0.95)
my_policy = MyPolicy()
linear_value = cherry.models.LinearValue(128)
adam = torch.optim.Adam(policy.parameters())
for step in range(1000):
   replay = collect_experience(policy)
   my_a2c.update(replay, my_policy, linear_value, adam)

Install

pip install cherry-rl

Changelog

A human-readable changelog is available in the CHANGELOG.md file.

Documentation

Documentation and tutorials are available on cherry’s website: http://cherry-rl.net.

Contributing

Here are a couple of guidelines we strive to follow.

  • It's always a good idea to open an issue first, where we can discuss how to best proceed.
  • If you want to contribute a new example using cherry, it would preferably stand in a single file.
  • If you would like to contribute a new feature to the core library, we suggest to first implement an example showcasing your new functionality. Doing so is quite useful:
    • it allows for automatic testing,
    • it ensures that the functionality is correctly implemented,
    • it shows users how to use your functionality, and
    • it gives a concrete example when discussing the best way to merge your implementation.

Acknowledgements

Cherry draws inspiration from many reinforcement learning implementations, including

Why 'cherry' ?

Because it's the sweetest part of the cake.