causal-rl/load_model.py at main · CharlesCDietzel/causal-rl · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
"""
Model testing

"""

import gymnasium as gym
import torch
from causal_learn import CausalWorldModel
from causallearn.utils.GraphUtils import GraphUtils
from carla_.carla_env import CarlaEnv
import click
import os
import matplotlib.pyplot as plt
import numpy as np


@click.command()
@click.option(
    "--env-name",
    default="MountainCarContinuous-v0",
    type=str,
    help="Name of RL environment to use",
)

# Create environment to populate from saved model
def main(
    env_name: str,
):

    if env_name == "carla":
        base_env = CarlaEnv()
    else:
        base_env = gym.make(env_name)
    state_dim = base_env.observation_space.shape[0]
    action_dim = (
        base_env.action_space.n
        if isinstance(base_env.action_space, gym.spaces.Discrete)
        else base_env.action_space.shape[0]
    )
    world_model_hidden_dim = 64
    world_model = CausalWorldModel(state_dim, action_dim, world_model_hidden_dim)

    # Load saved model

    model_path = f"models/{env_name}/"
    checkpoint = torch.load(
        f"{model_path}world_model_with_graphs.pt", weights_only=False
    )
    world_model.load_state_dict(checkpoint["model_state_dict"])
    world_model.causal_graphs = checkpoint["causal_graphs"]
    print(f"Loaded world model with {len(world_model.causal_graphs)} causal graphs")

    # Save all causal graphs from model
    causal_graph_path = f"output/{env_name}/causal_graphs/"
    os.makedirs(causal_graph_path, exist_ok=True)
    for filename in os.listdir(causal_graph_path):
        file_path = os.path.join(causal_graph_path, filename)
        if os.path.isfile(file_path):  # Check if it's a file (not a subdirectory)
            try:
                os.remove(file_path)
                # print(f"Deleted: {file_path}")
            except OSError as e:
                print(f"Error deleting {file_path}: {e}")
    causal_matrix_history = []
    for i, graph in enumerate(world_model.causal_graphs):
        pyd = GraphUtils.to_pydot(graph.G)
        pyd.write_png(f"{causal_graph_path}causal_graph_{i}.png")

        full_adj_matrix = graph.G.graph
        input_dim = world_model.state_dim + world_model.action_dim
        causal_matrix = full_adj_matrix[
            :input_dim, input_dim : input_dim + world_model.state_dim
        ]
        causal_matrix = (np.abs(causal_matrix) > 0).astype(float)
        causal_matrix_history.append(causal_matrix)
    causal_matrix_mean = np.mean(causal_matrix_history, axis=0)

    causal_matrix_path = f"output/{env_name}/causal_adjacency_matrix.png"
    # os.makedirs(causal_matrix_path, exist_ok=True)

    fig, ax = plt.subplots(figsize=(8, 8))
    im = ax.imshow(causal_matrix_mean, cmap="Blues", vmin=0, vmax=1)
    cbar = plt.colorbar(im, ax=ax)
    cbar.set_label(
        "Frequency of causal graph edge", rotation=270, labelpad=20, fontsize=16
    )
    cbar.ax.tick_params(labelsize=12)
    ax.set_title(f"Average Causal Adjacency Matrix for {env_name}", fontsize=16, pad=16)
    ax.set_xlabel("Next step states", fontsize=16)
    ax.set_ylabel("Current step states and actions", fontsize=16)
    ax.set_xticks(np.arange(-0.5, causal_matrix_mean.shape[1], 1), minor=True)
    ax.set_yticks(np.arange(-0.5, causal_matrix_mean.shape[0], 1), minor=True)
    ax.set_xticks(np.arange(0, causal_matrix_mean.shape[1], 1), minor=False)
    ax.set_yticks(np.arange(0, causal_matrix_mean.shape[0], 1), minor=False)
    ax.grid(color="black", linestyle="-", linewidth=1.0, alpha=1.0, which="minor")
    x_axis_labels = []
    y_axis_labels = []
    for i in range(state_dim):
        x_axis_labels.append(f"s{i},t")
        y_axis_labels.append(f"s{i},t-1")
    for i in range(action_dim):
        y_axis_labels.append(f"a{i},t-1")
    ax.set_xticklabels(x_axis_labels, fontsize=12)
    ax.set_yticklabels(y_axis_labels, fontsize=12)
    ax.set_aspect("equal")
    plt.tight_layout()
    plt.savefig(causal_matrix_path)


if __name__ == "__main__":
    main()