Towards Lifelong RL through Zero-Shot Composition

Jun 29, 20264 min

https://geraudnt.github.io/index.html

Goal: Multipurpose robots that are (1) flexible, (2) instructable, (3) reliable.

  • What is a task? → MDP: S,A,P,R\langle S, A, P, R \rangle
  • What solution should an agent learn? → Policy
    • Standard RL: π=argmaxπQ(s,a)\pi^* = \arg\max_\pi Q^*(s,a) is the solution to a single task
  • ⚠️ Problem: learning is still not compositional - learning separately how to go to blue shapes and squares in a game doesn't yield an easy zero-shot way to obtain a policy for going to blue squares
    • Can we combine the QQ^* functions learned directly (e.g. min(Q(blue),Q(square))\min (Q^*(\text{blue}), Q^*(\text{square}))? - no info about common goal though
    • 🔖 No function ff that takes in the knowledge you've learned about blue (e.g. Q(blue)Q^*(\text{blue})) and knowledge you've learned about square (Q(square)Q^*(\text{square})) and spits out Q(blue and square)Q^*(\text{blue and square})
      • But you can do this for value function! max(V1,V2)=V12\max(V_1, V_2) = V_{1 \cup 2}

Word value functions (WVFs)

Intuition: While solving one task, we should learn about other tasks that we may need to solve in the future.

  • Take the standard RL loop, and only change what is in the agent:
    • Agent learns a general value function
    • Agent decides what goal to achieve this time (Gtset of terminal statesG_t \in \text{set of terminal states}), what action to take, and what reward it receives

https://www.raillab.org/publication/nangue-2025-composition/nangue-2025-composition.pdf

  • Policies: π(s,a,g)\pi(s, a, g)
  • Reward function: rˉ(s,g,a)={rˉminifgsG(achievedagoalthatwasntthegoalstate)r(s,a)otherwise\bar{r}(s, g, a) = \cases{\bar{r}_{\min} & if g \neq s \in G (achieved a goal that wasn't the goal state) \\r(s,a) & otherwise}
  • Q - function:
  • The QQ^* learned off of this reward ^ accumulates information for multiple goals!
    • WVFs.png
    • But now we have lots of goal function values! How does the agent figure out how to combine them? → max over WVFs!
    • ❓ Are we learning the layers of WVFs individually?
      • Learned at once - π=argmaxπQπ(s,g,a)\pi^* = \arg\max_\pi Q^\pi(s, g, a) - layers of WVFs are gotten for plugging in different gg's
      • ^ pin down one of the gg's and you get a slice of the big WVF

Skills are WVFs!

  • To get the WVF for "blue + square", take the pointwise min of Q(blue)Q^*(\text{blue}) and Q(square)Q^*(\text{square})
    • This way spots that correspond to both "blue" and "square" still have high value (after composing via min)
    • min turns out to be the provably correct way of composing WVFs → gets you the optimal QQ of intersections
  • For a task space - what is the minimal set of primitives you can compose to get all tasks in the space?
    • "basis" set of primitives
    • "minimum generating set" in Boolean algebra
    • goal & minimal task are different! e.g. goal = "blue square", minimal task = "square" → essentially factoring task space into primitives, training WVFs with each of these primitives as "goals", and then when using the trained Q(s,a,g)Q^*(s, a, g) can plug in primitives to gg and then compose into the "goal" (e.g. blue square)
  • Many different compositions work on WVFs!
    • AND, OR, XOR, NOT
    • NOT requires the agent to learn what "do nothing" means in the task space

Extensions

  • Continuous goal domain
  • Language
    • Method 1: T5 translation model to convert natural language to Boolean expression
    • Method 2: LLM to convert natural language to Boolean expression
  • Zero-shot temporal logics
  • WVFs in brains? Spatial maps in brains - grid cells
    • SVD of WVFs gives you similar map as grid cells in brain images
    • Grid cells from WVF!
  • Next steps:
    • SafeRL
    • POMDPs

Questions

  • Doing this vs. hindsight experience replay?
    • Whole Boolean lattice of possible tasks available in this method - you could do this in HER, but your goal space in HER would need to contain all the exponential possibilities to match this method
    • Compositions for the win :)!
    • exp growth curve WVFs.png ← "can solve" = we have π\pi^* for all tasks
    • ^ Super-exponential explosion of skills! (since there are exponential compositions of blue, square, beige, and we do some online stuff, but still yay!)
  • ❓ How to extend this method to continuous goal state scenarios?
    • Representation learning - learn mapping from continuous goal space to finite discrete goal space for the model to think in
    • Energy-based models - continuous optimization (MCMC)
  • ❓ Can you learn tasks as they come (or does WVF need to learn everything upfront)? → Yes, incrementality works
    • ⚠️ Problem is that if I want a robot to make toast, it should come to my house and be able to do this (not try to learn it fresh right now)
  • ❓ What's the downside? Any scalability issues?
    • Need to learn WVF (value function is like amortization of knowledge - takes time to learn it)
      • Time it takes to learn WVF to converge takes super long in comparison to learning value function for a specific task
      • ^ but as you try to do more tasks, the cumulative time you spend learning WVF starts making sense when you amortize it across tasks to do (polynomial growth in learning time, exponential growth in your ability to do things)