122 lines
4.1 KiB
Python
122 lines
4.1 KiB
Python
"""Load and manage Disco103 pre-trained weights.
|
|
|
|
This module handles loading the Disco103 meta-parameters from the npz file
|
|
provided in the DiscoRL repository and integrating them with DiscoRL agents.
|
|
"""
|
|
|
|
import os
|
|
from typing import Any, Dict, Tuple
|
|
import numpy as np
|
|
import jax
|
|
import jax.numpy as jnp
|
|
|
|
|
|
def load_disco103_weights(disco_rl_path: str = None) -> Dict[str, Any]:
|
|
"""Load Disco103 pre-trained weights.
|
|
|
|
Args:
|
|
disco_rl_path: path to disco_rl repo root. If None, search from cwd.
|
|
|
|
Returns:
|
|
dict with keys like 'meta_params' or structure matching the npz file
|
|
"""
|
|
if disco_rl_path is None:
|
|
# Try to find disco_rl in standard locations
|
|
possible_paths = [
|
|
'disco_rl/disco_rl/update_rules/weights/disco_103.npz',
|
|
'../disco_rl/disco_rl/update_rules/weights/disco_103.npz',
|
|
'../../disco_rl/disco_rl/update_rules/weights/disco_103.npz',
|
|
]
|
|
npz_path = None
|
|
for p in possible_paths:
|
|
if os.path.exists(p):
|
|
npz_path = p
|
|
break
|
|
if npz_path is None:
|
|
raise FileNotFoundError(
|
|
'Could not find disco_103.npz. Please provide disco_rl_path '
|
|
'or ensure disco_rl/ is accessible.'
|
|
)
|
|
else:
|
|
npz_path = os.path.join(
|
|
disco_rl_path, 'disco_rl/update_rules/weights/disco_103.npz'
|
|
)
|
|
|
|
if not os.path.exists(npz_path):
|
|
raise FileNotFoundError(f'disco_103.npz not found at {npz_path}')
|
|
|
|
# Load the npz file
|
|
data = np.load(npz_path, allow_pickle=True)
|
|
|
|
# Convert to dictionary; npz files can be accessed as dict-like
|
|
weights = {}
|
|
for key in data.files:
|
|
item = data[key]
|
|
# Some items might be numpy object arrays (e.g., nested structures)
|
|
# Try to convert to jax arrays where possible
|
|
if isinstance(item, np.ndarray):
|
|
weights[key] = jnp.asarray(item)
|
|
else:
|
|
weights[key] = item
|
|
|
|
print(f'Loaded Disco103 weights from {npz_path}')
|
|
print(f' Keys: {list(weights.keys())}')
|
|
for key, val in weights.items():
|
|
if hasattr(val, 'shape'):
|
|
print(f' {key}: shape={val.shape}, dtype={val.dtype}')
|
|
else:
|
|
print(f' {key}: {type(val)}')
|
|
|
|
return weights
|
|
|
|
|
|
def unflatten_disco_weights(flat_dict: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""Convert flat npz weight dict to nested structure expected by DiscoRL.
|
|
|
|
The exact structure depends on how the weights were saved. This is a
|
|
placeholder; you may need to adjust based on the actual npz structure.
|
|
|
|
For now, we assume the npz contains the meta_params directly or under
|
|
a 'meta_params' key.
|
|
"""
|
|
# If there's a 'meta_params' key, use it; otherwise assume flat_dict IS the params
|
|
if 'meta_params' in flat_dict:
|
|
meta_params = flat_dict['meta_params']
|
|
else:
|
|
# Try to reconstruct nested structure from flat keys
|
|
# This is environment-specific; adjust as needed
|
|
meta_params = flat_dict
|
|
|
|
return meta_params
|
|
|
|
|
|
def merge_weights_with_agent(
|
|
agent_meta_state: Dict[str, Any],
|
|
disco_weights: Dict[str, Any],
|
|
) -> Dict[str, Any]:
|
|
"""Merge loaded Disco103 weights into agent's meta_state.
|
|
|
|
This updates the meta_state's rnn_state and other components with
|
|
pre-trained weights if available.
|
|
|
|
Args:
|
|
agent_meta_state: the agent's initial meta_state dict
|
|
disco_weights: loaded weights dict
|
|
|
|
Returns:
|
|
updated agent_meta_state
|
|
"""
|
|
# For now, we mainly care about the meta_params (update_rule weights)
|
|
# The rnn_state and ema_state are often initialized fresh during
|
|
# agent creation, but we can override them if they're in disco_weights.
|
|
|
|
updated_state = dict(agent_meta_state)
|
|
|
|
# If the npz has useful rnn_state or other components, merge them
|
|
# This is a placeholder; adjust based on actual npz structure
|
|
for key in ['rnn_state', 'adv_ema_state', 'td_ema_state']:
|
|
if key in disco_weights:
|
|
updated_state[key] = disco_weights[key]
|
|
|
|
return updated_state
|