<jupyter_start><jupyter_text>MultiTask NMF (Discrete)<jupyter_code>n_abs = 10
l1_ratio=0.0 # this is currently not working... since alpha is not set
state_dim = 16
action_dim = np.array((5, 5))
horizon = 100
n_trajs = 10

def get_expert(weight_path, state_dim, action_dim):
    expert = CategoricalActorCriticNet(
        4,
        state_dim,
        action_dim.prod(),
        FCBody(
            state_dim, 
            hidden_units=(16,)
        ),
        SplitBody(
            MultiLinear(16, action_dim.sum(), 4, key='task_id', w_scale=1e-3),
            2,
        ),
    )
    # load weight
    weight_dict = expert.state_dict()
    loaded_weight_dict = {k: v for k, v in torch.load(
        weight_path,
        map_location=lambda storage, loc: storage)['network'].items()
        if k in weight_dict}
    weight_dict.update(loaded_weight_dict)
    expert.load_state_dict(weight_dict)
    return expert

set_seed(0)

expert_dict = {
    1: '../log/reacher.1_corner/fc_discrete.baseline/split/0.190315-202731/models/step-704000-mean--6.36',
    2: '../log/reacher.2_corner/fc_discrete.baseline/split/0.190315-203310/models/step-704000-mean--17.16',
    3: '../log/reacher.3_corner/fc_discrete.baseline/split/0.190315-203532/models/step-704000-mean--11.25',
    #1: '../log/reacher.ng.1_corner/fc_discrete.baseline/ng/0.190316-161504/models/step-704000-mean--6.08',
    #2: '../log/reacher.ng.2_corner/fc_discrete.baseline/ng/0.190316-161535/models/step-704000-mean--26.97',
    #3: '../log/reacher.ng.3_corner/fc_discrete.baseline/ng/0.190316-162113/models/step-704000-mean--16.12',
}

envs = [DiscretizeActionEnv(
    MultiGoalReacherEnv(
        [
            [0.15, 0.0],
            [-0.15, 0.0],
            [0.0, 0.15],
            [0.0, -0.15],
        ],
        sample_indices=[i],
        with_goal_pos=True,
    ),
    n_bins=[5, 5],
) for i in range(3)]
decomposer = MTNMF(n_abs, max_iter=5000, tol=0.0001)

states = []
experts = dict()

for goal_idx, weight_path in expert_dict.items():
    experts[goal_idx] = get_expert(weight_path, state_dim, action_dim)
    for _ in range(n_trajs // 2):
        states.append(rollout(envs[goal_idx-1], experts[goal_idx], horizon=horizon))
    for _ in range(n_trajs - (n_trajs // 2)):
        states.append(rollout(envs[goal_idx-1], None, horizon=horizon))
states = np.concatenate(states)
print('states shape:', states.shape)
    
pvs = []
    
for goal_idx in expert_dict:
    infos = {'task_id': [goal_idx-1] * len(states)}
    pv = F.softmax(experts[goal_idx].get_logits(states, infos), dim=-1).cpu().detach().numpy()
    print(pv.shape)
    pv = pv.reshape(pv.shape[0], -1)
    pvs.append(pv)

pvs = np.stack(pvs, 0)
A, S, info = MTNMF(n_abs, max_iter=5000, l1_ratio=l1_ratio).fit(pvs.transpose(0, 2, 1))
print(pvs.shape)

fsave(
    dict(
        abs=S.T,
        policies=list(pvs.reshape(pvs.shape[0], pvs.shape[1], 2, -1)),
        states=[states for _ in range(len(pvs))],
        infos=list([[{'task_id': i} for _ in range(len(states))] for i in range(3)]),
    ),
    '../data/nmf_sample/reacher/split.mix.{}'.format(n_abs),
    'pkl',
)<jupyter_output>[33mWARN: gym.spaces.Box autodetected dtype as <class 'numpy.float32'>. Please provide explicit dtype.[0m
[33mWARN: gym.spaces.Box autodetected dtype as <class 'numpy.float32'>. Please provide explicit dtype.[0m
[33mWARN: gym.spaces.Box autodetected dtype as <class 'numpy.float32'>. Please provide explicit dtype.[0m
states shape: (3000, 16)
(3000, 2, 5)
(3000, 2, 5)
(3000, 2, 5)
NMF loss: 4324.044033279368
NMF loss: 384.95037468093733
NMF loss: 337.3501477731719
NMF loss: 318.22234011513086
NMF loss: 312.56009875841835
NMF loss: 309.3346029731073
NMF loss: 308.1361014290645
NMF loss: 305.28736316119796
NMF loss: 305.03279790818027
NMF loss: 304.78534478715176
NMF loss: 304.6378301112521
NMF loss: 304.51718465325223
NMF loss: 304.45668584609433
NMF loss: 304.41464215417295
NMF loss: 304.387090696981
NMF loss: 304.3220786372614
NMF loss: 304.24312767222114
NMF loss: 304.19769493610545
NMF loss: 304.1641330567978
NMF loss: 304.146955503876
(3, 3000, 10)<jupyter_text>MultiTask NMF (continuous)<jupyter_code>n_abs = 6
l1_ratio=0.0 # this is currently not working... since alpha is not set
state_dim = 8
action_dim = 2
horizon = 100
n_trajs = 10
n_bins = 10

set_seed(0)

def get_expert(weight_path, state_dim, action_dim):
    expert = GaussianActorCriticNet(
        4,
        state_dim,
        action_dim,
        FCBody(
            state_dim, 
            hidden_units=(32,)
        ),
    )
    # load weight
    weight_dict = expert.state_dict()
    loaded_weight_dict = {k: v for k, v in torch.load(
        weight_path,
        map_location=lambda storage, loc: storage)['network'].items()
        if k in weight_dict}
    weight_dict.update(loaded_weight_dict)
    expert.load_state_dict(weight_dict)
    return expert

expert_dict = {
    1: '../log/reacher.ng.1_corner/fc_discrete.gaussian/cont/0.190316-212125/models/step-128000-mean--18.15',
    2: '../log/reacher.ng.2_corner/fc_discrete.gaussian/cont/0.190316-212209/models/step-128000-mean--11.82',
    3: '../log/reacher.ng.3_corner/fc_discrete.gaussian/cont/0.190316-214339/models/step-128000-mean--5.18',
}

envs = [MultiGoalReacherEnv(
        [
            [0.15, 0.0],
            [-0.15, 0.0],
            [0.0, 0.15],
            [0.0, -0.15],
        ],
        sample_indices=[i],
        with_goal_pos=False,
) for i in range(3)]
decomposer = MTNMF(n_abs, max_iter=5000, tol=0.0001)

states = []
ax = np.linspace(-1, 1, n_bins)
ay = np.linspace(-1, 1, n_bins)
ax, ay = np.meshgrid(ax, ay)
actions = np.stack([ax, ay], -1)

experts = dict()

for goal_idx, weight_path in expert_dict.items():
    experts[goal_idx] = get_expert(weight_path, state_dim, action_dim)
    for _ in range(n_trajs):
        states.append(rollout(envs[goal_idx-1], experts[goal_idx], horizon=horizon))
states = np.concatenate(states)
print('states shape:', states.shape)
    
pvs = []
    
for goal_idx in expert_dict:
    infos = {'task_id': [goal_idx-1] * len(states)}
    mean = experts[goal_idx](states, infos)['mean'].cpu().detach().numpy()
    std = experts[goal_idx].std.expand(len(states), *std.shape)
    print(mean.shape, std.shape)
    #pv = F.softmax(experts[goal_idx].get_logits(states, infos), dim=-1).cpu().detach().numpy()
    pv = pv.reshape(pv.shape[0], -1)
    pvs.append(pv)

pvs = np.stack(pvs, 0)
A, S, info = MTNMF(n_abs, max_iter=5000, l1_ratio=l1_ratio).fit(pvs.transpose(0, 2, 1))

fsave(
    dict(
        abs=S.T,
        policies=list(pvs.reshape(pvs.shape[0], pvs.shape[1], 2, -1)),
        states=[states for _ in range(len(pvs))],
        infos=list([[{'task_id': i} for _ in range(len(states))] for i in range(3)]),
    ),
    '../data/nmf_sample/reacher/cont.ng.{}'.format(n_abs),
    'pkl',
)

import torch
import numpy as np

print(torch.__version__)

a = torch.Tensor([1, 2, 3])
print(a.expand(3, 3))

# (1, 2, 3)
#probs = torch.Tensor([[[0.2, 0.3, 0.5], [0.1, 0.3, 0.6]]])
#print(torch.nn.functional.softmax(probs, dim=2))
# probs = torch.Tensor([[0.2, 0.3, 0.5]])
# print(probs.shape[0])
# log_probs = torch.log(probs)
# print(log_probs)
# dist = torch.distributions.Categorical(logits=log_probs)
# action = dist.sample()
# print(action)
# log_prob = dist.log_prob(action)
# print(log_prob)<jupyter_output>0.4.0
tensor([[ 1.,  2.,  3.],
        [ 1.,  2.,  3.],
        [ 1.,  2.,  3.]])
{'fraction_non_alphanumeric': 0.1615685217625657, 'fraction_numerical': 0.10618515024929255, 'mean_word_length': 3.231470923603193, 'pattern_counts': {'":': 0, '<': 10, '<?xml version=': 0, '>': 10, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '425749', 'license': 'Apache-2.0', 'n_tokens_mistral': 3480, 'n_tokens_neox': 2987, 'n_words': 535, 'path': 'ipynb/NMF_reacher.ipynb', 'repo_name': 'lchenat/TSA', 'size': '7421'}