DOCSEARCHbyAlgolia
Submit Docs

q-transformer

repository·main·Indexed 19 days ago

https://github.com/lucidrains/q-transformer

An implementation of Scalable Offline Reinforcement Learning via Autoregressive Q-Functions. It provides the QRoboticTransformer attention model, a QLearner for training on ReplayMemoryDataset, and an Agent class for environment interaction. The library supports custom environments via BaseEnvironment and allows for the prediction of optimal actions based on video tensors and text instructions.

Tokens
1.2K
Snippets
5
Records
6
Agent score
15%

What's inside q-transformer

  1. Implement a custom environment by overriding BaseEnvironment

    main

    To use the agent and learner, you must provide an environment by overriding BaseEnvironment. The environment must implement two specific methods:

    1. env.init(): Must return a tuple containing instructions (string) and the initial state (Tensor).
      • Return type: Tuple[str, Tensor[*state_shape]]
    2. env(actions): Must accept actions and return a tuple containing rewards, the next state, and a done flag.
      • Return type: Tuple[Tensor[()], Tensor[*state_shape], Tensor[()]]

  3. Generate replay memory using the Agent class

    main

    The Agent class facilitates interaction between the QRoboticTransformer model and your custom environment to generate a ReplayMemoryDataset for training.

    Parameters:

    • model: The QRoboticTransformer instance.
    • environment: Your custom environment instance.
    • num_episodes: Total number of episodes to run.
    • max_num_steps_per_episode: Maximum steps allowed per episode.
    from q_transformer import Agent

agent = Agent(
    model,
    environment = env,
    num_episodes = 1000,
    max_num_steps_per_episode = 100,
)

agent()

  4. Get optimal actions from the model

    main

    Once trained, use the get_optimal_actions method on the QRoboticTransformer model to perform tasks. This method takes a video tensor and a list of text instructions and returns the predicted optimal actions.

    Arguments:

    • video: A tensor representing the visual input (e.g., shape [batch, channels, frames, height, width]).
    • instructions: A list of strings containing the task instructions.
    import torch

video = torch.randn(2, 3, 6, 224, 224)
instructions = [
    'bring me that apple sitting on the table',
    'please pass the butter'
]

actions = model.get_optimal_actions(video, instructions)

  5. Initialize QRoboticTransformer model

    main

    The QRoboticTransformer is the core attention model. It requires configuration for a Vision Transformer (vit) backbone, the number of actions, and action discretization parameters.

    Key parameters include:

    • vit: A dictionary configuring the vision backbone (e.g., num_classes, dim, depth, window_size).
    • num_actions: The number of actions to be performed.
    • action_bins: The number of discrete bins for actions.
    • dueling: Boolean to enable dueling architecture.
    from q_transformer import QRoboticTransformer

model = QRoboticTransformer(
    vit = dict(
        num_classes = 1000,
        dim_conv_stem = 64,
        dim = 64,
        dim_head = 64,
        depth = (2, 2, 5, 2),
        window_size = 7,
        mbconv_expansion_rate = 4,
        mbconv_shrinkage_rate = 0.25,
        dropout = 0.1
    ),
    num_actions = 8,
    action_bins = 256,
    depth = 1,
    heads = 8,
    dim_head = 64,
    cond_drop_prob = 0.2,
    dueling = True
)

  6. Train the model using QLearner

    main

    The QLearner performs Q-learning on a provided ReplayMemoryDataset using the QRoboticTransformer model.

    Parameters:

    • model: The QRoboticTransformer instance.
    • dataset: An instance of ReplayMemoryDataset containing collected experiences.
    • num_train_steps: Total number of training steps.
    • learning_rate: The optimizer learning rate.
    • batch_size: Number of samples per batch.
    • grad_accum_every: Number of steps to accumulate gradients before updating.
    from q_transformer import QLearner, ReplayMemoryDataset

q_learner = QLearner(
    model,
    dataset = ReplayMemoryDataset(),
    num_train_steps = 10000,
    learning_rate = 3e-4,
    batch_size = 4,
    grad_accum_every = 16,
)

q_learner()