Перейти к основному содержимому

search

Этот модуль предоставляет функцию для выполнения лучевого поиска и нахождения решений для заданного состояния.

Функция: beam_search: Выполняет лучевой поиск для нахождения решений в среде кубика Рубика.

MAX_BATCH_SIZE​

Максимальное количество состояний, одновременно обрабатываемых нейронной сетью (DNN) за один проход.

def beam_search(
    env,
    model,
    beam_width,
    ergonomic_bias=None,
    extra_depths=0,
    max_depth=100
)

Выполняет лучевой поиск для нахождения решений для заданного разобранного состояния.

Аргументы:

  • env Cube3 - Среда кубика Рубика, представляющая разобранное состояние.
  • model torch.nn.Module - Нейронная сеть (DNN), используемая для предсказания распределения вероятностей следующих ходов для каждого состояния.
  • beam_width int - Максимальное количество кандидатов, сохраняемых на каждом шаге поиска.
  • ergonomic_bias dict или None - Словарь, задающий эргономическое смещение для ходов, если доступно.
  • extra_depths int - Количество дополнительных уровней глубины для поиска после нахождения первого решения.
  • max_depth int - Максимальная глубина поиска, должна быть равна или больше числа Бога (20 для кубика Рубика в метрике HTM).

Возвращает:

  • dict | None: При наличии хотя бы одного решения, возвращает словарь со следующими ключами:
  1. "solutions": Список оптимальных или близких к оптимальным решений, найденных в ходе поиска.
  2. "num_nodes": Общее количество узлов, развернутых в ходе поиска.
  3. "time": Время, затраченное на поиск (в секундах).

В противном случае, None.

_reflect_setup​

def _reflect_setup(ergonomic_bias, env)

Инициализирует эргономическое смещение, если оно предоставлено.

Аргументы:

  • ergonomic_bias dict или None - Словарь, задающий эргономическое смещение для ходов, если доступно.
  • env Cube3 - Среда кубика Рубика, представляющая разобранное состояние.

Возвращает:

  • ergonomic_bias numpy.ndarray - Эргономическое смещение для ходов, если доступно.
  • env Cube3 - Среда кубика Рубика, представляющая разобранное состояние.

predict​

@torch.inference_mode()
def predict(model, batch_x, ergonomic_bias, env)

Предсказывает распределение вероятностей следующих ходов для каждого состояния.

Аргументы:

  • model torch.nn.Module - Нейронная сеть (DNN), используемая для предсказания распределения вероятностей следующих ходов для каждого состояния.
  • batch_x numpy.ndarray - Пакет состояний.
  • ergonomic_bias dict или None - Словарь, задающий эргономическое смещение для ходов, если доступно.
  • env Cube3 - Среда кубика Рубика, представляющая разобранное состояние.

Возвращает:

  • batch_logprob numpy.ndarray - Логарифм вероятности каждого хода для каждого состояния.
примечание

Вывод (inference) с использованием Automatic Mixed Precision по некоторым причинам работает немного быстрее, чем простой вывод с половинной точностью (с model.half()).

update_candidates​

def update_candidates(candidates, batch_logprob, env, depth, beam_width)

Расширяет пути кандидатов с использованием предсказанных вероятностей следующих ходов.

Аргументы:

  • candidates dict - Словарь, содержащий пути кандидатов, накопленные вероятности и состояния.
  • batch_logprob numpy.ndarray - Логарифм вероятности каждого хода для каждого состояния.
  • env Cube3 - Среда кубика Рубика, представляющая разобранное состояние.
  • depth int - Текущая глубина поиска.
  • beam_width int - Максимальное количество кандидатов, сохраняемых на каждом шаге поиска.

Возвращает:

  • candidates dict - Обновленный словарь, содержащий пути кандидатов, накопленные вероятности и состояния.

_get_prune_idx​

def _get_prune_idx(candidates_paths, allow_wide, depth)

Получает индексы кандидатов для отсечения на основе предыдущих ходов.

Аргументы:

  • candidates_paths numpy.ndarray - Пути состояний-кандидатов.
  • allow_wide bool - Разрешать ли широкие ходы (wide moves).
  • depth int - Текущая глубина поиска.

Возвращает:

  • prune_idx numpy.ndarray - Индексы кандидатов для отсечения.
примечание

Использование numba.jit на самом деле замедляет эту функцию.

_update_states​

def _update_states(candidate_states, candidate_paths, env)

Обновляет состояния на основе расширенных путей.

Аргументы:

  • candidate_states numpy.ndarray - Состояния кандидатов.
  • candidate_paths numpy.ndarray - Пути состояний-кандидатов.
  • env Cube3 - Среда кубика Рубика, представляющая разобранное состояние.

Возвращает:

  • candidate_states numpy.ndarray - Обновленные состояния кандидатов.

_map_state​

def _map_state(candidate_states, target_ix, source_ix)

Выполняет замену наклеек на уровне пакета (batch).

Аргументы:

  • candidate_states numpy.ndarray - Состояния кандидатов.
  • target_ix numpy.ndarray - Целевые индексы для замены наклеек.
  • source_ix numpy.ndarray - Исходные индексы для замены наклеек.

Возвращает:

  • candidate_states numpy.ndarray - Обновленные состояния кандидатов.
примечание

Использование numba.jit на самом деле замедляет эту функцию.