DOCSEARCHbyAlgolia
Submit Docs

LightZero Documentation

repository·main·Indexed 23 days ago

https://github.com/opendilab/lightzero

A reinforcement learning framework supporting algorithm families such as AlphaZero, MuZero, UniZero, and ReZero. It includes specialized tools for continuous action space learning, long-term strategic planning (e.g., Billiards RL via pooltool), and multi-task learning. The framework provides entry functions for training and evaluation, a Worker Module for experience collection and policy assessment, and a LossLandscape API for visualizing neural network loss surfaces.

Tokens
27.8K
Snippets
41
Records
155
Agent score
80%

What's inside LightZero

  1. Overview of LightZero Worker components

    main

    The lzero.worker module contains the core components for reinforcement learning algorithms, specifically focusing on data collection and performance evaluation. It is divided into two main types of components:

    1. Collectors: Responsible for gathering experience data through environment interaction during the training process.
    2. Evaluators: Responsible for assessing the performance of the trained policy at regular intervals during training.

    All worker components support distributed training (multi-process/multi-GPU), TensorBoard logging, multi-task learning (via the task_id parameter), and configurable collection/evaluation frequencies.

  2. Overview of LightZero

    main
    LightZero is an open-source, lightweight, efficient, and easy-to-understand algorithm library that combines Monte Carlo Tree Search (MCTS) with Reinforcement Learning (RL). It aims to standardize the MCTS algorithm family to accelerate research and applications in decision-making problems. The library is implemented using PyTorch and is built upon the DI-engine framework, utilizing Cython and C++ for high-performance MCTS implementations.

  3. Explore supported environments in the LightZero Zoo

    main

    The LightZero Zoo provides a wide variety of reinforcement learning environments categorized by their action space types (discrete or continuous) and domain. Supported environments include:

    Board Games (Discrete)

    • Tic-tac-toe: board_games/tictactoe
    • Gomoku: board_games/gomoku
    • Connect Four: board_games/connect4
    • Chess: board_games/chess
    • Go: board_games/go

    Classic Control & Physics (Discrete/Continuous)

    • CartPole: classic_control/cartpole (Discrete)
    • Pendulum: classic_control/pendulum (Continuous)
    • MountainCar: classic_control/mountain_car (Discrete)
    • LunarLander: box2d/lunarlander (Supports both Discrete and Continuous)
    • BipedalWalker: box2d/bipedalwalker (Continuous)
    • MuJoCo: mujoco (Continuous)
    • DMC2Gym: dmc2gym (Continuous)

    Specialized Environments

    • Atari: atari (Discrete)
    • MiniGrid: minigrid (Discrete)
    • Memory Tasks: memory (Discrete)
    • Jericho: jericho (Text-based adventure, Discrete)
    • PoolTool: pooltool/sum_to_three (Continuous)
    • CrowdSim: crowd_sim (Continuous)
    • MetaDrive: metadrive (Continuous)
    • Memory Maze: memory_maze (Discrete)
    • BSuite: bsuite (Discrete)
    • 2048: game_2048 (Discrete)

  4. Understand the Worker Module components

    main

    The Worker Module in LightZero provides two primary types of components for reinforcement learning:

    1. Collectors: Responsible for gathering experience data during training through environment interaction (self-play or environment interaction).
    2. Evaluators: Responsible for assessing the performance of trained policies at regular intervals during the training process.

    All workers support distributed training (multi-process/multi-GPU), TensorBoard logging, multi-task learning scenarios (using the task_id parameter), and configurable frequencies for collection and evaluation.

  5. Choose the correct Collector for your algorithm

    main

    LightZero provides different collectors depending on the algorithm and the required data granularity:

    • alphazero_collector.py: Used for AlphaZero. It uses Episode-based collection and is designed for perfect information games (e.g., board games).
    • muzero_collector.py: Used for MuZero, EfficientZero, or Gumbel MuZero. It uses Episode-based collection and supports both perfect and imperfect information environments.
    • muzero_segment_collector.py: Used for MuZero, EfficientZero, or Gumbel MuZero. It uses Segment-based collection, meaning it collects a specified number of game segments rather than complete episodes, providing more fine-grained control.

  6. Customize models and workers

    main

    If your custom algorithm requires more than just a new policy, you can extend other parts of the framework:

    • Models: If you need specialized neural architectures, implement them in the model folder. You can leverage common structures from model.common (like RepresentationNetwork or PredictionNetwork).
    • Workers: If your data collection logic differs (e.g., you need to preprocess transitions), implement a custom worker. For example, you can modify the collect function in a collector to process data using a method like self._policy.get_train_sample(transitions) when a trajectory is complete.
    if timestep.done:
    # Prepare trajectory data.
    transitions = to_tensor_transitions(self._traj_buffer[env_id])
    # Use ``get_train_sample`` to process the data.
    train_sample = self._policy.get_train_sample(transitions)
    return_data.extend(train_sample)
    self._traj_buffer[env_id].clear()

  7. Best practices for custom environment design

    main

    When designing or customizing environments for LightZero, consider the following three pillars:

    1. State Representation: Determine how to represent the environment state in the observation space. Use low-dimensional continuous states for simple environments, or high-dimensional discrete states (like images) for complex ones.
    2. Observation Preprocessing: Apply appropriate preprocessing based on the observation type (e.g., scaling, cropping, grayscale conversion, or normalization) to reduce dimensionality and accelerate learning.
    3. Reward Design: Design reward functions that align with your goals. It is recommended to normalize extrinsic rewards to the range [0, 1]. This normalization helps in more effectively tuning hyperparameters, such as the intrinsic reward weight in RND (Random Network Distillation) algorithms.

  8. How ROPE (Relative Position Encoding) works in UniZero

    main

    When self.config.rotary_emb = True, the model employs Rotary Position Embedding (ROPE).

    • Implementation: It applies rotation to Query and Key tensors using precomputed frequency components, integrating position information directly into self-attention.
    • Indexing Logic: Indices are assigned based on the episode time step. In environments where states (s) and actions (a) alternate, each time step occupies two indices. For example, a 50-step episode with alternating (s, a) pairs results in position indices 1, 2, 3, ..., 100.
    • Advantages: ROPE provides higher flexibility for variable sequence lengths and allows inter-token dependency to decay as relative distance increases. It is recommended for environments with long-range dependencies, whereas Absolute Encoding is sufficient for short-dependency environments like Pong or DMC Cartpole-Swingup.

  9. Implement specialized logic for board game environments

    main

    Board games in LightZero require specific implementations to support multi-player logic and MCTS.

    1. Operating Modes

    Implement logic to handle these three modes:

    • self_play_mode: Standard setup. Each step performs one move. Reward is +1 for a win, 0 otherwise.
    • play_with_bot_mode: The agent plays against a built-in bot. The agent is player 1, the bot is player 2. Reward is +1 if agent wins, -1 if bot wins, 0 for draw.
    • eval_mode: Used for evaluation. Can use a bot or a human (command line input) as the opponent.

    2. Legal Actions

    Implement legal_actions() to return a list of valid moves. This is critical for MCTS to generate valid child nodes.

    def legal_actions(self) -> List[int]:
    return [i for i in range(7) if self.board[i] == 0]

    3. Bot and Random Actions

    To support play_with_bot_mode and stochasticity, implement:

    • bot_action(): Returns an action generated by a bot (e.g., rule-based or MCTS-based) based on self.bot_action_type.
    • random_action(): Returns a random action from the current legal_actions list.
    def bot_action(self) -> int:
    if np.random.rand() < self.prob_random_action_in_bot:
        return self.random_action()
    else:
        if self.bot_action_type == 'rule':
            return self.rule_bot.get_rule_bot_action(self.board, self._current_player)
        elif self.bot_action_type == 'mcts':
            return self.mcts_bot.get_actions(self.board, player_index=self.current_player_index)

  10. How LightZero environment observations differ from DI-engine

    main

    In LightZero, the environment observation (obs) is a dictionary rather than a simple array (unlike DI-engine). This design accommodates board games where state information must include more than just raw observations.

    For all environments (including non-board games), the obs dictionary must contain:

    • 'observation': The actual state/observation data.
    • 'action_mask': A mask indicating legal actions.
    • 'to_play': An integer indicating the current player. For single-player (non-board) environments, set this to -1. For multi-player board games, set this to the player index.

    Board Game Specifics: For board games, the dictionary may also include additional keys like 'board' and 'current_player_index' to support MCTS (Monte Carlo Tree Search) workflows.

  11. Understand the Sum to Three game rules

    main

    Sum to Three is a simplified billiards game designed for reinforcement learning.

    Rules:

    1. Played on a table with no pockets.
    2. Contains 2 balls: a cue ball and an object ball.
    3. The player must hit the object ball with the cue ball.
    4. A point is scored if the total number of times a ball hits a cushion is exactly 3.
    5. The player takes 10 shots; the final score is the total number of points achieved.

  12. Understand LightZero's file directory structure

    main

    When running an experiment in LightZero, the framework organizes output into a specific directory structure. The main components are:

    • ckpt/: Stores model checkpoints (e.g., ckpt_best.pth.tar and periodic iteration_N.pth.tar files).
    • log/: Contains detailed logs categorized by component:
      • buffer/: Buffer-related logs.
      • collector/: Logs from the data collection stage.
      • evaluator/: Logs from the evaluation stage.
      • learner/: Logs from the model training process.
      • serial/: Contains Tensorboard event files for real-time monitoring.
    • total_config.py & formatted_total_config.py: Record the configuration used for the experiment.
    cartpole_muzero
├── ckpt
│   ├── ckpt_best.pth.tar
│   ├── iteration_0.pth.tar
│   └── iteration_10000.pth.tar
├── log  
│   ├── buffer
│   │   └── buffer_logger.txt
│   ├── collector
│   │   └── collector_logger.txt
│   ├── evaluator
│   │   └── evaluator_logger.txt
│   ├── learner
│   │   └── learner_logger.txt
│   └── serial
│       └── events.out.tfevents.1626453528.CN0014009700M.local
├── formatted_total_config.py
└── total_config.py