【重磅Nature+代码开源】DeepMind提出AlphaDev, 用强化学习将排序算法提升70%

"""Pseudocode description of the AlphaDev algorithm."""

###########################
########## Content ########
# 1. Environment
# 2. Networks
#   2.1 Network helpers
#   2.2 Representation network
#   2.3 Prediction network (correctness and latency values and policy)
# 3. Helpers
# 4. Part 1: Self-Play
# 5. Part 2: Training
###########################

import collections
import functools
import math
from typing import Any, Callable, Dict, NamedTuple, Optional, Sequence

import chex
import haiku as hk
import jax
import jax.lax
import jax.numpy as jnp
import ml_collections
import numpy
import optax


############################
###### 1. Environment ######


class TaskSpec(NamedTuple):
  max_program_size: int
  num_inputs: int
  num_funcs: int
  num_locations: int
  num_actions: int
  correct_reward: float
  correctness_reward_weight: float
  latency_reward_weight: float
  latency_quantile: float


class AssemblyGame(object):
  """The environment AlphaDev is interacting with."""

  class AssemblyInstruction(object):
    pass

  class AssemblySimulator(object):

    # pylint: disable-next=unused-argument
    def apply(self, instruction):
      return {}

    def measure_latency(self, program) -> float:
      pass

  def __init__(self, task_spec):
    self.task_spec = task_spec
    self.program = []
    self.simulator = self.AssemblySimulator(task_spec)
    self.previous_correct_items = 0

  def step(self, action):
    instruction = self.AssemblyInstruction(action)
    self.program.append(instruction)
    self.execution_state = self.simulator.apply(instruction)
    return self.observation(), self.correctness_reward()

  def observation(self):
    return {
        'program': self.program,
        'program_length': len(self.program),
        'memory': self.execution_state.memory,
        'registers': self.execution_state.registers,
    }

  def correctness_reward(self) -> float:
    """Computes a reward based on the correctness of the output."""
    make_expected_outputs = lambda: []
    expected_outputs = make_expected_outputs()
    state = self.execution_state

    # Weighted sum of correctly placed items
    correct_items = 0
    for output, expected in zip(state.memory, expected_outputs):
      correct_items += output.weight * sum(
          output[i] == expected[i] for i in range(len(output))
      )
    reward = self.task_spec.correctness_reward_weight * (
        correct_items - self.previous_correct_items
    )
    self.previous_correct_items = correct_items

    # Bonus for fully correct programs
    all_correct = all(
        output == expected
        for output, expected in zip(state.memory, expected_outputs)
    )
    reward += self.task_spec.correct_reward * all_correct

    return reward

  def latency_reward(self) -> float:
    latency_samples = [
        self.simulator.measure_latency(self.program)
        for _ in range(self.task_spec.num_latency_simulation)
    ]
    return (
        numpy.quantile(latency_samples, self.task_spec.latency_quantile)
        * self.task_spec.latency_reward_weight
    )

  def clone(self):
    pass

神经网络部分

class Action(object):
  """Action representation."""

  def __init__(self, index: int):
    self.index = index

  def __hash__(self):
    return self.index

  def __eq__(self, other):
    return self.index == other.index

  def __gt__(self, other):
    return self.index > other.index


class NetworkOutput(NamedTuple):
  value: float
  correctness_value_logits: jnp.ndarray
  latency_value_logits: jnp.ndarray
  policy_logits: Dict[Action, float]


class Network(object):
  """Wrapper around Representation and Prediction networks."""

  def __init__(self, hparams: ml_collections.ConfigDict, task_spec: TaskSpec):
    self.representation = hk.transform(RepresentationNet(
        hparams, task_spec, hparams.embedding_dim
    ))
    self.prediction = hk.transform(PredictionNet(
        task_spec=task_spec,
        value_max=hparams.value.max,
        value_num_bins=hparams.value.num_bins,
        embedding_dim=hparams.embedding_dim,
    ))
    rep_key, pred_key = jax.random.PRNGKey(42).split()
    self.params = {
        'representation': self.representation.init(rep_key),
        'prediction': self.prediction.init(pred_key),
    }

  def inference(self, params: Any, observation: jnp.array) -> NetworkOutput:
    # representation + prediction function
    embedding = self.representation.apply(params['representation'], observation)
    return self.prediction.apply(params['prediction'], embedding)

  def get_params(self):
    # Returns the weights of this network.
    return self.params

  def update_params(self, updates: Any) -> None:
    # Update network weights internally.
    self.params = jax.tree_map(lambda p, u: p + u, self.params, updates)

  def training_steps(self) -> int:
    # How many steps / batches the network has been trained for.
    return 0


class UniformNetwork(object):
  """Network representation that returns uniform output."""

  # pylint: disable-next=unused-argument
  def inference(self, observation) -> NetworkOutput:
    # representation + prediction function
    return NetworkOutput(0, 0, 0, {})

  def get_params(self):
    # Returns the weights of this network.
    return self.params

  def update_params(self, updates: Any) -> None:
    # Update network weights internally.
    self.params = jax.tree_map(lambda p, u: p + u, self.params, updates)

  def training_steps(self) -> int:
    # How many steps / batches the network has been trained for.
    return 0

排序测试函数

using TestCases = std::vector<std::pair<std::vector<int>, std::vector<int>>>;

TestCases GenerateSortTestCases(int items_to_sort) {
  TestCases test_cases;
  auto add_all_permutations = [&test_cases](const std::vector<int>& initial) {
    std::vector<int> perm(initial);
    do {
      std::vector<int> expected = perm;
      std::sort(expected.begin(), expected.end());
      test_cases.push_back({perm, expected});
    } while (std::next_permutation(perm.begin(), perm.end()));
  };
  // Loop over all possible configurations of binary relations on sorted items.
  // Between each two consecutive items we can insert either '==' or '<'. Then,
  // for each configuration we generate all possible permutations.
  for (int i = 0; i < 1 << (items_to_sort - 1); ++i) {
    std::vector<int> relation = {1};
    for (int mask = i, j = 0; j < items_to_sort - 1; mask /= 2, ++j) {
      relation.push_back(mask % 2 == 0 ? relation.back() : relation.back() + 1);
    }
    add_all_permutations(relation);
  }
  return test_cases;
}

TestCases GenerateVariableSortTestCases(int max_items_to_sort) {
  TestCases test_cases;
  for (int num_items = 1; num_items <= max_items_to_sort; ++num_items) {
    TestCases base_test_cases = GenerateSortTestCases(num_items);
    for (auto [input, expected] : base_test_cases) {
      input.insert(input.begin(), num_items);
      expected.insert(expected.begin(), num_items);
      test_cases.push_back({input, expected});
    }
  }
  return test_cases;
}

void VerifyFunction(const TestCases& test_cases, std::function<void(int*)> fn) {
  for (const auto& [input, expected_output] : test_cases) {
    std::vector<int> output = input;
    fn(&output[0]);
    EXPECT_EQ(output, expected_output);
  }
}

TEST(SortingFunctionsTest, VerifyTestCases) {
  TestCases expected_test_cases = {
      {{1, 1, 1}, {1, 1, 1}},  //
      {{1, 2, 2}, {1, 2, 2}},  //
      {{2, 1, 2}, {1, 2, 2}},  //
      {{2, 2, 1}, {1, 2, 2}},  //
      {{1, 1, 2}, {1, 1, 2}},  //
      {{1, 2, 1}, {1, 1, 2}},  //
      {{2, 1, 1}, {1, 1, 2}},  //
      {{1, 2, 3}, {1, 2, 3}},  //
      {{1, 3, 2}, {1, 2, 3}},  //
      {{2, 1, 3}, {1, 2, 3}},  //
      {{2, 3, 1}, {1, 2, 3}},  //
      {{3, 1, 2}, {1, 2, 3}},  //
      {{3, 2, 1}, {1, 2, 3}},  //
  };
  EXPECT_EQ(expected_test_cases, GenerateSortTestCases(3));
}

本文参考文献

https://www.nature.com/articles/s41586-023-06004-9

https://github.com/deepmind/alphadev

https://www.nature.com/articles/d41586-023-01812-5

https://www.deepmind.com/blog/alphadev-discovers-faster-sorting-algorithms