pi-zero-pytorch

To orchestrate online learning, wrap your π0 model with the EFPO class. This allows you to gather experience from an environment and learn from those memories.

Workflow

1. Initialize EFPO: Pass your trained or untrained model to EFPO(model).
2. Gather Experience: Use epo.gather_experience_from_env(env, steps=N) to collect memories from your environment.
3. Learn: Use epo.learn_agent(memories, batch_size=N) to update the model using the gathered memories.

from pi_zero_pytorch import π0, EFPO

# 1. Initialize model
model = π0(
    dim = 512,
    dim_action_input = 6,
    dim_joint_state = 12,
    num_tokens = 20_000
)

# 2. Setup environment (replace with your own environment)
from pi_zero_pytorch.mock_env import Env
mock_env = Env((256, 256), 2, 32, 1024, 12)

# 3. Wrap model with EFPO
epo = EFPO(model)

# 4. Gather memories from environment
memories = epo.gather_experience_from_env(mock_env, steps = 10)

# 5. Learn from memories
epo.learn_agent(memories, batch_size = 2)

The Agent class wraps a PiZero model and manages its training components. It initializes an actor (the policy) and a critic (the value function).

Evolutionary Learning: If num_latent_genes > 1, the agent enables evolutionary learning by maintaining a LatentGenePool. This allows the agent to evolve latent representations. You can trigger a genetic algorithm step using take_genetic_algorithm_step_(fitnesses).

Optimization: The agent maintains separate optimizers for the actor (actor_optim) and the critic (critic_optim), allowing for different learning rates and weight decay settings.

from pi_zero_pytorch.pi_zero import Agent

# Creating an agent with evolutionary capabilities
agent = Agent(
    model=my_pi_zero_model,
    num_latent_genes=5,  # Enables evolutionary learning
    actor_lr=3e-4,
    critic_lr=3e-4,
    actor_fpo_loss_fn=torch.nn.functional.huber_loss
)

# Perform a genetic step if evolutionary learning is active
agent.take_genetic_algorithm_step_(fitnesses=my_fitness_scores)

The π0 class implements the robotic foundation model architecture.

Training

During training, pass vision, commands, joint_state, and actions to the model. It returns a loss tensor.

Inference (Sampling)

After training, you can sample actions by providing vision, commands, and joint_state, and specifying a trajectory_length. This returns the sampled actions.

Parameters

• dim: Dimension of the model.
• dim_action_input: Dimension of the action input.
• dim_joint_state: Dimension of the joint state.
• num_tokens: Number of tokens.

import torch
from pi_zero_pytorch import π0

model = π0(
    dim = 512,
    dim_action_input = 6,
    dim_joint_state = 12,
    num_tokens = 20_000
)

vision = torch.randn(1, 1024, 512)
commands = torch.randint(0, 20_000, (1, 1024))
joint_state = torch.randn(1, 12)
actions = torch.randn(1, 32, 6)

# Training step
loss, _ = model(vision, commands, joint_state, actions)
loss.backward()

# Inference/Sampling
sampled_actions = model(vision, commands, joint_state, trajectory_length = 32) # (1, 32, 6)

Use sample_actions to generate a sequence of actions (a trajectory) from visual, language, and state inputs. This method implements a flow-matching sampling process (ODE integration).

Key features:

• Inpainting/Real-Time Chunking: If frozen_actions are provided, the model can perform inpainting to ensure the generated trajectory is consistent with already executed actions.
• Classifier-Free Guidance (CFG): If reward_tokens are provided, you can use cond_scale to steer the sampling towards higher rewards.
• Critic Support: If a critic module is passed, the method also returns predicted critic values for the sampled trajectory.
• RTC Guidance: Supports Real-Time Chunking (RTC) guidance if rtc_guidance was configured during initialization.

# Basic sampling
actions = model.sample_actions(
    images=images,           # [b, nv, d] or [b, c, h, w]
    token_ids=token_ids,     # [b, nt]
    joint_states=joint_states, # [b, djs]
    trajectory_length=16,
    steps=18                 # Number of ODE steps
)

# Sampling with reward guidance and frozen actions (for chunking)
actions, values = model.sample_actions(
    images=images,
    token_ids=token_ids,
    joint_states=joint_states,
    trajectory_length=16,
    frozen_actions=frozen_actions, # [b, nfa, da]
    reward_tokens=reward_tokens,    # [b, d]
    cond_scale=2.0,                # Guidance scale
    critic=critic_model             # Optional critic
)