Thanks to #27 we have model checkpoints from training. I'd like to do the following:

• Load these state dictionaries into a blank model
• Evaluate them
• Samples form them and write trajectories files which are much smaller than the training runs are for these bigger/messier tasks.

We currently have 5 probe environments for single timestep models and I'd like a prob environment to test if a model can learn:

1. to take the correct action as a function of a previous observation
2. to test that the model doesn't learn the correct action if it's outside the recurrence range or something (more speculative)

more details to add later or message me

Might be valuable for training performance.

My initial plan was to look at how decision transformers solved various tasks and compare this to trajectory transformers training with online RL (via PPO), however this may be difficult due to the instability of standard transformers.

On Transforming Reinforcement Learning with Transformers: The Development Trajectory - https://arxiv.org/pdf/2212.14164.pdf
STABILIZING TRANSFORMERS FOR REINFORCEMENT LEARNING - https://arxiv.org/pdf/1910.06764.pdf

In light of this, to make it easier to generate training data for the decision transformers. I want to try another architecture such as LSTM.

Emulating the architecture from the BabyAI paper would probably be a good way to go (handling mission statements would be nice)- https://github.com/mila-iqia/babyai/tree/master

Tasks:

• Implement an architecture such as LSTM in an agent class (such as LSTMAgent)
• Train this class with PPO
• Pass tests cases as well as showing it is performant on tasks like Memory Task and others
• Ensure that the resulting trajectories can be used to train decision transformers.

TrajectoryLSTM model used a BOW model that handles 7x7x3 encoding way better than it would work for the DT / BC models we use. Writing a converter between these formats would enable trajectories to be recorded with one type of observation to be converted to the other when training the subsequent offline agent.

At some point, I implemented weights and biases sweeps to find good hyperparameters when doing PPO training. I suspect this will be very useful in the future and would like it to also exist for the decision transformer script as well.

Examples from PPO:

• scripts/run_ppo_sweeps.sh (command to start it)
• https://github.com/jbloomAus/DecisionTransformerInterpretability/blob/main/sweeps/ppo_sweep_template.yml (template yml config file)

Tasks:

• make sure the PPO sweep works (I think there might be some funkiness with passing in the exp-name so might want to fix that first.
• copy the pattern and implement it for scripts/run_decision_transformer.sh

Currently store_model_checkpoint will add files to be uploaded to wandb to an artifact which is uploaded at the end of the ppo training cycle. It would be better if these we each uploaded as soon as they were created in case instances fail.

I wrote a few wrappers for minigrid environments. which should probably live in the Minigrid github repo.

If someone has the time to submit that as a PR and write whatever tests their maintainers require (probably not a huge amount, then this would be greatly appreciated.

Checklist:
[ ] make a PR to Minigrid with the wrappers (double check the view size wrapper doesn't exist, it could be that it just had a bug that I didn't want to wait for them to fix)
[ ] add tests to the PR
[ ] have the PR accepted
[ ] when the PR is accepted, update this repo to use the new minigrid version and remove the wrapper code from this repo.

After failing to generate a successful TransformerAC Model, I'm pivoting to generating a recurrent AC model based on the Baby Pytorch implementation (https://github.com/mila-iqia/babyai/blob/master/babyai/model.py)

The tasks involved are:

• Port over the model and get it running using our PPO methods, rewrite any if needed
• Write basic tests to ensure it works
• Ensure we can generate trajectories from it that we can train our decision transformers on

This is an important task I wouldn't want to do if there was lots of other work in flight or before I had some more results. However, I wanted to mark it here to keep track of any thoughts I have on it before I get to it.

It really bothers me that the minigrid wrapper for tokenising the observations also modifies the observations space for images. These should be untied.

class DictObservationSpaceWrapper(ObservationWrapper):
    """
    Transforms the observation space (that has a textual component) to a fully numerical observation space,
    where the textual instructions are replaced by arrays representing the indices of each word in a fixed vocabulary.

    This wrapper is not applicable to BabyAI environments, given that these have their own language component.

    Example:
        >>> import miniworld
        >>> import gymnasium as gym
        >>> import matplotlib.pyplot as plt
        >>> from minigrid.wrappers import DictObservationSpaceWrapper
        >>> env = gym.make("MiniGrid-LavaCrossingS11N5-v0")
        >>> obs, _ = env.reset()
        >>> obs['mission']
        'avoid the lava and get to the green goal square'
        >>> env_obs = DictObservationSpaceWrapper(env)
        >>> obs, _ = env_obs.reset()
        >>> obs['mission'][:10]
        [19, 31, 17, 36, 20, 38, 31, 2, 15, 35]
    """

    def __init__(self, env, max_words_in_mission=50, word_dict=None):
        """
        max_words_in_mission is the length of the array to represent a mission, value 0 for missing words
        word_dict is a dictionary of words to use (keys=words, values=indices from 1 to < max_words_in_mission),
                  if None, use the Minigrid language
        """
        super().__init__(env)

        if word_dict is None:
            word_dict = self.get_minigrid_words()

        self.max_words_in_mission = max_words_in_mission
        self.word_dict = word_dict

        image_observation_space = spaces.Box(
            low=0,
            high=255,
            shape=(self.agent_view_size, self.agent_view_size, 3),
            dtype="uint8",
        )
        self.observation_space = spaces.Dict(
            {
                "image": image_observation_space,
                "direction": spaces.Discrete(4),
                "mission": spaces.MultiDiscrete(
                    [len(self.word_dict.keys())] * max_words_in_mission
                ),
            }
        )

    @staticmethod
    def get_minigrid_words():
        colors = ["red", "green", "blue", "yellow", "purple", "grey"]
        objects = [
            "unseen",
            "empty",
            "wall",
            "floor",
            "box",
            "key",
            "ball",
            "door",
            "goal",
            "agent",
            "lava",
        ]

        verbs = [
            "pick",
            "avoid",
            "get",
            "find",
            "put",
            "use",
            "open",
            "go",
            "fetch",
            "reach",
            "unlock",
            "traverse",
        ]

        extra_words = [
            "up",
            "the",
            "a",
            "at",
            ",",
            "square",
            "and",
            "then",
            "to",
            "of",
            "rooms",
            "near",
            "opening",
            "must",
            "you",
            "matching",
            "end",
            "hallway",
            "object",
            "from",
            "room",
        ]

        all_words = colors + objects + verbs + extra_words
        assert len(all_words) == len(set(all_words))
        return {word: i for i, word in enumerate(all_words)}

    def string_to_indices(self, string, offset=1):
        """
        Convert a string to a list of indices.
        """
        indices = []
        # adding space before and after commas
        string = string.replace(",", " , ")
        for word in string.split():
            if word in self.word_dict.keys():
                indices.append(self.word_dict[word] + offset)
            else:
                raise ValueError(f"Unknown word: {word}")
        return indices

    def observation(self, obs):
        obs["mission"] = self.string_to_indices(obs["mission"])
        assert len(obs["mission"]) < self.max_words_in_mission
        obs["mission"] += [0] * (self.max_words_in_mission - len(obs["mission"]))

        return obs

I have two main ideas for this:

Implement a variation of the BOW encoding that is used by BabyAI but add position to avoid building a convnet.

Not discussed in the BabyAI paper itself, it seems like they actually used a much better tokenization scheme than I am using and this could plausibly be causing many of my problems.

class ImageBOWEmbedding(nn.Module):
   def __init__(self, max_value, embedding_dim):
       super().__init__()
       self.max_value = max_value
       self.embedding_dim = embedding_dim
       self.embedding = nn.Embedding(3 * max_value, embedding_dim)
       self.apply(initialize_parameters)

   def forward(self, inputs):
       offsets = torch.Tensor([0, self.max_value, 2 * self.max_value]).to(inputs.device)
       inputs = (inputs + offsets[None, :, None, None]).long()
       return self.embedding(inputs).sum(1).permute(0, 3, 1, 2)

This takes each position and it represents it as a unique embedding that is independent of position. For example, Key + yellow + close (object, state,color) -> 13 or something -> maps to a specific vector. They then pass this into a convolutional network (which I was hoping to avoid).

I can do better than this though in the context of my model. I can have, bear with me, 5 separate embedding matrices. Consider the current model, it's actually, 3 embedding matrices:

• 1 state
• 1 object
• 1 color
• 1 row position
• 1 column position.

Then any given object,color,state at any given position starts of with a unique representation. We can then look at how weight regularization acts as a feature selector over these and how the embeddings of each evolve in response to the circuits which use them.

One complication here is that for the positional embeddings, it seems like I should use something like sinusoidal embeddings but I want to ensure they are orthogonal to the previous sinusoidal embeddings. I have some ideas about how to do this but I will google it anyway.

Eg: this from the gato paper.

Haven't tried. suspect there'll be issue, possibly silent. Not a priority unless you are actually interested in a model that uses LNPre.

#43 was insufficient to create an optimal behaviour agent with the current rollout data.

To falsify the hypothesis that our failure to generate good trajectories (meet the criteria of both solving the task and being reasonable RTG modulated, at all, calibration unimportant for the time being) is a function of the data generating distribution with the current rollout methodology and PPO agent, I would like upgrade our ability to sample from PPO agents.

This will include:

1. Adding an option to, rather than randomly sample from the posterior distribution indicated by the logits, sample deterministically from the best action.
2. Adding an option to sample with temperature. This will enable us to calibrate the right amount of randomness.
3. Adding an option to sample with from anything except the maximal logit (unlike temperature, this will reliably produce anti-good trajectories rather than consistently good trajectories).

Furthermore, the utility for this should also provide some feedback on the sampled trajectories (for example, average RTG achieved, RTG distribution, trajectory length distribution etc).

Tasks:

• Write the end-end utility (ie, wrap what we have currently in a runner)
• Add each of the options above with tests
  • Deterministic Sampling
  • Temperature Sampling
  • Bottom-k sampling
• Make it possible to select more than of these and specify the number of rollouts to generate with each.
• Make it possible to view some example rollouts when doing so/also have some performance metrics.

In DT/BC training I use a weighted random sampler. I'm kinda suspicious of it and want to see what happens if I turn it off. I'd also potentially like to visualize the distribution of sampling probabilities

• Add a plotly graph upload of a histogram of the sampling probabilities
• See if basic DT's on Dynamics Obstacles learn different when it's off/on. Also try memory env.

While cloning or otherwise working with the repo I saw errors like these ones:

error: external filter 'git-lfs filter-process' failed. I assume it is related to this file fatal: trajectories/MiniGrid-Dynamic-Obstacles-8x8-v0bd60729d-dc0b-4294-9110-8d5f672aa82c.pkl:

Error downloading object: trajectories/MiniGrid-Dynamic-Obstacles-8x8-v0bd60729d-dc0b-4294-9110-8d5f672aa82c.pkl (8c06bc4): Smudge error: Error downloading trajectories/MiniGrid-Dynamic-Obstacles-8x8-v0bd60729d-dc0b-4294-9110-8d5f672aa82c.pkl

This error message is seemingly related to Git Large File Storage (LFS). The error message indicates that the git-lfs filter-process command failed. This error can occur if Git LFS is not installed or if the Git LFS filters are not configured correctly.

This solution appears to work for me:

Check if git-lfs is installed and install it if it is not.

git lfs install --skip-smudge
git clone <your-repo-url>
git lfs pull
git lfs install --force

The --skip-smudge option tells Git LFS to skip the smudge filter, which is responsible for downloading the large files. The git clone command clones the repository. The git lfs pull command downloads the large files. Finally, the --force option reinstates the smudge filter.

Maybe it would be helpful to add something on git lfs in the Readme or somewhere else.

Add static interpretability visualizations to wandb dashboard. Seems like a cool idea I just had.

Some stuff: QK/OV circuit viz, different attribution embeddings, time embedding, L2 Norms of different components.

From Gato paper: "Regularization: We train with an AdamW weight decay parameter of 0.1. Additionally, we use stochastic depth (Huang et al., 2016) during pretraining, where each of the transformer sub-layers (i.e. each Multi-Head Attention and Dense Feedforward layer) is skipped with a probability of 0.1."

Gao Huang, Yu Sun, Zhuang Liu, Daniel Sedra, and Kilian Weinberger. Deep networks with stochastic depth. Preprint arXiv:1603.09382, 2016.

Stochastic depth seems plausibly super valuable to me via intutions. I should read that paper at some point - Joseph

I used SyncVectorEnv to speed up game simulations elsewhere but rushed through writing the eval code and therefore didn't use it there.

Reimplementing evaluate_dt_agent with parallelization would speed up both the decision transformer training script and the calibration evaluation script.

Reasonable steps might look something like:

• Write a test which ensures that evaluate_dt_agent works
• Rewrite the function with parallelization (get statistics for each episode correctly, flatten them)
• See that the test passes
• Regenerate a calibration curve for models/MiniGrid-Dynamic-Obstacles-8x8-v0/demo_model_one_hot_overnight.pt, which should look like the one in the post (more/less).

Follow this tutorial: https://ljvmiranda921.github.io/notebook/2018/06/21/precommits-using-black-and-flake8/

We can sample rollouts using sample from agents. However, it would be good to be able to merge trajectory datasets. In order to do this, it might be worth cleaning up the offline_dataset code and working out the right way to do this. A fast way would be to use torch utils to concat datasets but this break the visualizer we have for the dataset. Possible not worth doing anything else.

Based on my understanding of GLUs from reading these three papers:

• https://arxiv.org/pdf/1612.08083.pdf (original, in conv net)
• https://arxiv.org/pdf/2002.05202.pdf (says it's good in transformers)
• https://arxiv.org/pdf/2301.07608.pdf (show's it's useful in similar systems)

It seems likely that such modification on our transformers will be valuable. Moreover, Neel told me that they are used in models including T5 (see here: https://huggingface.co/docs/transformers/model_doc/t5, they use gated gelu).

I think this shouldn't be too hard to add as a PR, and I could theoretically accept it myself but could talk to Neel first. At a high level we have:

• 1. makes sure you can specific this in the config reasonably
• 2. make sure you can implement it without any problems (check it works as expected). have tests.
• 3. make sure you get hooks for the relevant layers (we will need to add an activation)
• 4. check if the default d_mlp should change if you do this.

And then a nice optional thing to do is to make a tutorial talking about how to interpret it. I think this would be best done once we have loaded T5 and were doing interp on it. Or we can train a very small model that uses it and use that as a demo. I also want to to know what other people link about this.

https://wandb.ai/arena-ldn/PPO-MiniGrid/artifacts/trajectory/4eb3c096-8836-4d0f-973a-67685b89d0f0.gz/12e1daaebfd9c806051f

python -m src.run_decision_transformer \
    --exp_name MiniGrid-MemoryS7FixedStart-v0 \
    --trajectory_path trajectories/4eb3c096-8836-4d0f-973a-67685b89d0f0.gz \
    --d_model 128 \
    --n_heads 2 \
    --d_mlp 256 \
    --n_layers 1 \
    --learning_rate 0.0001 \
    --batch_size 128 \
    --train_epochs 5 \
    --test_epochs 1 \
    --n_ctx 23 \
    --pct_traj 1 \
    --weight_decay 0.001 \
    --seed 1 \
    --wandb_project_name DecisionTransformerInterpretability \
    --test_frequency 1000 \
    --eval_frequency 1000 \
    --eval_episodes 10 \
    --initial_rtg -1 \
    --initial_rtg 0 \
    --initial_rtg 1 \
    --prob_go_from_end 0.1 \
    --eval_max_time_steps 1000 \
    --track True

  File "/Users/josephbloom/miniforge3/envs/decision_transformer_interpretability/lib/python3.10/runpy.py", line 196, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/Users/josephbloom/miniforge3/envs/decision_transformer_interpretability/lib/python3.10/runpy.py", line 86, in _run_code
    exec(code, run_globals)
  File "/Users/josephbloom/GithubRepositories/DecisionTransformerInterpretability/src/run_decision_transformer.py", line 55, in <module>
    run_decision_transformer(
  File "/Users/josephbloom/GithubRepositories/DecisionTransformerInterpretability/src/decision_transformer/runner.py", line 114, in run_decision_transformer
    model = train(
  File "/Users/josephbloom/GithubRepositories/DecisionTransformerInterpretability/src/decision_transformer/train.py", line 142, in train
    evaluate_dt_agent(
  File "/Users/josephbloom/GithubRepositories/DecisionTransformerInterpretability/src/decision_transformer/train.py", line 342, in evaluate_dt_agent
    new_obs, new_reward, terminated, truncated, info = env.step(action)
  File "/Users/josephbloom/miniforge3/envs/decision_transformer_interpretability/lib/python3.10/site-packages/gymnasium/vector/vector_env.py", line 203, in step
    return self.step_wait()
  File "/Users/josephbloom/miniforge3/envs/decision_transformer_interpretability/lib/python3.10/site-packages/gymnasium/vector/sync_vector_env.py", line 149, in step_wait
    ) = env.step(action)
  File "/Users/josephbloom/miniforge3/envs/decision_transformer_interpretability/lib/python3.10/site-packages/gymnasium/wrappers/record_video.py", line 155, in step
    ) = self.env.step(action)
  File "/Users/josephbloom/miniforge3/envs/decision_transformer_interpretability/lib/python3.10/site-packages/gymnasium/core.py", line 408, in step
    return self.env.step(action)
  File "/Users/josephbloom/miniforge3/envs/decision_transformer_interpretability/lib/python3.10/site-packages/gymnasium/wrappers/record_episode_statistics.py", line 89, in step
    ) = self.env.step(action)
  File "/Users/josephbloom/miniforge3/envs/decision_transformer_interpretability/lib/python3.10/site-packages/gymnasium/wrappers/order_enforcing.py", line 56, in step
    return self.env.step(action)
  File "/Users/josephbloom/miniforge3/envs/decision_transformer_interpretability/lib/python3.10/site-packages/gymnasium/wrappers/env_checker.py", line 49, in step
    return self.env.step(action)
  File "/Users/josephbloom/GithubRepositories/DecisionTransformerInterpretability/src/environments/memory.py", line 167, in step
    if action == Actions.pickup:
RuntimeError: Boolean value of Tensor with more than one value is ambiguous

See discussion in #43

The collect demonstrations utility is responsible for collecting example trajectories from a trained agent (one of the 3 ppo agent architectures supported, only two work). It provides a few different sampling procedures for doing this such as basic (proportional to softmax), temperature sampling, bottomk and topk. These enable us to train decision transformers on a broader distribution of actions/observations encouraging more robust features to be learned and a better calibrated RTG - behavior relationship (if you do behavioral cloning, you are restricted to training on only good trajectories with the offline agent which doesn't lead to good features).

It's not super clear yet how useful this is, but there's very obvious next steps to improve this utility that seem like a good engineering practice if anyone wants to help. The major goals are:

• Hook up wandb tracking (add an arg for track, then log metrics to the dashboard)
• media to log: videos of each of the rollouts (same as ppo rollout out code somewhat). We want to see qualitatively what kinds of trajectories we're sampling.
• Metric to log: reward/time to finish for different rollout configs, we want to know which config setting sample what kinds of trajectories/outcomes.

It's possible that is an opportunity for the code around uploading videos which is very janky can be improved.

I've explained how to use the repo in the docs, but this could be way better.

Some examples of pages what might be useful to add to the sphinx docs.

1. A page explaining what Minigrid is, and where to look at each MiniGrid environment (maybe include a table of environments by row with columns for different models (dt, bc, and ppo-traj, ppo standard).
2. A page explaining the difference between the online method (PPO) and offline methods, explaining why we need the trajectories for the latter.
3. A page/diagram explaining how trajectories are generated with the ppo method used to train offline agents.
4. MOST IMPORTANT: pages for stuff you think I won't have thought to explain since I wrote the repo.

The app likely doesn't work with behavioural clone models and this should be fixed so that we can do a comparison.

A couple of weeks ago I worked with Navpreet to start writing a Maze Environment for Minigrid which is mostly working. Finishing this PR and possibly adding a version that uses Kruskal's algorithm and not Prim's could be a really valuable contribution.

The exciting thing about these environments being part of Minigrid is that it would offer opportunities to study algorithmic distillation and precise, memory-based reasoning in transformers. I think doing this could be the start of a lot of cool work.

I currently host a public version of the interpretability app with the streamlit community, which is very slow.

I would be willing to pay for a service to host the app, but I need to know that deployment will be very easy.

Checklist:

• find a third party app hosting service
• show that it is easy to deploy the streamlit app to this service (https://github.com/jbloomAus/DecisionTransformerInterpretability/blob/c84edb381c53b3f9ef2fa9517e34914a52e15fbd/app.py)
• show that the service is faster that streamlit community deployment
• make a PR with a github action which deploys the app

This is a bug in the trajectorywriter/offline dataset where we end up truncating some trajectories when we finish online training and this leads to having “short” truncated trajectories, which are bad for our data. It would be good to remove them. They are visible in the visualization of the reward over traj-lengths as spots on the x-axis but not at max-length.

A link to the method I use to ensure that these get labelled as truncated to avoid bugs:

See #46, investigate effect on memory model.

• try one or two at different ratios
• see if it makes a big difference.

/Just wanted to have meta card to track progress on these things with links:

• LayerNorm (I'll probably only try layernorm pre) (#52)
• AdamW Optimizer
• Adding a warmup stage with LambdaLR scheduler or cosine annealing
• Implement gated MLP's (https://arxiv.org/pdf/2002.05202.pdf). Might need to be done in TransformerLens.
• Make it possible to use GeLU not ReLU (try that out as well).
• Better encode state. #61
• Look into current init ranges for all the model components and consider proper init ranges
• Look into where all the parameters are and consider how we can make a sparser model
• Implement wandb sweeps for DT training (likely already exists a card for this so I should find it)
• Implement masking rather than just having different tokens during padding. Might be important?

If we've implemented all of those and still no success with the memory env training, possibly try either much longer training runs, more variable sampling methods, or ask for advice (or go bug hunting).

I recently wrote a version of get_minibatches in the memory class of the ppo subpackage.

TLDNR: This is important for sampling sections of trajectories which is necessary for online training of trajectory models as opposed to models which only respond to the latest observation. I have a few ideas for what to do here:

Keep the logic more or less the same, but vectorize it. It's way to serialized and it doesn't have to be. Obviously write lots of tests.

Similar to the Maze experiments done by the shard theory team, calculate a vector corresponding to the key/ball in the memory env (or make a more general tool) that can be used in the app to analyse object configured behaviors.

PM me if you are interested in working on this.

We aren't currently masking padding tokens which I'm not sure how much of a big deal it is. I'll reach out to ask someone for feedback, but in the meantime, I'm going to set the RTG in the padded training data to the true RTG because the padding RTG is 0 which I think might have confused the model.

This is hopefully going to be the different between a memory task solving DT completely solving the task and it's current status of solving it but not using RTG to modulate behavior for the balls (it works if the key is the goal strangely).

• update offline training data to padd with true RTG
• message someone about how essential padding is (I'm not sure why it isn't already implemented on Hooked Transformer class, maybe it is and I missed it)

As part of #53, I'm going to explore some techniques which may prove useful for improving the performance of DTs / BCs.

• Add a wandb graph showing the lr of the scheduler over time (this will ensure that we can validate the methods we use
• Then add a LR scheduler utility (I think we'll just make use of the huggingface transformer utilities?)
• Validate each work and see if any affect converge speed of Dynamic Obstacles.
• Can try out on memory env if desired.

Our PPO models are now stored with checkpoints but our offline trained models aren't. Creating some parity here would be good.

Please ensure:

• It remains easy to load a model into the app
• All of the required parameters to instantiate a model are saved with that model
• Nice utilities facilitate saving/loading of checkpoints as is the case with the PPO models
• Don't duplicate anything you don't have to with the PPO checkpoint set up.

The end-end tests show how to pass arguments to the different PPO models (FC, Transformer, LSTM) but it might nice to have a command line tool that handles all of these.

I'm nervous that I haven't been disciplined enough with the initialization of the weights in my models. I should research how people think about initialization and what to do about it.

I wrote a class that mixes environments/hyperparameters called a Multienv which is kinda useful for training, however it should probably live in the Minigrid github repo.

If someone has the time to submit that as a PR and write whatever tests their maintainers require (probably not a huge amount, then this would be greatly appreciated.

Checklist:
[ ] make a PR to Minigrid with the Multienv
[ ] add tests to the PR
[ ] have the PR accepted
[ ] when the PR is accepted, update this repo to use the new minigrid version and remove the multienv code from this repo.

The BOWEmbedding from BabyAI has very large vectors on init. I wonder if that's really terrible and was slowing down my training of demo generating models. Need to investigate at some point.

It's now possible to allow columns (up to one level of nesting) inside columns in streamlit apps thanks to streamlit/streamlit#5941. I think this creates an oppurtunity to make the streamlit interpretability app way more visually appealing. I don't have strong opinions now but basically suspect someone could have a lot of fun with this and it'll be really enjoyable later.

For reference, the current streamlit app uses a portrait view and a sidebar for selecting which analyses appear where. I think it'd be interesting to redesign this. Happy to do a "user interview" to describe how I use stuff and what I don't like about how I've currently implemented it.

I failed to get the TransformerPPO model working. I suspect this is because of bugs, but it could also be that transformers are inherently unstable.

If someone is interested in attempting to fix my implementation then I would greatly appreciate it. For the time being, I'm going to pivot to a recurrent architecture to get progress on the project going again.

Start here:

Big tasks are:

• Get the Traj PPO agent to pass the memory test.
• passing some of the skipped acceptance tests might be useful feedback on the way.

In order to have enough training data to give decision transformers a good distribution over good/bad trajectories. It would good to be able to sample from an arbitrary quality agent (which means we need to store these agents).

As an initial step, storing model checkpoints during training of PPO agents would be useful. Storing these on wandb seems like a good first step. (there are examples in the code of the uploaded artifact).

Analysis features

Static

Composition

• Make composition maps
• Replace composition scores with strip plots?
• Create a meta-composition score. Something that measures total influence?
• How do we check for composition between MLP_in and W_out? (seems expensive?, maybe tie to very specific hypotheses)

Dynamic

Logit Lens

• By Layer
• By Layer accumulated
• By Head

Attention Maps:

• Make it easier to export a nice visualization of the attention map (cv is actually not great for that).
• Make it possible to calculate the rank(k) approximation to the attention map.

Causal

Activation Patching (features)

• Set up component
• Set up RTG Metric
• Residual stream patching.
• Patching via Attn and MLP
• Head All Pos Patching
• Head Specific Pos Patching (do later)
• Head All Pos by Component
• MLP at different Positions
• Show counterfactual attention map (ie: show difference in attention given intervention)
• Show what the logit diff is for each metric score.
  Activation Patching (token variations):
• Action (fairly easy)
• Key/Ball (important!)
• Timestep (also fairly easy)

RTG Scan

• Switch to using t-lens for decomp
• Provide more than one level of decomp
• Add a clustergram to show heads which mediate a similar relationship between RTG and logits/logit diff

Congruence -> If features aren't in superposition, what effect do they have on the predictions?

• - Pos
• - Time
• - W_in
• - W_Out
• - MLP Out

Renew old features:

• QK circuit visualizations for action and RTG embeddings

SVD Decomp / Explore ways to use dimensionality reduction to quickly understand what heads are doing.

• QK Circuit SVD SVD Decomp / Explore ways to use dimensionality reduction to quickly understand what heads are doing.#69
• OV Circuit SVD

Cache Characterization?

• Plot L2 norm of residual streams (along with mean and std)

Advanced

Implement Path Patching

• Understand Callum's code.

Implement AVEC

• Reread post to see if we can find.

Several things I feel are missing which are required for exploratory analysis to be more complete:

• visualise dot product of time embeddings with each other
• visualise dot product of positional embeddings with each other
• Use Jay's head type analysis but write specific patterns for attending to RTG, attending to positive RTG, attending to states, and attending to actions.

Several things I feel will be required for falsifying predictions of how the model is working:

• implement a variant of path patching for DTs either in a notebook or as part of the app.
• CaSc, not sure how feasible this is but it has always been the goal.

I attempted this but weirdly got a clash with black even if profile = black with isort.

• Psychological eval
• Activation Patching for instruction and RTG. -> try to explain
• Work out how to tackle targets (patching same object multiple frames.
• [ ]

Link to excalidraw: https://link.excalidraw.com/l/92qNDVgxtkX/A4R4OHfBqAu

I think the make_env function should take in the environment_config class we're using everywhere.

A link the the make_env function:

A link to the config class it should take:

Be sure to make sure the tests all pass! I think this will involve touching a lot of different pieces of code but it's still conceptually very simple.

Selective Noise Injection (SNI) and Information Bottleneck Actor-Critic (IBAC) make models better at generalising (including in at least one MiniGrid environment). It seems like a fun hack-dayish kind of effort to test this out.

This would currently be bottlenecked by the TrajectoryPPO class possibly not being perfectly working yet but I'm putting this here in case someone is ambitious and wanted to give it a shot.

References:

• https://github.com/microsoft/IBAC-SNI/
• https://arxiv.org/pdf/1910.12911.pdf

I consider this somewhat similar spiritually to Engineering Monosemanticity in Toy Models although I have no idea if this is exactly true.

Reading papers suggests this should help. I'm going to try it. It's already implemented except maybe in the app so I will:

• Set layer-norm to LNPre
• Training a few different models with this config. (the memory env and maybe one of our original models.
• Observing any differences in training metrics
• If the model turns out to be more performant ensure app compatibility now, otherwise make a card for it.