Fix embedding

Fix various bugs in the embedding.

• Allow for missing x in batch: use other variables instead of just always relying on x to be there
• Properly get the features for the embedding
• Fix some formatting
• Fix some typos